| 1 | //! A real MCP client that speaks JSON-RPC to a spawned child process. |
| 2 | //! |
| 3 | //! `codewhale mcp-server` used to wire every configured server to an |
| 4 | //! in-memory stub, so a user's `command`/`args`/`env` were never executed and |
| 5 | //! every health probe answered `{"status": "ok"}` from a hardcoded literal |
| 6 | //! (#4727). A fabricated success is the worst possible answer here: it is |
| 7 | //! indistinguishable from a working integration. This module replaces it with |
| 8 | //! an actual subprocess connection, and every failure path below returns an |
| 9 | //! error naming the server rather than a plausible-looking value. |
| 10 | |
| 11 | use std::collections::HashSet; |
| 12 | use std::io::{self, BufRead, BufReader, Write}; |
| 13 | use std::process::{Child, ChildStdin, Command, Stdio}; |
| 14 | use std::sync::mpsc::{Receiver, RecvTimeoutError, sync_channel}; |
| 15 | use std::sync::{Arc, Mutex}; |
| 16 | use std::thread; |
| 17 | use std::time::{Duration, Instant}; |
| 18 | |
| 19 | #[cfg(unix)] |
| 20 | use std::os::unix::process::CommandExt; |
| 21 | #[cfg(windows)] |
| 22 | use std::os::windows::io::{AsRawHandle, FromRawHandle, OwnedHandle}; |
| 23 | #[cfg(windows)] |
| 24 | use std::os::windows::process::CommandExt; |
| 25 | |
| 26 | use anyhow::{Context, Result, anyhow, bail}; |
| 27 | use serde_json::{Value, json}; |
| 28 | |
| 29 | use crate::{ |
| 30 | MCP_PROTOCOL_VERSION, MCP_SUPPORTED_PROTOCOL_VERSIONS, McpManagedClient, McpResourceDescriptor, |
| 31 | McpServerConfig, McpToolDescriptor, |
| 32 | }; |
| 33 | |
| 34 | /// Budget for spawn + `initialize` + `notifications/initialized`. Generous |
| 35 | /// because a first `npx`/`uvx` launch may download the server package. |
| 36 | const HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(30); |
| 37 | |
| 38 | /// Budget for a single request once the server is up. |
| 39 | const REQUEST_TIMEOUT: Duration = Duration::from_secs(120); |
| 40 | |
| 41 | /// How long a dropped client waits for a graceful exit after closing stdin |
| 42 | /// before it kills the child. |
| 43 | const SHUTDOWN_GRACE: Duration = Duration::from_millis(500); |
| 44 | |
| 45 | /// How long a failure path waits for a dying child's exit status before giving |
| 46 | /// up and reporting the failure without one. |
| 47 | const EXIT_STATUS_GRACE: Duration = Duration::from_millis(200); |
| 48 | |
| 49 | /// Hard bounds for paginated catalog discovery. These match the TUI's MCP |
| 50 | /// catalog guardrails so the proxy cannot turn one malformed catalog into an |
| 51 | /// unbounded allocation or wait. |
| 52 | const MAX_LIST_PAGES: usize = 64; |
| 53 | const MAX_LIST_ITEMS: usize = 4_096; |
| 54 | const MAX_LIST_BYTES: usize = 32 * 1024 * 1024; |
| 55 | /// A single JSON-RPC message cannot be larger than the entire catalog budget. |
| 56 | /// Enforce this while consuming the child's buffered stdout, before a newline |
| 57 | /// (or EOF) can make an unbounded `read_line` allocation visible to the proxy. |
| 58 | pub(crate) const MAX_JSONRPC_LINE_BYTES: usize = MAX_LIST_BYTES; |
| 59 | /// Bound queued child output as well as individual lines. When the consumer is |
| 60 | /// slower than the child, the reader blocks and the OS pipe supplies the rest |
| 61 | /// of the backpressure instead of allowing an unbounded heap queue. A |
| 62 | /// zero-capacity channel is deliberate: at most the line currently being read |
| 63 | /// and the line currently being handled can be resident in userspace. |
| 64 | const MAX_PENDING_CHILD_MESSAGES: usize = 0; |
| 65 | |
| 66 | fn valid_tool_input_schema(input_schema: &Value) -> bool { |
| 67 | let Some(schema) = input_schema.as_object() else { |
| 68 | return false; |
| 69 | }; |
| 70 | if schema.get("type").and_then(Value::as_str) != Some("object") { |
| 71 | return false; |
| 72 | } |
| 73 | if let Some(properties) = schema.get("properties") { |
| 74 | let Some(properties) = properties.as_object() else { |
| 75 | return false; |
| 76 | }; |
| 77 | if properties.values().any(|property| !property.is_object()) { |
| 78 | return false; |
| 79 | } |
| 80 | } |
| 81 | if let Some(required) = schema.get("required") { |
| 82 | let Some(required) = required.as_array() else { |
| 83 | return false; |
| 84 | }; |
| 85 | if required.iter().any(|name| !name.is_string()) { |
| 86 | return false; |
| 87 | } |
| 88 | } |
| 89 | true |
| 90 | } |
| 91 | |
| 92 | fn valid_annotations(annotations: &Value) -> bool { |
| 93 | let Some(annotations) = annotations.as_object() else { |
| 94 | return false; |
| 95 | }; |
| 96 | if let Some(audience) = annotations.get("audience") { |
| 97 | let Some(audience) = audience.as_array() else { |
| 98 | return false; |
| 99 | }; |
| 100 | if audience |
| 101 | .iter() |
| 102 | .any(|role| !matches!(role.as_str(), Some("user") | Some("assistant"))) |
| 103 | { |
| 104 | return false; |
| 105 | } |
| 106 | } |
| 107 | if let Some(priority) = annotations.get("priority") { |
| 108 | let Some(priority) = priority.as_f64() else { |
| 109 | return false; |
| 110 | }; |
| 111 | if !priority.is_finite() || !(0.0..=1.0).contains(&priority) { |
| 112 | return false; |
| 113 | } |
| 114 | } |
| 115 | true |
| 116 | } |
| 117 | |
| 118 | fn optional_string_field( |
| 119 | fields: &serde_json::Map<String, Value>, |
| 120 | field: &str, |
| 121 | context: &str, |
| 122 | ) -> Result<Option<String>> { |
| 123 | match fields.get(field) { |
| 124 | None => Ok(None), |
| 125 | Some(Value::String(value)) => Ok(Some(value.clone())), |
| 126 | Some(_) => bail!("{context}.{field} must be a string"), |
| 127 | } |
| 128 | } |
| 129 | |
| 130 | fn parse_tool_entry( |
| 131 | server_name: &str, |
| 132 | tool: &Value, |
| 133 | index: usize, |
| 134 | allow_legacy_schema_omission: bool, |
| 135 | ) -> Result<(McpToolDescriptor, Value)> { |
| 136 | let context = format!("MCP server '{server_name}': tools/list tools[{index}]"); |
| 137 | let fields = tool |
| 138 | .as_object() |
| 139 | .with_context(|| format!("{context} must be an object"))?; |
| 140 | let tool_name = fields |
| 141 | .get("name") |
| 142 | .and_then(Value::as_str) |
| 143 | .with_context(|| format!("{context}.name must be a string"))? |
| 144 | .to_string(); |
| 145 | let description = optional_string_field(fields, "description", &context)?; |
| 146 | let input_schema = match fields.get("inputSchema") { |
| 147 | Some(schema) if valid_tool_input_schema(schema) => schema.clone(), |
| 148 | Some(_) => bail!("{context}.inputSchema must be a valid object-shaped MCP input schema"), |
| 149 | None if allow_legacy_schema_omission => { |
| 150 | // Servers that omit the entire initialize capabilities object are |
| 151 | // already identified as legacy by `validate_initialize_result`. |
| 152 | // Preserve omission compatibility only for that explicit mode; |
| 153 | // a standard advertised tools capability must send inputSchema. |
| 154 | json!({"type": "object", "properties": {}}) |
| 155 | } |
| 156 | None => bail!("{context}.inputSchema is required for an advertised MCP tool"), |
| 157 | }; |
| 158 | Ok(( |
| 159 | McpToolDescriptor { |
| 160 | server_name: server_name.to_string(), |
| 161 | // The manager owns qualification; report the raw name and let it |
| 162 | // build `mcp__server__tool`. |
| 163 | qualified_name: tool_name.clone(), |
| 164 | tool_name, |
| 165 | description, |
| 166 | }, |
| 167 | input_schema, |
| 168 | )) |
| 169 | } |
| 170 | |
| 171 | fn parse_resource_entry( |
| 172 | server_name: &str, |
| 173 | resource: &Value, |
| 174 | index: usize, |
| 175 | ) -> Result<(McpResourceDescriptor, Value)> { |
| 176 | let context = format!("MCP server '{server_name}': resources/list resources[{index}]"); |
| 177 | let fields = resource |
| 178 | .as_object() |
| 179 | .with_context(|| format!("{context} must be an object"))?; |
| 180 | let uri = fields |
| 181 | .get("uri") |
| 182 | .and_then(Value::as_str) |
| 183 | .with_context(|| format!("{context}.uri must be a string"))? |
| 184 | .to_string(); |
| 185 | let name = fields |
| 186 | .get("name") |
| 187 | .and_then(Value::as_str) |
| 188 | .with_context(|| format!("{context}.name must be a string"))? |
| 189 | .to_string(); |
| 190 | let description = optional_string_field(fields, "description", &context)?; |
| 191 | let mime_type = optional_string_field(fields, "mimeType", &context)?; |
| 192 | |
| 193 | let mut metadata = json!({"name": name}); |
| 194 | if let Some(mime_type) = mime_type { |
| 195 | metadata["mimeType"] = Value::String(mime_type); |
| 196 | } |
| 197 | if let Some(size) = fields.get("size") { |
| 198 | if size.as_i64().is_none() && size.as_u64().is_none() { |
| 199 | bail!("{context}.size must be an integer"); |
| 200 | } |
| 201 | metadata["size"] = size.clone(); |
| 202 | } |
| 203 | if let Some(annotations) = fields.get("annotations") { |
| 204 | if !valid_annotations(annotations) { |
| 205 | bail!("{context}.annotations is not a valid MCP annotations object"); |
| 206 | } |
| 207 | metadata["annotations"] = annotations.clone(); |
| 208 | } |
| 209 | |
| 210 | Ok(( |
| 211 | McpResourceDescriptor { |
| 212 | server_name: server_name.to_string(), |
| 213 | uri, |
| 214 | description, |
| 215 | }, |
| 216 | metadata, |
| 217 | )) |
| 218 | } |
| 219 | |
| 220 | /// Read one newline-delimited child message without ever retaining more than |
| 221 | /// `max_bytes`. On an oversized line the reader is intentionally abandoned; |
| 222 | /// continuing after losing JSON-RPC framing would be unsafe. |
| 223 | pub(crate) fn read_bounded_line<R: BufRead>( |
| 224 | reader: &mut R, |
| 225 | max_bytes: usize, |
| 226 | ) -> io::Result<Option<String>> { |
| 227 | let mut line = Vec::new(); |
| 228 | loop { |
| 229 | let available = reader.fill_buf()?; |
| 230 | if available.is_empty() { |
| 231 | if line.is_empty() { |
| 232 | return Ok(None); |
| 233 | } |
| 234 | break; |
| 235 | } |
| 236 | |
| 237 | if let Some(newline) = available.iter().position(|byte| *byte == b'\n') { |
| 238 | if line.len().saturating_add(newline) > max_bytes { |
| 239 | return Err(io::Error::new( |
| 240 | io::ErrorKind::InvalidData, |
| 241 | format!("JSON-RPC line exceeded the {max_bytes}-byte limit"), |
| 242 | )); |
| 243 | } |
| 244 | line.extend_from_slice(&available[..newline]); |
| 245 | reader.consume(newline + 1); |
| 246 | break; |
| 247 | } |
| 248 | |
| 249 | if line.len().saturating_add(available.len()) > max_bytes { |
| 250 | return Err(io::Error::new( |
| 251 | io::ErrorKind::InvalidData, |
| 252 | format!("JSON-RPC line exceeded the {max_bytes}-byte limit"), |
| 253 | )); |
| 254 | } |
| 255 | line.extend_from_slice(available); |
| 256 | let consumed = available.len(); |
| 257 | reader.consume(consumed); |
| 258 | } |
| 259 | |
| 260 | if line.last() == Some(&b'\r') { |
| 261 | line.pop(); |
| 262 | } |
| 263 | String::from_utf8(line).map(Some).map_err(|err| { |
| 264 | io::Error::new( |
| 265 | io::ErrorKind::InvalidData, |
| 266 | format!("child stdout was not valid UTF-8: {err}"), |
| 267 | ) |
| 268 | }) |
| 269 | } |
| 270 | |
| 271 | enum ChildStdoutMessage { |
| 272 | Line(String), |
| 273 | Invalid(String), |
| 274 | } |
| 275 | |
| 276 | fn response_to_server_request(message: &Value) -> Option<Value> { |
| 277 | let method = message.get("method").and_then(Value::as_str)?; |
| 278 | let id = message.get("id")?; |
| 279 | Some(match method { |
| 280 | "ping" => json!({"jsonrpc": "2.0", "id": id, "result": {}}), |
| 281 | _ => json!({ |
| 282 | "jsonrpc": "2.0", |
| 283 | "id": id, |
| 284 | "error": { |
| 285 | "code": -32601, |
| 286 | "message": format!("Method not found: {method}") |
| 287 | } |
| 288 | }), |
| 289 | }) |
| 290 | } |
| 291 | |
| 292 | fn write_jsonrpc_line<W: Write>(writer: &mut W, message: &Value) -> Result<()> { |
| 293 | let mut line = serde_json::to_string(message)?; |
| 294 | line.push('\n'); |
| 295 | writer.write_all(line.as_bytes())?; |
| 296 | writer.flush()?; |
| 297 | Ok(()) |
| 298 | } |
| 299 | |
| 300 | struct ListBudget { |
| 301 | method: String, |
| 302 | pages: usize, |
| 303 | items: usize, |
| 304 | bytes: usize, |
| 305 | max_pages: usize, |
| 306 | max_items: usize, |
| 307 | max_bytes: usize, |
| 308 | seen_cursors: HashSet<String>, |
| 309 | deadline: Instant, |
| 310 | overall_timeout: Duration, |
| 311 | } |
| 312 | |
| 313 | impl ListBudget { |
| 314 | fn new(method: &str, timeout: Duration) -> Self { |
| 315 | Self { |
| 316 | method: method.to_string(), |
| 317 | pages: 0, |
| 318 | items: 0, |
| 319 | bytes: 0, |
| 320 | max_pages: MAX_LIST_PAGES, |
| 321 | max_items: MAX_LIST_ITEMS, |
| 322 | max_bytes: MAX_LIST_BYTES, |
| 323 | seen_cursors: HashSet::new(), |
| 324 | deadline: Instant::now() + timeout, |
| 325 | overall_timeout: timeout, |
| 326 | } |
| 327 | } |
| 328 | |
| 329 | #[cfg(test)] |
| 330 | fn with_limits( |
| 331 | method: &str, |
| 332 | timeout: Duration, |
| 333 | max_pages: usize, |
| 334 | max_items: usize, |
| 335 | max_bytes: usize, |
| 336 | ) -> Self { |
| 337 | Self { |
| 338 | max_pages, |
| 339 | max_items, |
| 340 | max_bytes, |
| 341 | ..Self::new(method, timeout) |
| 342 | } |
| 343 | } |
| 344 | |
| 345 | fn remaining_timeout(&self) -> Result<Duration> { |
| 346 | let remaining = self.deadline.saturating_duration_since(Instant::now()); |
| 347 | if remaining.is_zero() { |
| 348 | bail!( |
| 349 | "MCP {} exceeded its overall {:?} deadline", |
| 350 | self.method, |
| 351 | self.overall_timeout |
| 352 | ); |
| 353 | } |
| 354 | Ok(remaining) |
| 355 | } |
| 356 | |
| 357 | fn observe_page(&mut self, page: &Value, field: &str) -> Result<Option<String>> { |
| 358 | let values = page.get(field).and_then(Value::as_array).with_context(|| { |
| 359 | format!( |
| 360 | "MCP {} response did not contain a '{field}' array", |
| 361 | self.method |
| 362 | ) |
| 363 | })?; |
| 364 | |
| 365 | self.pages = self.pages.saturating_add(1); |
| 366 | self.items = self.items.saturating_add(values.len()); |
| 367 | self.bytes = self.bytes.saturating_add(serde_json::to_vec(page)?.len()); |
| 368 | if self.pages > self.max_pages { |
| 369 | bail!( |
| 370 | "MCP {} exceeded the {}-page catalog limit", |
| 371 | self.method, |
| 372 | self.max_pages |
| 373 | ); |
| 374 | } |
| 375 | if self.items > self.max_items { |
| 376 | bail!( |
| 377 | "MCP {} exceeded the {}-item catalog limit", |
| 378 | self.method, |
| 379 | self.max_items |
| 380 | ); |
| 381 | } |
| 382 | if self.bytes > self.max_bytes { |
| 383 | bail!( |
| 384 | "MCP {} exceeded the {}-byte aggregate catalog limit", |
| 385 | self.method, |
| 386 | self.max_bytes |
| 387 | ); |
| 388 | } |
| 389 | |
| 390 | let next_cursor = match page.get("nextCursor") { |
| 391 | None => None, |
| 392 | Some(Value::String(cursor)) => Some(cursor.clone()), |
| 393 | Some(_) => bail!("MCP {} returned a non-string nextCursor", self.method), |
| 394 | }; |
| 395 | if let Some(cursor) = next_cursor.as_ref() |
| 396 | && !self.seen_cursors.insert(cursor.clone()) |
| 397 | { |
| 398 | bail!("MCP {} repeated a pagination cursor", self.method); |
| 399 | } |
| 400 | if next_cursor.is_some() && self.pages >= self.max_pages { |
| 401 | bail!( |
| 402 | "MCP {} exceeded the {}-page catalog limit", |
| 403 | self.method, |
| 404 | self.max_pages |
| 405 | ); |
| 406 | } |
| 407 | Ok(next_cursor) |
| 408 | } |
| 409 | } |
| 410 | |
| 411 | /// What the server said it supports in its `initialize` response. |
| 412 | /// |
| 413 | /// `None` means the server sent no `capabilities` object at all. Those are |
| 414 | /// treated as legacy servers and probed optimistically; an explicit |
| 415 | /// capabilities object is honoured, because a tools-only server answers |
| 416 | /// `resources/list` with a "method not found" error that would otherwise fail |
| 417 | /// the whole aggregated listing. |
| 418 | #[derive(Debug, Clone, Copy)] |
| 419 | struct ServerCapabilities { |
| 420 | tools: bool, |
| 421 | resources: bool, |
| 422 | } |
| 423 | |
| 424 | fn validate_initialize_result( |
| 425 | server_name: &str, |
| 426 | result: &Value, |
| 427 | ) -> Result<Option<ServerCapabilities>> { |
| 428 | let response = result.as_object().with_context(|| { |
| 429 | format!("MCP server '{server_name}': initialize result must be an object") |
| 430 | })?; |
| 431 | let protocol_version = response |
| 432 | .get("protocolVersion") |
| 433 | .and_then(Value::as_str) |
| 434 | .with_context(|| { |
| 435 | format!("MCP server '{server_name}': initialize result omitted protocolVersion") |
| 436 | })?; |
| 437 | // Negotiation per spec: we advertise the newest revision and accept any |
| 438 | // dated revision we still implement; anything else ends the handshake. |
| 439 | if !MCP_SUPPORTED_PROTOCOL_VERSIONS.contains(&protocol_version) { |
| 440 | bail!( |
| 441 | "MCP server '{server_name}': unsupported protocol version '{protocol_version}' (supported: {})", |
| 442 | MCP_SUPPORTED_PROTOCOL_VERSIONS.join(", ") |
| 443 | ); |
| 444 | } |
| 445 | |
| 446 | let server_info = response |
| 447 | .get("serverInfo") |
| 448 | .and_then(Value::as_object) |
| 449 | .with_context(|| { |
| 450 | format!("MCP server '{server_name}': initialize result omitted serverInfo") |
| 451 | })?; |
| 452 | for field in ["name", "version"] { |
| 453 | if !server_info.get(field).is_some_and(Value::is_string) { |
| 454 | bail!("MCP server '{server_name}': initialize serverInfo.{field} must be a string"); |
| 455 | } |
| 456 | } |
| 457 | |
| 458 | match response.get("capabilities") { |
| 459 | None => Ok(None), |
| 460 | Some(Value::Object(capabilities)) => Ok(Some(ServerCapabilities { |
| 461 | tools: capabilities.contains_key("tools"), |
| 462 | resources: capabilities.contains_key("resources"), |
| 463 | })), |
| 464 | Some(_) => bail!("MCP server '{server_name}': initialize capabilities must be an object"), |
| 465 | } |
| 466 | } |
| 467 | |
| 468 | /// A live connection to one MCP server subprocess. |
| 469 | pub struct ChildProcessMcpClient { |
| 470 | server_name: String, |
| 471 | capabilities: Option<ServerCapabilities>, |
| 472 | connection: Mutex<Connection>, |
| 473 | request_timeout: Duration, |
| 474 | } |
| 475 | |
| 476 | impl std::fmt::Debug for ChildProcessMcpClient { |
| 477 | fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { |
| 478 | f.debug_struct("ChildProcessMcpClient") |
| 479 | .field("server_name", &self.server_name) |
| 480 | .finish_non_exhaustive() |
| 481 | } |
| 482 | } |
| 483 | |
| 484 | impl ChildProcessMcpClient { |
| 485 | /// Spawn `config.command` with `config.args`/`config.env` and complete the |
| 486 | /// MCP handshake. |
| 487 | /// |
| 488 | /// Returns `Err` — never a degraded-but-usable client — when the command |
| 489 | /// cannot be executed, exits immediately, or does not answer `initialize` |
| 490 | /// within `HANDSHAKE_TIMEOUT`. |
| 491 | pub fn spawn(config: &McpServerConfig) -> Result<Self> { |
| 492 | Self::spawn_with_timeouts(config, HANDSHAKE_TIMEOUT, REQUEST_TIMEOUT) |
| 493 | } |
| 494 | |
| 495 | fn spawn_with_timeouts( |
| 496 | config: &McpServerConfig, |
| 497 | handshake_timeout: Duration, |
| 498 | request_timeout: Duration, |
| 499 | ) -> Result<Self> { |
| 500 | let server_name = config.name.clone(); |
| 501 | if config.command.trim().is_empty() { |
| 502 | bail!("MCP server '{server_name}' has no command configured"); |
| 503 | } |
| 504 | |
| 505 | let mut command = Command::new(&config.command); |
| 506 | command |
| 507 | .args(&config.args) |
| 508 | .envs(&config.env) |
| 509 | .stdin(Stdio::piped()) |
| 510 | .stdout(Stdio::piped()) |
| 511 | // The child's diagnostics belong on our stderr: stdout is the |
| 512 | // JSON-RPC channel and must not be polluted, and swallowing the |
| 513 | // child's stderr is how a misconfigured server becomes a silent |
| 514 | // one. |
| 515 | .stderr(Stdio::inherit()); |
| 516 | |
| 517 | // Own descendants too: shell/package launchers can exit before the |
| 518 | // actual server, leaving it alive with inherited protocol pipes. |
| 519 | #[cfg(unix)] |
| 520 | command.process_group(0); |
| 521 | #[cfg(windows)] |
| 522 | command.creation_flags(windows::Win32::System::Threading::CREATE_SUSPENDED.0); |
| 523 | |
| 524 | let child = command.spawn().with_context(|| { |
| 525 | format!( |
| 526 | "MCP server '{server_name}': failed to spawn command '{}'", |
| 527 | config.command |
| 528 | ) |
| 529 | })?; |
| 530 | |
| 531 | #[cfg(windows)] |
| 532 | let job = match contain_windows_child(&child) { |
| 533 | Ok(job) => job, |
| 534 | Err(error) => { |
| 535 | // It has not been allowed to run without containment. |
| 536 | let mut child = child; |
| 537 | let _ = child.kill(); |
| 538 | let _ = child.wait(); |
| 539 | return Err(error.context(format!( |
| 540 | "MCP server '{server_name}': failed to contain child" |
| 541 | ))); |
| 542 | } |
| 543 | }; |
| 544 | let (sender, responses) = sync_channel(MAX_PENDING_CHILD_MESSAGES); |
| 545 | let mut child = child; |
| 546 | let stdin = child |
| 547 | .stdin |
| 548 | .take() |
| 549 | .with_context(|| format!("MCP server '{server_name}': child stdin unavailable"))?; |
| 550 | let stdout = child |
| 551 | .stdout |
| 552 | .take() |
| 553 | .with_context(|| format!("MCP server '{server_name}': child stdout unavailable"))?; |
| 554 | let mut connection = Connection { |
| 555 | child: Some(child), |
| 556 | stdin: None, |
| 557 | responses, |
| 558 | next_id: 1, |
| 559 | #[cfg(windows)] |
| 560 | _job: job, |
| 561 | }; |
| 562 | |
| 563 | // A dedicated reader thread keeps `recv_timeout` able to bound a wait |
| 564 | // that a blocking read on the child would not. The custom line reader |
| 565 | // also caps memory before a hostile child can complete an oversized |
| 566 | // stdout line. |
| 567 | let stdin = Arc::new(Mutex::new(stdin)); |
| 568 | connection.stdin = Some(Arc::clone(&stdin)); |
| 569 | let response_stdin = Arc::downgrade(&stdin); |
| 570 | drop(stdin); |
| 571 | let response_server_name = server_name.clone(); |
| 572 | thread::spawn(move || { |
| 573 | let mut reader = BufReader::new(stdout); |
| 574 | loop { |
| 575 | match read_bounded_line(&mut reader, MAX_JSONRPC_LINE_BYTES) { |
| 576 | Ok(Some(line)) => { |
| 577 | // Server requests must be answered even while no |
| 578 | // client request is in flight. Route them before the |
| 579 | // zero-capacity response channel: blocking there until |
| 580 | // the next client call would violate MCP ping's prompt |
| 581 | // response requirement. Valid notifications need no |
| 582 | // response and likewise must not occupy the rendezvous. |
| 583 | if let Ok(message) = serde_json::from_str::<Value>(&line) |
| 584 | && message.get("jsonrpc").and_then(Value::as_str) == Some("2.0") |
| 585 | && message.get("method").and_then(Value::as_str).is_some() |
| 586 | { |
| 587 | let Some(response) = response_to_server_request(&message) else { |
| 588 | continue; |
| 589 | }; |
| 590 | let Some(stdin) = response_stdin.upgrade() else { |
| 591 | break; |
| 592 | }; |
| 593 | let result = stdin |
| 594 | .lock() |
| 595 | .map_err(|_| { |
| 596 | anyhow!("connection stdin poisoned by an earlier panic") |
| 597 | }) |
| 598 | .and_then(|mut stdin| write_jsonrpc_line(&mut *stdin, &response)); |
| 599 | // Never retain the temporary strong handle while a |
| 600 | // failure waits on the rendezvous channel. Drop |
| 601 | // must remain able to close child stdin promptly. |
| 602 | drop(stdin); |
| 603 | if let Err(error) = result { |
| 604 | let _ = sender.send(ChildStdoutMessage::Invalid(format!( |
| 605 | "MCP server '{response_server_name}': failed to answer idle child request: {error:#}" |
| 606 | ))); |
| 607 | break; |
| 608 | } |
| 609 | continue; |
| 610 | } |
| 611 | if sender.send(ChildStdoutMessage::Line(line)).is_err() { |
| 612 | break; |
| 613 | } |
| 614 | } |
| 615 | Ok(None) => break, |
| 616 | Err(err) => { |
| 617 | let _ = sender.send(ChildStdoutMessage::Invalid(err.to_string())); |
| 618 | break; |
| 619 | } |
| 620 | } |
| 621 | } |
| 622 | }); |
| 623 | |
| 624 | let initialize = connection.request( |
| 625 | &server_name, |
| 626 | "initialize", |
| 627 | json!({ |
| 628 | "protocolVersion": MCP_PROTOCOL_VERSION, |
| 629 | "clientInfo": { |
| 630 | "name": "codewhale-mcp-server", |
| 631 | "version": env!("CARGO_PKG_VERSION") |
| 632 | }, |
| 633 | "capabilities": { |
| 634 | "tools": {}, |
| 635 | "resources": {} |
| 636 | } |
| 637 | }), |
| 638 | handshake_timeout, |
| 639 | )?; |
| 640 | let capabilities = validate_initialize_result(&server_name, &initialize)?; |
| 641 | |
| 642 | connection |
| 643 | .send(&json!({ |
| 644 | "jsonrpc": "2.0", |
| 645 | "method": "notifications/initialized" |
| 646 | })) |
| 647 | .with_context(|| { |
| 648 | format!("MCP server '{server_name}': failed to confirm initialization") |
| 649 | })?; |
| 650 | |
| 651 | Ok(Self { |
| 652 | server_name, |
| 653 | capabilities, |
| 654 | connection: Mutex::new(connection), |
| 655 | request_timeout, |
| 656 | }) |
| 657 | } |
| 658 | |
| 659 | fn supports_tools(&self) -> bool { |
| 660 | self.capabilities.is_none_or(|caps| caps.tools) |
| 661 | } |
| 662 | |
| 663 | fn supports_resources(&self) -> bool { |
| 664 | self.capabilities.is_none_or(|caps| caps.resources) |
| 665 | } |
| 666 | |
| 667 | fn request(&self, method: &str, params: Value) -> Result<Value> { |
| 668 | self.request_with_timeout(method, params, self.request_timeout) |
| 669 | } |
| 670 | |
| 671 | fn request_with_timeout( |
| 672 | &self, |
| 673 | method: &str, |
| 674 | params: Value, |
| 675 | timeout: Duration, |
| 676 | ) -> Result<Value> { |
| 677 | let mut connection = self.connection.lock().map_err(|_| { |
| 678 | anyhow!( |
| 679 | "MCP server '{}': connection poisoned by an earlier panic", |
| 680 | self.server_name |
| 681 | ) |
| 682 | })?; |
| 683 | connection.request(&self.server_name, method, params, timeout) |
| 684 | } |
| 685 | |
| 686 | /// Drive a paginated `*/list` method to exhaustion, validating every item |
| 687 | /// before requesting the next page. A malformed page therefore fails at |
| 688 | /// its source rather than being silently thinned or discovered only after |
| 689 | /// the rest of the catalog has been fetched. |
| 690 | fn list_paginated<T, F>(&self, method: &str, field: &str, mut parse: F) -> Result<Vec<T>> |
| 691 | where |
| 692 | F: FnMut(&Value, usize) -> Result<T>, |
| 693 | { |
| 694 | let mut items = Vec::new(); |
| 695 | let mut cursor: Option<String> = None; |
| 696 | let mut budget = ListBudget::new(method, self.request_timeout); |
| 697 | loop { |
| 698 | let params = match &cursor { |
| 699 | Some(cursor) => json!({ "cursor": cursor }), |
| 700 | None => json!({}), |
| 701 | }; |
| 702 | let page = self.request_with_timeout(method, params, budget.remaining_timeout()?)?; |
| 703 | let next_cursor = budget.observe_page(&page, field)?; |
| 704 | let values = page |
| 705 | .get(field) |
| 706 | .and_then(Value::as_array) |
| 707 | .expect("ListBudget validated the catalog field"); |
| 708 | for value in values { |
| 709 | let index = items.len(); |
| 710 | items.push(parse(value, index)?); |
| 711 | } |
| 712 | match next_cursor { |
| 713 | Some(next_cursor) => cursor = Some(next_cursor), |
| 714 | None => break, |
| 715 | } |
| 716 | } |
| 717 | Ok(items) |
| 718 | } |
| 719 | } |
| 720 | |
| 721 | impl McpManagedClient for ChildProcessMcpClient { |
| 722 | fn list_tools(&self) -> Result<Vec<McpToolDescriptor>> { |
| 723 | Ok(self |
| 724 | .list_tools_with_input_schemas()? |
| 725 | .into_iter() |
| 726 | .map(|(tool, _)| tool) |
| 727 | .collect()) |
| 728 | } |
| 729 | |
| 730 | fn list_tools_with_input_schemas(&self) -> Result<Vec<(McpToolDescriptor, Value)>> { |
| 731 | if !self.supports_tools() { |
| 732 | return Ok(Vec::new()); |
| 733 | } |
| 734 | let allow_legacy_schema_omission = self.capabilities.is_none(); |
| 735 | self.list_paginated("tools/list", "tools", |tool, index| { |
| 736 | parse_tool_entry(&self.server_name, tool, index, allow_legacy_schema_omission) |
| 737 | }) |
| 738 | } |
| 739 | |
| 740 | fn call_tool(&self, tool_name: &str, arguments: Value) -> Result<Value> { |
| 741 | // The server's result is returned verbatim, including an `isError` |
| 742 | // content payload: reinterpreting it here would replace what the |
| 743 | // server actually said with our guess about it. |
| 744 | self.request( |
| 745 | "tools/call", |
| 746 | json!({ |
| 747 | "name": tool_name, |
| 748 | "arguments": arguments |
| 749 | }), |
| 750 | ) |
| 751 | } |
| 752 | |
| 753 | fn list_resources(&self) -> Result<Vec<McpResourceDescriptor>> { |
| 754 | Ok(self |
| 755 | .list_resources_with_metadata()? |
| 756 | .into_iter() |
| 757 | .map(|(resource, _)| resource) |
| 758 | .collect()) |
| 759 | } |
| 760 | |
| 761 | fn list_resources_with_metadata(&self) -> Result<Vec<(McpResourceDescriptor, Value)>> { |
| 762 | if !self.supports_resources() { |
| 763 | return Ok(Vec::new()); |
| 764 | } |
| 765 | self.list_paginated("resources/list", "resources", |resource, index| { |
| 766 | parse_resource_entry(&self.server_name, resource, index) |
| 767 | }) |
| 768 | } |
| 769 | |
| 770 | fn read_resource(&self, uri: &str) -> Result<Value> { |
| 771 | self.request("resources/read", json!({ "uri": uri })) |
| 772 | } |
| 773 | } |
| 774 | |
| 775 | struct Connection { |
| 776 | /// `None` once `Drop` hands the child to the reaper thread. |
| 777 | child: Option<Child>, |
| 778 | stdin: Option<Arc<Mutex<ChildStdin>>>, |
| 779 | responses: Receiver<ChildStdoutMessage>, |
| 780 | next_id: u64, |
| 781 | #[cfg(windows)] |
| 782 | _job: OwnedHandle, |
| 783 | } |
| 784 | |
| 785 | /// Use the same suspended-spawn, kill-on-close Job Object ownership as hook |
| 786 | /// and shell children. The TUI depends on this crate, so its private guards |
| 787 | /// cannot be imported here without creating a dependency cycle. |
| 788 | #[cfg(windows)] |
| 789 | fn contain_windows_child(child: &Child) -> Result<OwnedHandle> { |
| 790 | use windows::Win32::Foundation::HANDLE; |
| 791 | use windows::Win32::System::Diagnostics::ToolHelp::{ |
| 792 | CreateToolhelp32Snapshot, TH32CS_SNAPTHREAD, THREADENTRY32, Thread32First, Thread32Next, |
| 793 | }; |
| 794 | use windows::Win32::System::JobObjects::{ |
| 795 | AssignProcessToJobObject, CreateJobObjectW, JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE, |
| 796 | JOBOBJECT_EXTENDED_LIMIT_INFORMATION, JobObjectExtendedLimitInformation, |
| 797 | SetInformationJobObject, |
| 798 | }; |
| 799 | use windows::Win32::System::Threading::{OpenThread, ResumeThread, THREAD_SUSPEND_RESUME}; |
| 800 | use windows::core::PCWSTR; |
| 801 | |
| 802 | unsafe { |
| 803 | let handle = CreateJobObjectW(None, PCWSTR::null())?; |
| 804 | let job = OwnedHandle::from_raw_handle(handle.0); |
| 805 | let mut limits = JOBOBJECT_EXTENDED_LIMIT_INFORMATION::default(); |
| 806 | limits.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE; |
| 807 | SetInformationJobObject( |
| 808 | handle, |
| 809 | JobObjectExtendedLimitInformation, |
| 810 | &limits as *const _ as *const core::ffi::c_void, |
| 811 | std::mem::size_of_val(&limits) as u32, |
| 812 | )?; |
| 813 | AssignProcessToJobObject(handle, HANDLE(child.as_raw_handle()))?; |
| 814 | |
| 815 | let snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0)?; |
| 816 | let snapshot = OwnedHandle::from_raw_handle(snapshot.0); |
| 817 | let mut entry = THREADENTRY32 { |
| 818 | dwSize: std::mem::size_of::<THREADENTRY32>() as u32, |
| 819 | ..Default::default() |
| 820 | }; |
| 821 | let mut next = Thread32First(HANDLE(snapshot.as_raw_handle()), &mut entry); |
| 822 | let mut resumed = 0usize; |
| 823 | while next.is_ok() { |
| 824 | if entry.th32OwnerProcessID == child.id() { |
| 825 | let thread = OpenThread(THREAD_SUSPEND_RESUME, false, entry.th32ThreadID)?; |
| 826 | let thread = OwnedHandle::from_raw_handle(thread.0); |
| 827 | if ResumeThread(HANDLE(thread.as_raw_handle())) == u32::MAX { |
| 828 | return Err(io::Error::last_os_error().into()); |
| 829 | } |
| 830 | resumed += 1; |
| 831 | } |
| 832 | next = Thread32Next(HANDLE(snapshot.as_raw_handle()), &mut entry); |
| 833 | } |
| 834 | if resumed == 0 { |
| 835 | bail!("suspended MCP child had no resumable thread"); |
| 836 | } |
| 837 | Ok(job) |
| 838 | } |
| 839 | } |
| 840 | |
| 841 | impl Connection { |
| 842 | fn send(&mut self, message: &Value) -> Result<()> { |
| 843 | let stdin = self |
| 844 | .stdin |
| 845 | .as_ref() |
| 846 | .context("connection stdin already closed")?; |
| 847 | let mut stdin = stdin |
| 848 | .lock() |
| 849 | .map_err(|_| anyhow!("connection stdin poisoned by an earlier panic"))?; |
| 850 | write_jsonrpc_line(&mut *stdin, message) |
| 851 | } |
| 852 | |
| 853 | fn request( |
| 854 | &mut self, |
| 855 | server: &str, |
| 856 | method: &str, |
| 857 | params: Value, |
| 858 | timeout: Duration, |
| 859 | ) -> Result<Value> { |
| 860 | let id = self.next_id; |
| 861 | self.next_id += 1; |
| 862 | if let Err(err) = self.send(&json!({ |
| 863 | "jsonrpc": "2.0", |
| 864 | "id": id, |
| 865 | "method": method, |
| 866 | "params": params |
| 867 | })) { |
| 868 | // A child that has already exited leaves us racing two symptoms of |
| 869 | // the same fact: either the reader thread sees EOF first, or our |
| 870 | // write loses the race and returns EPIPE. Which one wins is |
| 871 | // platform- and timing-dependent (macOS reliably reports the write |
| 872 | // error where Linux reports the EOF), so both report the death the |
| 873 | // same way rather than leaking a bare "Broken pipe". |
| 874 | if is_broken_pipe(&err) { |
| 875 | bail!( |
| 876 | "MCP server '{server}': process closed stdin before answering {method}{}", |
| 877 | self.exit_note() |
| 878 | ); |
| 879 | } |
| 880 | return Err(err) |
| 881 | .with_context(|| format!("MCP server '{server}': failed to send {method}")); |
| 882 | } |
| 883 | |
| 884 | let deadline = Instant::now() + timeout; |
| 885 | loop { |
| 886 | let remaining = deadline.saturating_duration_since(Instant::now()); |
| 887 | if remaining.is_zero() { |
| 888 | bail!("MCP server '{server}': {method} timed out after {timeout:?}"); |
| 889 | } |
| 890 | let line = match self.responses.recv_timeout(remaining) { |
| 891 | Ok(ChildStdoutMessage::Line(line)) => line, |
| 892 | Ok(ChildStdoutMessage::Invalid(error)) => { |
| 893 | bail!( |
| 894 | "MCP server '{server}': invalid child stdout while awaiting {method}: {error}" |
| 895 | ); |
| 896 | } |
| 897 | Err(RecvTimeoutError::Timeout) => { |
| 898 | bail!("MCP server '{server}': {method} timed out after {timeout:?}"); |
| 899 | } |
| 900 | Err(RecvTimeoutError::Disconnected) => { |
| 901 | bail!( |
| 902 | "MCP server '{server}': process closed stdout before answering {method}{}", |
| 903 | self.exit_note() |
| 904 | ); |
| 905 | } |
| 906 | }; |
| 907 | |
| 908 | // Servers occasionally emit banners or log lines on stdout. They |
| 909 | // are skipped; only the matching response ends the wait. Valid |
| 910 | // requests and notifications are handled by the reader before |
| 911 | // they reach this rendezvous. |
| 912 | let Ok(message) = serde_json::from_str::<Value>(&line) else { |
| 913 | continue; |
| 914 | }; |
| 915 | if message.get("jsonrpc").and_then(Value::as_str) != Some("2.0") { |
| 916 | bail!( |
| 917 | "MCP server '{server}': {method} received a child message without jsonrpc \"2.0\"" |
| 918 | ); |
| 919 | } |
| 920 | if message.get("id").and_then(Value::as_u64) != Some(id) { |
| 921 | continue; |
| 922 | } |
| 923 | if let Some(error) = message.get("error") { |
| 924 | bail!("MCP server '{server}': {method} failed: {error}"); |
| 925 | } |
| 926 | return message.get("result").cloned().with_context(|| { |
| 927 | format!( |
| 928 | "MCP server '{server}': {method} response contained neither result nor error" |
| 929 | ) |
| 930 | }); |
| 931 | } |
| 932 | } |
| 933 | |
| 934 | /// The child's exit status, when it has one, for appending to a failure |
| 935 | /// message. Polls briefly because both callers run at the moment the child |
| 936 | /// is dying: the write can return EPIPE, or stdout can hit EOF, before the |
| 937 | /// kernel has finished reaping the process. Bounded and error-path-only, |
| 938 | /// so the cost buys a real diagnostic ("exited with status 127" is the |
| 939 | /// difference between a crashed server and a missing one). |
| 940 | fn exit_note(&mut self) -> String { |
| 941 | let Some(child) = self.child.as_mut() else { |
| 942 | return String::new(); |
| 943 | }; |
| 944 | let deadline = Instant::now() + EXIT_STATUS_GRACE; |
| 945 | loop { |
| 946 | match child.try_wait() { |
| 947 | Ok(Some(status)) => return format!(" (process exited with {status})"), |
| 948 | Ok(None) if Instant::now() < deadline => { |
| 949 | thread::sleep(Duration::from_millis(5)); |
| 950 | } |
| 951 | _ => return String::new(), |
| 952 | } |
| 953 | } |
| 954 | } |
| 955 | } |
| 956 | |
| 957 | /// Whether an error chain bottoms out in a broken-pipe I/O error, i.e. we wrote |
| 958 | /// to a child that had already closed its end. |
| 959 | fn is_broken_pipe(err: &anyhow::Error) -> bool { |
| 960 | err.chain().any(|cause| { |
| 961 | cause |
| 962 | .downcast_ref::<std::io::Error>() |
| 963 | .is_some_and(|io| io.kind() == std::io::ErrorKind::BrokenPipe) |
| 964 | }) |
| 965 | } |
| 966 | |
| 967 | /// Reap stdio children off the dropping thread (#6211 R4). `Drop` runs on |
| 968 | /// whichever thread drops the client — manager reload, pool rebuild — and |
| 969 | /// the half-second grace wait must not stall it, least of all an executor |
| 970 | /// thread. One shared thread reaps every child; drops only send. What this |
| 971 | /// does not do: join the reaper at process exit, so a child that ignores |
| 972 | /// EOF can outlive a racing shutdown where the old synchronous `Drop` |
| 973 | /// would have killed it first. |
| 974 | fn child_reaper() -> &'static std::sync::mpsc::Sender<Child> { |
| 975 | static REAPER: std::sync::OnceLock<std::sync::mpsc::Sender<Child>> = std::sync::OnceLock::new(); |
| 976 | REAPER.get_or_init(|| { |
| 977 | let (tx, rx) = std::sync::mpsc::channel::<Child>(); |
| 978 | thread::Builder::new() |
| 979 | .name("mcp-stdio-reaper".to_string()) |
| 980 | .spawn(move || { |
| 981 | for mut child in rx { |
| 982 | reap_child(&mut child); |
| 983 | } |
| 984 | }) |
| 985 | .expect("MCP stdio reaper thread spawns"); |
| 986 | tx |
| 987 | }) |
| 988 | } |
| 989 | |
| 990 | /// Grace-wait a child, then kill what ignores stdin-close. Runs on the |
| 991 | /// reaper thread, or inline in `Drop` if the reaper is gone. |
| 992 | fn reap_child(child: &mut Child) { |
| 993 | let deadline = Instant::now() + SHUTDOWN_GRACE; |
| 994 | loop { |
| 995 | match child.try_wait() { |
| 996 | Ok(Some(_)) => break, |
| 997 | Ok(None) if Instant::now() < deadline => { |
| 998 | thread::sleep(Duration::from_millis(10)); |
| 999 | } |
| 1000 | _ => break, |
| 1001 | } |
| 1002 | } |
| 1003 | #[cfg(unix)] |
| 1004 | unsafe { |
| 1005 | // Also run after the immediate launcher exited. Descendants can |
| 1006 | // still own the group and the stdout/stderr pipe descriptors. |
| 1007 | let _ = libc::kill(-(child.id() as libc::pid_t), libc::SIGKILL); |
| 1008 | } |
| 1009 | let _ = child.kill(); |
| 1010 | let _ = child.wait(); |
| 1011 | } |
| 1012 | |
| 1013 | impl Drop for Connection { |
| 1014 | fn drop(&mut self) { |
| 1015 | // Closing stdin is the protocol-level shutdown signal for a stdio MCP |
| 1016 | // server; kill only the ones that ignore it, so servers get a chance |
| 1017 | // to flush state. |
| 1018 | self.stdin.take(); |
| 1019 | // The grace wait leaves the dropping thread: the child is reaped on |
| 1020 | // the shared reaper thread. If the reaper itself is gone, reap |
| 1021 | // inline — today's behavior — rather than leak the child. |
| 1022 | if let Some(child) = self.child.take() |
| 1023 | && let Err(mut failed) = child_reaper().send(child) |
| 1024 | { |
| 1025 | reap_child(&mut failed.0); |
| 1026 | } |
| 1027 | } |
| 1028 | } |
| 1029 | |
| 1030 | #[cfg(test)] |
| 1031 | mod tests { |
| 1032 | use std::collections::HashMap; |
| 1033 | use std::io::Cursor; |
| 1034 | |
| 1035 | use super::*; |
| 1036 | |
| 1037 | fn config(command: &str, args: &[&str]) -> McpServerConfig { |
| 1038 | McpServerConfig { |
| 1039 | name: "probe".to_string(), |
| 1040 | command: command.to_string(), |
| 1041 | args: args.iter().map(|arg| (*arg).to_string()).collect(), |
| 1042 | env: HashMap::new(), |
| 1043 | enabled: true, |
| 1044 | } |
| 1045 | } |
| 1046 | |
| 1047 | #[test] |
| 1048 | fn spawn_fails_loudly_when_the_command_does_not_exist() { |
| 1049 | let err = ChildProcessMcpClient::spawn(&config( |
| 1050 | "codewhale-nonexistent-mcp-server-binary", |
| 1051 | &["--stdio"], |
| 1052 | )) |
| 1053 | .unwrap_err(); |
| 1054 | let message = format!("{err:#}"); |
| 1055 | assert!( |
| 1056 | message.contains("failed to spawn command"), |
| 1057 | "spawn failure must name the command, got: {message}" |
| 1058 | ); |
| 1059 | assert!( |
| 1060 | message.contains("probe"), |
| 1061 | "spawn failure must name the server, got: {message}" |
| 1062 | ); |
| 1063 | } |
| 1064 | |
| 1065 | #[test] |
| 1066 | fn spawn_rejects_an_empty_command() { |
| 1067 | let err = ChildProcessMcpClient::spawn(&config(" ", &[])).unwrap_err(); |
| 1068 | assert!( |
| 1069 | format!("{err:#}").contains("no command configured"), |
| 1070 | "unexpected error: {err:#}" |
| 1071 | ); |
| 1072 | } |
| 1073 | |
| 1074 | #[test] |
| 1075 | fn standard_tool_input_schemas_are_required_and_fail_closed_when_malformed() { |
| 1076 | let fallback = json!({"type": "object", "properties": {}}); |
| 1077 | let (_, legacy_schema) = |
| 1078 | parse_tool_entry("legacy", &json!({"name": "legacy-tool"}), 0, true).unwrap(); |
| 1079 | assert_eq!(legacy_schema, fallback); |
| 1080 | |
| 1081 | for malformed_tool in [ |
| 1082 | json!({"name": "missing"}), |
| 1083 | json!({"name": "null", "inputSchema": null}), |
| 1084 | json!({"name": "empty", "inputSchema": {}}), |
| 1085 | json!({"name": "wrong-type", "inputSchema": {"type": "string"}}), |
| 1086 | json!({ |
| 1087 | "name": "bad-properties", |
| 1088 | "inputSchema": {"type": "object", "properties": "not-an-object"} |
| 1089 | }), |
| 1090 | json!({ |
| 1091 | "name": "bad-property", |
| 1092 | "inputSchema": {"type": "object", "properties": {"path": true}} |
| 1093 | }), |
| 1094 | json!({ |
| 1095 | "name": "bad-required-shape", |
| 1096 | "inputSchema": {"type": "object", "required": "path"} |
| 1097 | }), |
| 1098 | json!({ |
| 1099 | "name": "bad-required-entry", |
| 1100 | "inputSchema": {"type": "object", "required": [1]} |
| 1101 | }), |
| 1102 | ] { |
| 1103 | let error = parse_tool_entry("standard", &malformed_tool, 0, false).unwrap_err(); |
| 1104 | assert!( |
| 1105 | error.to_string().contains("inputSchema"), |
| 1106 | "malformed schema produced an unrelated error: {error:#}" |
| 1107 | ); |
| 1108 | } |
| 1109 | |
| 1110 | // Legacy mode tolerates only omission. Once a field is present, it is |
| 1111 | // standard data and must pass the same validation. |
| 1112 | let legacy_error = parse_tool_entry( |
| 1113 | "legacy", |
| 1114 | &json!({"name": "malformed", "inputSchema": null}), |
| 1115 | 0, |
| 1116 | true, |
| 1117 | ) |
| 1118 | .unwrap_err(); |
| 1119 | assert!(legacy_error.to_string().contains("inputSchema")); |
| 1120 | } |
| 1121 | |
| 1122 | #[test] |
| 1123 | fn catalog_entry_parsers_preserve_valid_fields_and_reject_malformed_entries() { |
| 1124 | let (tool, schema) = parse_tool_entry( |
| 1125 | "fixture", |
| 1126 | &json!({ |
| 1127 | "name": "read-file", |
| 1128 | "description": "Read one file", |
| 1129 | "inputSchema": { |
| 1130 | "type": "object", |
| 1131 | "properties": {"path": {"type": "string"}}, |
| 1132 | "required": ["path"] |
| 1133 | } |
| 1134 | }), |
| 1135 | 0, |
| 1136 | false, |
| 1137 | ) |
| 1138 | .unwrap(); |
| 1139 | assert_eq!(tool.tool_name, "read-file"); |
| 1140 | assert_eq!(tool.description.as_deref(), Some("Read one file")); |
| 1141 | assert_eq!(schema["required"], json!(["path"])); |
| 1142 | |
| 1143 | let (resource, metadata) = parse_resource_entry( |
| 1144 | "fixture", |
| 1145 | &json!({ |
| 1146 | "uri": "file:///guide.md", |
| 1147 | "name": "Guide", |
| 1148 | "description": "User guide", |
| 1149 | "mimeType": "text/markdown", |
| 1150 | "size": 42, |
| 1151 | "annotations": {"audience": ["assistant"], "priority": 0.8} |
| 1152 | }), |
| 1153 | 0, |
| 1154 | ) |
| 1155 | .unwrap(); |
| 1156 | assert_eq!(resource.uri, "file:///guide.md"); |
| 1157 | assert_eq!(resource.description.as_deref(), Some("User guide")); |
| 1158 | assert_eq!( |
| 1159 | metadata, |
| 1160 | json!({ |
| 1161 | "name": "Guide", |
| 1162 | "mimeType": "text/markdown", |
| 1163 | "size": 42, |
| 1164 | "annotations": {"audience": ["assistant"], "priority": 0.8} |
| 1165 | }) |
| 1166 | ); |
| 1167 | |
| 1168 | for malformed in [ |
| 1169 | Value::Null, |
| 1170 | json!({}), |
| 1171 | json!({"name": 7}), |
| 1172 | json!({"name": "tool", "description": []}), |
| 1173 | ] { |
| 1174 | assert!( |
| 1175 | parse_tool_entry("fixture", &malformed, 3, false).is_err(), |
| 1176 | "malformed tool entry passed: {malformed}" |
| 1177 | ); |
| 1178 | } |
| 1179 | for malformed in [ |
| 1180 | Value::Null, |
| 1181 | json!({}), |
| 1182 | json!({"uri": "file:///x"}), |
| 1183 | json!({"uri": 7, "name": "x"}), |
| 1184 | json!({"uri": "file:///x", "name": 7}), |
| 1185 | json!({"uri": "file:///x", "name": "x", "description": []}), |
| 1186 | json!({"uri": "file:///x", "name": "x", "mimeType": []}), |
| 1187 | json!({"uri": "file:///x", "name": "x", "size": 1.5}), |
| 1188 | json!({"uri": "file:///x", "name": "x", "annotations": []}), |
| 1189 | ] { |
| 1190 | assert!( |
| 1191 | parse_resource_entry("fixture", &malformed, 4).is_err(), |
| 1192 | "malformed resource entry passed: {malformed}" |
| 1193 | ); |
| 1194 | } |
| 1195 | } |
| 1196 | |
| 1197 | #[test] |
| 1198 | fn bounded_line_reader_rejects_an_oversized_line_before_completion() { |
| 1199 | let mut exact = BufReader::with_capacity(3, Cursor::new(b"12345678\nnext\n")); |
| 1200 | assert_eq!( |
| 1201 | read_bounded_line(&mut exact, 8).unwrap().as_deref(), |
| 1202 | Some("12345678") |
| 1203 | ); |
| 1204 | assert_eq!( |
| 1205 | read_bounded_line(&mut exact, 8).unwrap().as_deref(), |
| 1206 | Some("next") |
| 1207 | ); |
| 1208 | |
| 1209 | // The three-byte BufReader forces the bound check to span chunks. It |
| 1210 | // rejects when the ninth payload byte is buffered, before consuming |
| 1211 | // the newline or growing the retained line beyond eight bytes. |
| 1212 | let mut oversized = BufReader::with_capacity(3, Cursor::new(b"123456789\n")); |
| 1213 | let error = read_bounded_line(&mut oversized, 8).unwrap_err(); |
| 1214 | assert_eq!(error.kind(), io::ErrorKind::InvalidData); |
| 1215 | assert!(error.to_string().contains("8-byte limit")); |
| 1216 | } |
| 1217 | |
| 1218 | #[test] |
| 1219 | fn child_stdout_queue_is_rendezvous_backpressured() { |
| 1220 | let (sender, receiver) = sync_channel(MAX_PENDING_CHILD_MESSAGES); |
| 1221 | match sender.try_send(ChildStdoutMessage::Line("one".to_string())) { |
| 1222 | Err(std::sync::mpsc::TrySendError::Full(_)) => {} |
| 1223 | Err(std::sync::mpsc::TrySendError::Disconnected(_)) => { |
| 1224 | panic!("child stdout receiver disconnected") |
| 1225 | } |
| 1226 | Ok(()) => panic!("child stdout queue accepted a message without a waiting consumer"), |
| 1227 | } |
| 1228 | drop(receiver); |
| 1229 | } |
| 1230 | |
| 1231 | #[test] |
| 1232 | fn initialize_result_requires_supported_protocol_and_server_identity() { |
| 1233 | let valid = json!({ |
| 1234 | "protocolVersion": MCP_PROTOCOL_VERSION, |
| 1235 | "serverInfo": {"name": "fixture", "version": "1"}, |
| 1236 | "capabilities": {"tools": {}, "resources": {}} |
| 1237 | }); |
| 1238 | let capabilities = validate_initialize_result("fixture", &valid) |
| 1239 | .unwrap() |
| 1240 | .expect("capabilities"); |
| 1241 | assert!(capabilities.tools); |
| 1242 | assert!(capabilities.resources); |
| 1243 | |
| 1244 | // Negotiation accepts every dated revision still implemented, not only |
| 1245 | // the newest one advertised at initialize. |
| 1246 | for version in ["2025-03-26", "2024-11-05"] { |
| 1247 | let older = json!({ |
| 1248 | "protocolVersion": version, |
| 1249 | "serverInfo": {"name": "fixture", "version": "1"}, |
| 1250 | "capabilities": {"tools": {}} |
| 1251 | }); |
| 1252 | assert!( |
| 1253 | validate_initialize_result("fixture", &older).is_ok(), |
| 1254 | "supported revision {version} was rejected" |
| 1255 | ); |
| 1256 | } |
| 1257 | |
| 1258 | for invalid in [ |
| 1259 | json!({}), |
| 1260 | json!({ |
| 1261 | "protocolVersion": "2099-01-01", |
| 1262 | "serverInfo": {"name": "fixture", "version": "1"} |
| 1263 | }), |
| 1264 | json!({ |
| 1265 | "protocolVersion": MCP_PROTOCOL_VERSION, |
| 1266 | "serverInfo": {"name": "fixture"} |
| 1267 | }), |
| 1268 | json!({ |
| 1269 | "protocolVersion": MCP_PROTOCOL_VERSION, |
| 1270 | "serverInfo": {"name": "fixture", "version": "1"}, |
| 1271 | "capabilities": [] |
| 1272 | }), |
| 1273 | ] { |
| 1274 | assert!( |
| 1275 | validate_initialize_result("fixture", &invalid).is_err(), |
| 1276 | "invalid initialize result passed: {invalid}" |
| 1277 | ); |
| 1278 | } |
| 1279 | } |
| 1280 | |
| 1281 | #[test] |
| 1282 | fn catalog_budget_fails_instead_of_returning_partial_results() { |
| 1283 | let mut pages = |
| 1284 | ListBudget::with_limits("resources/list", Duration::from_secs(1), 2, 2, 256); |
| 1285 | assert_eq!( |
| 1286 | pages |
| 1287 | .observe_page( |
| 1288 | &json!({"resources": [{"uri": "file:///a", "name": "a"}], "nextCursor": "a"}), |
| 1289 | "resources" |
| 1290 | ) |
| 1291 | .unwrap(), |
| 1292 | Some("a".to_string()) |
| 1293 | ); |
| 1294 | let page_error = pages |
| 1295 | .observe_page( |
| 1296 | &json!({"resources": [{"uri": "file:///b", "name": "b"}], "nextCursor": "b"}), |
| 1297 | "resources", |
| 1298 | ) |
| 1299 | .unwrap_err(); |
| 1300 | assert!(page_error.to_string().contains("page catalog limit")); |
| 1301 | |
| 1302 | let mut items = ListBudget::with_limits("tools/list", Duration::from_secs(1), 4, 1, 256); |
| 1303 | let item_error = items |
| 1304 | .observe_page(&json!({"tools": [{}, {}]}), "tools") |
| 1305 | .unwrap_err(); |
| 1306 | assert!(item_error.to_string().contains("item catalog limit")); |
| 1307 | |
| 1308 | let mut bytes = ListBudget::with_limits("tools/list", Duration::from_secs(1), 4, 4, 8); |
| 1309 | let byte_error = bytes |
| 1310 | .observe_page(&json!({"tools": [{"name": "large"}]}), "tools") |
| 1311 | .unwrap_err(); |
| 1312 | assert!( |
| 1313 | byte_error |
| 1314 | .to_string() |
| 1315 | .contains("byte aggregate catalog limit") |
| 1316 | ); |
| 1317 | |
| 1318 | let mut malformed = |
| 1319 | ListBudget::with_limits("resources/list", Duration::from_secs(1), 4, 4, 256); |
| 1320 | assert!( |
| 1321 | malformed |
| 1322 | .observe_page(&json!({"resources": null}), "resources") |
| 1323 | .unwrap_err() |
| 1324 | .to_string() |
| 1325 | .contains("did not contain a 'resources' array") |
| 1326 | ); |
| 1327 | |
| 1328 | let expired = ListBudget::new("resources/list", Duration::ZERO); |
| 1329 | assert!( |
| 1330 | expired |
| 1331 | .remaining_timeout() |
| 1332 | .unwrap_err() |
| 1333 | .to_string() |
| 1334 | .contains("overall") |
| 1335 | ); |
| 1336 | } |
| 1337 | |
| 1338 | #[cfg(unix)] |
| 1339 | #[test] |
| 1340 | fn spawn_fails_when_the_child_exits_without_answering_initialize() { |
| 1341 | let err = ChildProcessMcpClient::spawn(&config("/bin/sh", &["-c", "exit 0"])).unwrap_err(); |
| 1342 | let message = format!("{err:#}"); |
| 1343 | // Which end reports the death first is a race the OS arbitrates — |
| 1344 | // stdout EOF on Linux, an EPIPE write on macOS — so assert on what is |
| 1345 | // actually contractual: the server is named, the failure is attributed |
| 1346 | // to the child dying before it answered, and no raw io::Error leaks. |
| 1347 | assert!( |
| 1348 | message.contains("probe") && message.contains("before answering initialize"), |
| 1349 | "unexpected error: {message}" |
| 1350 | ); |
| 1351 | assert!( |
| 1352 | !message.contains("Broken pipe"), |
| 1353 | "a dead child must not surface as a raw pipe error: {message}" |
| 1354 | ); |
| 1355 | } |
| 1356 | |
| 1357 | #[cfg(unix)] |
| 1358 | #[test] |
| 1359 | fn a_child_that_dies_mid_handshake_reports_its_exit_status() { |
| 1360 | // The child closes both pipes and exits nonzero: the diagnostic has to |
| 1361 | // carry the status, because "exited with 127" is what distinguishes a |
| 1362 | // crashed server from a missing one. |
| 1363 | let err = |
| 1364 | ChildProcessMcpClient::spawn(&config("/bin/sh", &["-c", "exit 127"])).unwrap_err(); |
| 1365 | let message = format!("{err:#}"); |
| 1366 | assert!( |
| 1367 | message.contains("before answering initialize") && message.contains("127"), |
| 1368 | "unexpected error: {message}" |
| 1369 | ); |
| 1370 | } |
| 1371 | |
| 1372 | #[cfg(unix)] |
| 1373 | #[test] |
| 1374 | fn handshake_times_out_when_the_child_never_answers() { |
| 1375 | let err = ChildProcessMcpClient::spawn_with_timeouts( |
| 1376 | &config("/bin/sh", &["-c", "sleep 30"]), |
| 1377 | Duration::from_millis(250), |
| 1378 | Duration::from_millis(250), |
| 1379 | ) |
| 1380 | .unwrap_err(); |
| 1381 | assert!( |
| 1382 | format!("{err:#}").contains("initialize timed out"), |
| 1383 | "unexpected error: {err:#}" |
| 1384 | ); |
| 1385 | } |
| 1386 | |
| 1387 | #[cfg(unix)] |
| 1388 | #[test] |
| 1389 | fn reaper_kills_a_child_that_ignores_stdin_close() { |
| 1390 | let mut child = Command::new("/bin/sh") |
| 1391 | .args(["-c", "sleep 30"]) |
| 1392 | .stdin(Stdio::piped()) |
| 1393 | .stdout(Stdio::piped()) |
| 1394 | .spawn() |
| 1395 | .unwrap(); |
| 1396 | // Close stdin like `Drop` does, then reap: the sleeper ignores it. |
| 1397 | drop(child.stdin.take()); |
| 1398 | reap_child(&mut child); |
| 1399 | assert!( |
| 1400 | child.try_wait().unwrap().is_some(), |
| 1401 | "reaper must have killed the child" |
| 1402 | ); |
| 1403 | } |
| 1404 | |
| 1405 | #[cfg(unix)] |
| 1406 | #[test] |
| 1407 | fn dropping_a_connection_with_a_live_child_returns_before_the_grace() { |
| 1408 | let child = Command::new("/bin/sh") |
| 1409 | .args(["-c", "sleep 30"]) |
| 1410 | .stdin(Stdio::piped()) |
| 1411 | .stdout(Stdio::piped()) |
| 1412 | .spawn() |
| 1413 | .unwrap(); |
| 1414 | let pid = child.id(); |
| 1415 | let (_tx, responses) = sync_channel(1); |
| 1416 | let connection = Connection { |
| 1417 | child: Some(child), |
| 1418 | stdin: None, |
| 1419 | responses, |
| 1420 | next_id: 1, |
| 1421 | }; |
| 1422 | let started = Instant::now(); |
| 1423 | drop(connection); |
| 1424 | assert!( |
| 1425 | started.elapsed() < SHUTDOWN_GRACE, |
| 1426 | "drop must hand off instead of grace-waiting" |
| 1427 | ); |
| 1428 | // The reaper owns the child now: it must die without anyone waiting |
| 1429 | // inline, so the handoff leaks nothing. |
| 1430 | let deadline = Instant::now() + Duration::from_secs(10); |
| 1431 | loop { |
| 1432 | let alive = unsafe { libc::kill(pid as libc::pid_t, 0) } == 0; |
| 1433 | if !alive { |
| 1434 | break; |
| 1435 | } |
| 1436 | assert!(Instant::now() < deadline, "reaper never reaped pid {pid}"); |
| 1437 | thread::sleep(Duration::from_millis(50)); |
| 1438 | } |
| 1439 | } |
| 1440 | |
| 1441 | #[cfg(unix)] |
| 1442 | fn assert_descendant_cleanup(handshake: bool, launcher_exits: bool) { |
| 1443 | use std::io::Read; |
| 1444 | |
| 1445 | let unique = std::time::SystemTime::now() |
| 1446 | .duration_since(std::time::UNIX_EPOCH) |
| 1447 | .unwrap() |
| 1448 | .as_nanos(); |
| 1449 | let root = std::env::temp_dir().join(format!( |
| 1450 | "codewhale-mcp-tree-{}-{unique}", |
| 1451 | std::process::id() |
| 1452 | )); |
| 1453 | std::fs::create_dir(&root).unwrap(); |
| 1454 | let pipe = root.join("pipe"); |
| 1455 | let pids = root.join("pids"); |
| 1456 | let path = std::ffi::CString::new(pipe.as_os_str().as_encoded_bytes()).unwrap(); |
| 1457 | assert_eq!(unsafe { libc::mkfifo(path.as_ptr(), 0o600) }, 0); |
| 1458 | let (eof_tx, eof_rx) = std::sync::mpsc::channel(); |
| 1459 | let reader = thread::spawn(move || { |
| 1460 | let result = |
| 1461 | std::fs::File::open(pipe).and_then(|mut pipe| pipe.read_to_end(&mut Vec::new())); |
| 1462 | let _ = eof_tx.send(result); |
| 1463 | }); |
| 1464 | let mut script = String::from( |
| 1465 | "exec 3>\"$MCP_TEST_PIPE\"\nsleep 30 &\nprintf '%s %s\\n' \"$$\" \"$!\" >\"$MCP_TEST_PIDS\"\n", |
| 1466 | ); |
| 1467 | if handshake { |
| 1468 | script.push_str( |
| 1469 | r#"IFS= read -r line |
| 1470 | id=$(printf '%s' "$line" | sed -n 's/.*"id":\([0-9][0-9]*\).*/\1/p') |
| 1471 | printf '{"jsonrpc":"2.0","id":%s,"result":{"protocolVersion":"2024-11-05","capabilities":{},"serverInfo":{"name":"tree","version":"1"}}}\n' "$id" |
| 1472 | IFS= read -r initialized |
| 1473 | "#, |
| 1474 | ); |
| 1475 | } |
| 1476 | script.push_str(if launcher_exits { "exit 0\n" } else { "wait\n" }); |
| 1477 | let mut cfg = config("/bin/sh", &["-c", &script]); |
| 1478 | cfg.env.insert( |
| 1479 | "MCP_TEST_PIPE".into(), |
| 1480 | root.join("pipe").display().to_string(), |
| 1481 | ); |
| 1482 | cfg.env |
| 1483 | .insert("MCP_TEST_PIDS".into(), pids.display().to_string()); |
| 1484 | |
| 1485 | let result = std::panic::catch_unwind(|| { |
| 1486 | let spawned = ChildProcessMcpClient::spawn_with_timeouts( |
| 1487 | &cfg, |
| 1488 | Duration::from_secs(2), |
| 1489 | Duration::from_secs(2), |
| 1490 | ); |
| 1491 | if handshake { |
| 1492 | let client = spawned.expect("normal handshake must succeed"); |
| 1493 | if launcher_exits { |
| 1494 | let deadline = Instant::now() + Duration::from_secs(2); |
| 1495 | loop { |
| 1496 | if client |
| 1497 | .connection |
| 1498 | .lock() |
| 1499 | .unwrap() |
| 1500 | .child |
| 1501 | .as_mut() |
| 1502 | .expect("child present before drop") |
| 1503 | .try_wait() |
| 1504 | .unwrap() |
| 1505 | .is_some() |
| 1506 | { |
| 1507 | break; |
| 1508 | } |
| 1509 | assert!(Instant::now() < deadline, "launcher did not exit"); |
| 1510 | thread::sleep(Duration::from_millis(10)); |
| 1511 | } |
| 1512 | } |
| 1513 | drop(client); |
| 1514 | } else { |
| 1515 | assert!(format!("{:#}", spawned.unwrap_err()).contains("initialize timed out")); |
| 1516 | } |
| 1517 | let ids: Vec<libc::pid_t> = std::fs::read_to_string(&pids) |
| 1518 | .unwrap() |
| 1519 | .split_whitespace() |
| 1520 | .map(|pid| pid.parse().unwrap()) |
| 1521 | .collect(); |
| 1522 | assert_eq!(ids.len(), 2, "must observe both launcher and descendant"); |
| 1523 | let deadline = Instant::now() + Duration::from_secs(2); |
| 1524 | for pid in ids { |
| 1525 | loop { |
| 1526 | let gone = unsafe { libc::kill(pid, 0) } != 0 |
| 1527 | && io::Error::last_os_error().raw_os_error() == Some(libc::ESRCH); |
| 1528 | // An adopted Linux zombie has exited; its reaping belongs |
| 1529 | // to init, not this client. It cannot retain a pipe or run. |
| 1530 | #[cfg(target_os = "linux")] |
| 1531 | let gone = gone |
| 1532 | || std::fs::read_to_string(format!("/proc/{pid}/stat")) |
| 1533 | .ok() |
| 1534 | .and_then(|stat| { |
| 1535 | stat.rsplit_once(") ") |
| 1536 | .map(|(_, tail)| tail.starts_with('Z')) |
| 1537 | }) |
| 1538 | .unwrap_or(false); |
| 1539 | if gone { |
| 1540 | break; |
| 1541 | } |
| 1542 | assert!( |
| 1543 | Instant::now() < deadline, |
| 1544 | "MCP process {pid} survived cleanup" |
| 1545 | ); |
| 1546 | thread::sleep(Duration::from_millis(10)); |
| 1547 | } |
| 1548 | } |
| 1549 | eof_rx |
| 1550 | .recv_timeout(Duration::from_secs(2)) |
| 1551 | .expect("descendant still holds inherited pipe") |
| 1552 | .expect("pipe read failed"); |
| 1553 | }); |
| 1554 | if result.is_err() { |
| 1555 | // Test failures must not leave their controlled sleepers running. |
| 1556 | if let Ok(ids) = std::fs::read_to_string(&pids) { |
| 1557 | for pid in ids |
| 1558 | .split_whitespace() |
| 1559 | .filter_map(|pid| pid.parse::<libc::pid_t>().ok()) |
| 1560 | { |
| 1561 | unsafe { |
| 1562 | libc::kill(pid, libc::SIGKILL); |
| 1563 | } |
| 1564 | } |
| 1565 | } |
| 1566 | } |
| 1567 | if result.is_ok() { |
| 1568 | reader.join().unwrap(); |
| 1569 | } |
| 1570 | let _ = std::fs::remove_dir_all(root); |
| 1571 | if let Err(error) = result { |
| 1572 | std::panic::resume_unwind(error); |
| 1573 | } |
| 1574 | } |
| 1575 | |
| 1576 | #[cfg(unix)] |
| 1577 | #[test] |
| 1578 | fn handshake_timeout_terminates_launcher_descendant_and_pipes() { |
| 1579 | assert_descendant_cleanup(false, false); |
| 1580 | } |
| 1581 | |
| 1582 | #[cfg(unix)] |
| 1583 | #[test] |
| 1584 | fn normal_connection_drop_terminates_descendant_and_pipes() { |
| 1585 | assert_descendant_cleanup(true, false); |
| 1586 | } |
| 1587 | |
| 1588 | #[cfg(unix)] |
| 1589 | #[test] |
| 1590 | fn connection_drop_after_launcher_exit_terminates_descendant_and_pipes() { |
| 1591 | assert_descendant_cleanup(true, true); |
| 1592 | } |
| 1593 | |
| 1594 | #[cfg(unix)] |
| 1595 | #[test] |
| 1596 | fn handshake_answers_a_same_id_child_ping_before_accepting_the_response() { |
| 1597 | let script = r#" |
| 1598 | IFS= read -r line |
| 1599 | id=$(printf '%s' "$line" | sed -n 's/.*"id":\([0-9][0-9]*\).*/\1/p') |
| 1600 | printf '{"jsonrpc":"2.0","id":%s,"method":"ping"}\n' "$id" |
| 1601 | IFS= read -r pong |
| 1602 | case "$pong" in |
| 1603 | *'"result":{}'*) ;; |
| 1604 | *) exit 9 ;; |
| 1605 | esac |
| 1606 | printf '{"jsonrpc":"2.0","id":%s,"result":{"protocolVersion":"2024-11-05","capabilities":{"tools":{}},"serverInfo":{"name":"pinging-child","version":"1"}}}\n' "$id" |
| 1607 | while IFS= read -r _line; do :; done |
| 1608 | "#; |
| 1609 | let client = ChildProcessMcpClient::spawn_with_timeouts( |
| 1610 | &config("/bin/sh", &["-c", script]), |
| 1611 | Duration::from_secs(2), |
| 1612 | Duration::from_secs(2), |
| 1613 | ) |
| 1614 | .expect("the child ping must be answered during initialize"); |
| 1615 | assert!(client.supports_tools()); |
| 1616 | } |
| 1617 | |
| 1618 | #[cfg(unix)] |
| 1619 | #[test] |
| 1620 | fn idle_child_requests_are_answered_before_the_next_client_request() { |
| 1621 | let script = r#" |
| 1622 | ready=0 |
| 1623 | while IFS= read -r line; do |
| 1624 | id=$(printf '%s' "$line" | sed -n 's/.*"id":\([0-9][0-9]*\).*/\1/p') |
| 1625 | method=$(printf '%s' "$line" | sed -n 's/.*"method":"\([^"]*\)".*/\1/p') |
| 1626 | case "$method" in |
| 1627 | initialize) |
| 1628 | printf '{"jsonrpc":"2.0","id":%s,"result":{"protocolVersion":"2024-11-05","capabilities":{"tools":{}},"serverInfo":{"name":"idle-request-child","version":"1"}}}\n' "$id" |
| 1629 | ;; |
| 1630 | notifications/initialized) |
| 1631 | printf '{"jsonrpc":"2.0","id":"idle-ping","method":"ping"}\n' |
| 1632 | (sleep 1; kill -KILL "$$") & |
| 1633 | watchdog=$! |
| 1634 | IFS= read -r pong || exit 9 |
| 1635 | kill "$watchdog" 2>/dev/null || : |
| 1636 | wait "$watchdog" 2>/dev/null || : |
| 1637 | case "$pong" in *'"id":"idle-ping"'*) ;; *) exit 10 ;; esac |
| 1638 | case "$pong" in *'"result":{}'*) ;; *) exit 11 ;; esac |
| 1639 | |
| 1640 | printf '{"jsonrpc":"2.0","id":"idle-unsupported","method":"client/unsupported"}\n' |
| 1641 | (sleep 1; kill -KILL "$$") & |
| 1642 | watchdog=$! |
| 1643 | IFS= read -r unsupported || exit 12 |
| 1644 | kill "$watchdog" 2>/dev/null || : |
| 1645 | wait "$watchdog" 2>/dev/null || : |
| 1646 | case "$unsupported" in *'"id":"idle-unsupported"'*) ;; *) exit 13 ;; esac |
| 1647 | case "$unsupported" in *'"code":-32601'*) ;; *) exit 14 ;; esac |
| 1648 | ready=1 |
| 1649 | ;; |
| 1650 | tools/list) |
| 1651 | [ "$ready" = 1 ] || exit 15 |
| 1652 | printf '{"jsonrpc":"2.0","id":%s,"result":{"tools":[]}}\n' "$id" |
| 1653 | ;; |
| 1654 | esac |
| 1655 | done |
| 1656 | printf 'stdin closed\n' > "$CODEWHALE_MCP_TEST_MARKER" |
| 1657 | "#; |
| 1658 | let marker = std::env::temp_dir().join(format!( |
| 1659 | "codewhale-mcp-idle-drop-{}-{}.marker", |
| 1660 | std::process::id(), |
| 1661 | std::time::SystemTime::now() |
| 1662 | .duration_since(std::time::UNIX_EPOCH) |
| 1663 | .expect("system clock after epoch") |
| 1664 | .as_nanos() |
| 1665 | )); |
| 1666 | let mut child_config = config("/bin/sh", &["-c", script]); |
| 1667 | child_config.env.insert( |
| 1668 | "CODEWHALE_MCP_TEST_MARKER".to_string(), |
| 1669 | marker.to_string_lossy().into_owned(), |
| 1670 | ); |
| 1671 | let client = ChildProcessMcpClient::spawn_with_timeouts( |
| 1672 | &child_config, |
| 1673 | Duration::from_secs(2), |
| 1674 | Duration::from_secs(2), |
| 1675 | ) |
| 1676 | .expect("handshake"); |
| 1677 | |
| 1678 | // The child terminates if either server-initiated request remains |
| 1679 | // unanswered for one second. Keep the client genuinely idle beyond |
| 1680 | // that deadline before making the next ordinary client request. |
| 1681 | thread::sleep(Duration::from_millis(1_500)); |
| 1682 | assert!(client.list_tools().unwrap().is_empty()); |
| 1683 | |
| 1684 | // The reader owns only a weak stdin handle. Dropping the client must |
| 1685 | // therefore still deliver EOF to the child and let it exit cleanly; |
| 1686 | // a strong reader-thread handle would force the kill fallback. Drop |
| 1687 | // hands the child to the reaper thread instead of waiting, so poll |
| 1688 | // for the marker the clean exit writes. |
| 1689 | drop(client); |
| 1690 | let deadline = Instant::now() + Duration::from_secs(5); |
| 1691 | loop { |
| 1692 | // The child creates and truncates the marker before writing it, |
| 1693 | // so a successful read can still land on an empty or partial |
| 1694 | // file. Poll until it carries the whole line a clean exit writes. |
| 1695 | if std::fs::read_to_string(&marker).is_ok_and(|body| body == "stdin closed\n") { |
| 1696 | break; |
| 1697 | } |
| 1698 | assert!( |
| 1699 | Instant::now() < deadline, |
| 1700 | "dropped client never delivered stdin EOF" |
| 1701 | ); |
| 1702 | thread::sleep(Duration::from_millis(20)); |
| 1703 | } |
| 1704 | std::fs::remove_file(marker).unwrap(); |
| 1705 | } |
| 1706 | |
| 1707 | #[cfg(unix)] |
| 1708 | #[test] |
| 1709 | fn handshake_rejects_a_child_response_without_jsonrpc_2_0() { |
| 1710 | let script = r#" |
| 1711 | IFS= read -r line |
| 1712 | id=$(printf '%s' "$line" | sed -n 's/.*"id":\([0-9][0-9]*\).*/\1/p') |
| 1713 | printf '{"id":%s,"result":{"protocolVersion":"2024-11-05","capabilities":{},"serverInfo":{"name":"invalid-child","version":"1"}}}\n' "$id" |
| 1714 | while IFS= read -r _line; do :; done |
| 1715 | "#; |
| 1716 | let error = ChildProcessMcpClient::spawn_with_timeouts( |
| 1717 | &config("/bin/sh", &["-c", script]), |
| 1718 | Duration::from_secs(2), |
| 1719 | Duration::from_secs(2), |
| 1720 | ) |
| 1721 | .unwrap_err(); |
| 1722 | assert!( |
| 1723 | format!("{error:#}").contains("without jsonrpc \"2.0\""), |
| 1724 | "unexpected error: {error:#}" |
| 1725 | ); |
| 1726 | } |
| 1727 | |
| 1728 | #[cfg(unix)] |
| 1729 | #[test] |
| 1730 | fn malformed_child_catalog_entries_fail_their_page_immediately() { |
| 1731 | let script = r#" |
| 1732 | while IFS= read -r line; do |
| 1733 | id=$(printf '%s' "$line" | sed -n 's/.*"id":\([0-9][0-9]*\).*/\1/p') |
| 1734 | method=$(printf '%s' "$line" | sed -n 's/.*"method":"\([^"]*\)".*/\1/p') |
| 1735 | [ -n "$id" ] || continue |
| 1736 | case "$method" in |
| 1737 | initialize) |
| 1738 | printf '{"jsonrpc":"2.0","id":%s,"result":{"protocolVersion":"2024-11-05","capabilities":{"tools":{},"resources":{}},"serverInfo":{"name":"malformed-catalog","version":"1"}}}\n' "$id" |
| 1739 | ;; |
| 1740 | tools/list) |
| 1741 | printf '{"jsonrpc":"2.0","id":%s,"result":{"tools":[{"name":"bad-schema","inputSchema":null}],"nextCursor":"more"}}\n' "$id" |
| 1742 | ;; |
| 1743 | resources/list) |
| 1744 | printf '{"jsonrpc":"2.0","id":%s,"result":{"resources":[{"uri":"file:///missing-name"}],"nextCursor":"more"}}\n' "$id" |
| 1745 | ;; |
| 1746 | esac |
| 1747 | done |
| 1748 | "#; |
| 1749 | let client = |
| 1750 | ChildProcessMcpClient::spawn(&config("/bin/sh", &["-c", script])).expect("handshake"); |
| 1751 | |
| 1752 | let tools_error = client.list_tools().unwrap_err(); |
| 1753 | assert!( |
| 1754 | tools_error.to_string().contains("tools[0].inputSchema"), |
| 1755 | "the malformed first page must fail before following its cursor: {tools_error:#}" |
| 1756 | ); |
| 1757 | let resources_error = client.list_resources().unwrap_err(); |
| 1758 | assert!( |
| 1759 | resources_error.to_string().contains("resources[0].name"), |
| 1760 | "the malformed first page must fail before following its cursor: {resources_error:#}" |
| 1761 | ); |
| 1762 | } |
| 1763 | |
| 1764 | #[cfg(unix)] |
| 1765 | #[test] |
| 1766 | fn a_real_child_answers_tools_and_resources() { |
| 1767 | let script = crate::test_support::write_fake_mcp_server("stdio_client_roundtrip"); |
| 1768 | let client = ChildProcessMcpClient::spawn(&config( |
| 1769 | "/bin/sh", |
| 1770 | &[script.path().to_str().expect("utf-8 script path")], |
| 1771 | )) |
| 1772 | .expect("fake MCP server should complete the handshake"); |
| 1773 | |
| 1774 | let tools = client.list_tools_with_input_schemas().unwrap(); |
| 1775 | assert_eq!(tools.len(), 1); |
| 1776 | assert_eq!(tools[0].0.tool_name, "add"); |
| 1777 | assert_eq!(tools[0].1["required"], json!(["a", "b"])); |
| 1778 | |
| 1779 | let result = client.call_tool("add", json!({"a": 2, "b": 3})).unwrap(); |
| 1780 | assert_eq!( |
| 1781 | result["content"][0]["text"], "5", |
| 1782 | "the answer must come from the child process: {result}" |
| 1783 | ); |
| 1784 | |
| 1785 | let resources = client.list_resources_with_metadata().unwrap(); |
| 1786 | assert_eq!(resources.len(), 1); |
| 1787 | assert_eq!(resources[0].0.uri, "file:///fake/readme.txt"); |
| 1788 | assert_eq!(resources[0].1["name"], "Fake readme"); |
| 1789 | assert_eq!(resources[0].1["mimeType"], "text/plain"); |
| 1790 | assert_eq!(resources[0].1["size"], 16); |
| 1791 | assert_eq!( |
| 1792 | resources[0].1["annotations"]["audience"], |
| 1793 | json!(["assistant"]) |
| 1794 | ); |
| 1795 | |
| 1796 | let resource = client.read_resource("file:///fake/readme.txt").unwrap(); |
| 1797 | assert_eq!(resource["contents"][0]["text"], "spawned-resource"); |
| 1798 | |
| 1799 | // The stub's fabricated tools must be gone. |
| 1800 | assert!(client.call_tool("health", json!({})).is_err()); |
| 1801 | } |
| 1802 | } |
| 1803 |