From bf990f8b6a910e85a5fd41ee69d361e68828ea94 Mon Sep 17 00:00:00 2001 From: "Dylan M. Taylor" Date: Tue, 3 Nov 2026 11:30:53 +0000 Subject: [PATCH] mcp: bound a batch's entries and the bytes of results it gathers --- crates/mcp/src/server.rs | 28 +++++++++++++++++- crates/mcp/src/tools.rs | 17 +++++++++-- crates/mcp/tests/it/batch_budget.rs | 45 +++++++++++++++++++++++++++++ crates/mcp/tests/it/main.rs | 2 ++ 4 files changed, 89 insertions(+), 3 deletions(-) create mode 100644 crates/mcp/tests/it/batch_budget.rs diff --git a/crates/mcp/src/server.rs b/crates/mcp/src/server.rs index 747df0152..d28f184ff 100644 --- a/crates/mcp/src/server.rs +++ b/crates/mcp/src/server.rs @@ -16,6 +16,12 @@ const INVALID_REQUEST: i64 = -32600; const METHOD_NOT_FOUND: i64 = -32601; const INVALID_PARAMS: i64 = -32602; const INTERNAL_ERROR: i64 = -32603; + +/// Entries one batch (a JSON-RPC batch, or the `batch` tool's commands) may hold. +pub(crate) const MAX_BATCH_ITEMS: usize = 1000; +/// Bytes of results one batch may gather before the rest of it is refused: each result can copy +/// much of the document (an inspection holds every note), so a batch holds at most this much. +pub(crate) const MAX_BATCH_BYTES: usize = 64 << 20; const RESOURCE_NOT_FOUND: i64 = -32002; /// Keep request frames in step with the connected backend's reply limit. const MAX_FRAME_BYTES: usize = 16 * 1024 * 1024; @@ -93,8 +99,28 @@ impl Server { // Batches were removed in 2025-06-18; still answer older clients sensibly. if batch.is_empty() { Some(error(Value::Null, INVALID_REQUEST, "empty batch")) + } else if batch.len() > MAX_BATCH_ITEMS { + Some(error(Value::Null, INVALID_REQUEST, format!("a batch holds at most {MAX_BATCH_ITEMS} messages"))) } else { - let replies: Vec = batch.into_iter().filter_map(|m| self.handle(m)).collect(); + let mut used = 0usize; + let mut replies: Vec = Vec::new(); + for m in batch { + if used > MAX_BATCH_BYTES { + // Not run: answer each remaining request with a short error. + if let Some(id) = m.get("id").filter(|id| id.is_string() || id.is_number()) { + replies.push(error(id.clone(), INTERNAL_ERROR, "the batch's replies passed 64 MiB; not run")); + } + continue; + } + if let Some(r) = self.handle(m) { + used = used.saturating_add(r.to_string().len()); + replies.push(if used > MAX_BATCH_BYTES { + error(r.get("id").cloned().unwrap_or(Value::Null), INTERNAL_ERROR, "the batch's replies passed 64 MiB; reply dropped") + } else { + r + }); + } + } (!replies.is_empty()).then_some(Value::Array(replies)) } } diff --git a/crates/mcp/src/tools.rs b/crates/mcp/src/tools.rs index 2bf09b8a9..3bae69d25 100644 --- a/crates/mcp/src/tools.rs +++ b/crates/mcp/src/tools.rs @@ -357,12 +357,25 @@ pub fn call_tool(b: &mut dyn Backend, name: &str, a: &Value) -> ToolResult { wrap(exec(b, c, a.get("params").cloned().unwrap_or(json!({})))) } "batch" => { + use crate::server::{MAX_BATCH_BYTES, MAX_BATCH_ITEMS}; let keep = a.get("keepGoing").and_then(Value::as_bool).unwrap_or(false); + let commands = a.get("commands").and_then(Value::as_array).map(Vec::as_slice).unwrap_or_default(); + if commands.len() > MAX_BATCH_ITEMS { + return ToolResult::error(format!("a batch holds at most {MAX_BATCH_ITEMS} commands")); + } let mut out = vec![]; - for c in a.get("commands").and_then(Value::as_array).cloned().unwrap_or_default() { + let mut used = 0usize; + for c in commands { let id = c.get("command").and_then(Value::as_str).unwrap_or(""); match exec(b, id, c.get("params").cloned().unwrap_or(json!({}))) { - Ok(v) => out.push(json!({"command": id, "ok": true, "result": v})), + Ok(v) => { + used = used.saturating_add(v.to_string().len()); + if used > MAX_BATCH_BYTES { + out.push(json!({"command": id, "ok": false, "error": "the batch's results passed 64 MiB; stopped here"})); + return ToolResult { content: vec![json!({"type": "text", "text": Value::Array(out).to_string()})], is_error: true }; + } + out.push(json!({"command": id, "ok": true, "result": v})); + } Err(e) => { out.push(json!({"command": id, "ok": false, "error": e})); if !keep { diff --git a/crates/mcp/tests/it/batch_budget.rs b/crates/mcp/tests/it/batch_budget.rs new file mode 100644 index 000000000..3794c1551 --- /dev/null +++ b/crates/mcp/tests/it/batch_budget.rs @@ -0,0 +1,45 @@ +//! A batch (JSON-RPC or the `batch` tool) is bounded in entries and in the bytes of results it +//! gathers, so a short request cannot copy a large document into memory hundreds of times. + +use deckcraft_mcp::{Headless, Server}; +use serde_json::{Value, json}; + +fn server_with_notes(bytes: usize) -> Server { + let mut server = Server::new(Box::new(Headless::with_document())); + let line = json!({"jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": {"name": "run_command", "arguments": {"command": "slide.notes", "params": {"text": "n".repeat(bytes)}}}}); + let reply: Value = serde_json::from_str(&server.handle_line(&line.to_string()).unwrap()).unwrap(); + assert_ne!(reply["result"]["isError"], true, "{reply}"); + server +} + +#[test] +fn json_rpc_batches_are_bounded_in_entries_and_reply_bytes() { + let mut server = server_with_notes(1 << 20); + let ping = |i: usize| json!({"jsonrpc": "2.0", "id": i, "method": "ping"}); + let too_many: Vec = (0..1001).map(ping).collect(); + let reply: Value = serde_json::from_str(&server.handle_line(&Value::Array(too_many).to_string()).unwrap()).unwrap(); + assert_eq!(reply["error"]["code"], -32600, "{reply}"); + // 100 inspections of a slide with 1 MB of notes: about 64 are answered, the rest refused. + let inspect = |i: usize| json!({"jsonrpc": "2.0", "id": i, "method": "tools/call", "params": {"name": "inspect_slide", "arguments": {}}}); + let line = Value::Array((0..100).map(inspect).collect()).to_string(); + let reply = server.handle_line(&line).unwrap(); + assert!(reply.len() < (64 << 20) + (4 << 20), "{} bytes", reply.len()); + let replies: Vec = serde_json::from_str(&reply).unwrap(); + assert_eq!(replies.len(), 100, "every request is answered"); + assert!(replies.last().unwrap()["error"]["message"].as_str().unwrap().contains("64 MiB")); +} + +#[test] +fn the_batch_tool_is_bounded_in_commands_and_result_bytes() { + let mut server = server_with_notes(1 << 20); + let call = |server: &mut Server, commands: Vec| { + let line = json!({"jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": {"name": "batch", "arguments": {"commands": commands, "keepGoing": true}}}); + let reply: Value = serde_json::from_str(&server.handle_line(&line.to_string()).unwrap()).unwrap(); + (reply["result"]["isError"] == true, reply["result"]["content"][0]["text"].as_str().unwrap_or_default().to_string()) + }; + let (err, text) = call(&mut server, vec![json!({"command": "document.inspect"}); 1001]); + assert!(err && text.contains("at most 1000"), "{text}"); + let (err, text) = call(&mut server, vec![json!({"command": "slide.inspect"}); 100]); + assert!(err && text.contains("64 MiB"), "{}", &text[text.len().saturating_sub(200)..]); + assert!(text.len() < (64 << 20) + (4 << 20)); +} diff --git a/crates/mcp/tests/it/main.rs b/crates/mcp/tests/it/main.rs index 0314d68c8..1b9af7d2d 100644 --- a/crates/mcp/tests/it/main.rs +++ b/crates/mcp/tests/it/main.rs @@ -1,3 +1,5 @@ // Integration tests: unwrapping and panicking on failure is fine here, unlike in shipped code (AGENTS.md). #[allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)] mod automation_roots; +#[allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)] +mod batch_budget;