From b5801bb93c15acf0af14cb0d623d383735f079e1 Mon Sep 17 00:00:00 2001
From: "Dylan M. Taylor" <dylan@dylanmtaylor.com>
Date: Wed, 4 Nov 2026 04:33:40 +0000
Subject: [PATCH] address review: charge batch error entries against the byte
 budget

---
 crates/mcp/src/tools.rs             | 31 +++++++++++++----------------
 crates/mcp/tests/it/batch_budget.rs | 12 +++++++++++
 2 files changed, 26 insertions(+), 17 deletions(-)

diff --git a/crates/mcp/src/tools.rs b/crates/mcp/src/tools.rs
index b5a05493e..e6c4ae7d0 100644
--- a/crates/mcp/src/tools.rs
+++ b/crates/mcp/src/tools.rs
@@ -367,23 +367,20 @@ pub fn call_tool(b: &mut dyn Backend, name: &str, a: &Value) -> ToolResult {
             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) => {
-                        let entry = json!({"command": id, "ok": true, "result": v});
-                        // What the reply carries: the entry as text, escaped again inside the JSON-RPC reply.
-                        used = used.saturating_add(json!(entry.to_string()).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(entry);
-                    }
-                    Err(e) => {
-                        out.push(json!({"command": id, "ok": false, "error": e}));
-                        if !keep {
-                            return ToolResult { content: vec![json!({"type": "text", "text": Value::Array(out).to_string()})], is_error: true };
-                        }
-                    }
+                let (entry, failed) = match exec(b, id, c.get("params").cloned().unwrap_or(json!({}))) {
+                    Ok(v) => (json!({"command": id, "ok": true, "result": v}), false),
+                    Err(e) => (json!({"command": id, "ok": false, "error": e}), true),
+                };
+                // What the reply carries: the entry as text, escaped again inside the JSON-RPC reply.
+                // Failures count too: with keepGoing, a remote backend's large errors add up.
+                used = used.saturating_add(json!(entry.to_string()).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(entry);
+                if failed && !keep {
+                    return ToolResult { content: vec![json!({"type": "text", "text": Value::Array(out).to_string()})], is_error: true };
                 }
             }
             ToolResult::json(&Value::Array(out))
diff --git a/crates/mcp/tests/it/batch_budget.rs b/crates/mcp/tests/it/batch_budget.rs
index 7172ddca0..ff98188c0 100644
--- a/crates/mcp/tests/it/batch_budget.rs
+++ b/crates/mcp/tests/it/batch_budget.rs
@@ -56,3 +56,15 @@ fn the_batch_tool_charges_the_escaped_reply() {
     assert!(reply.len() < (64 << 20) + (4 << 20), "{} bytes", reply.len());
     assert!(reply.contains("64 MiB"));
 }
+
+#[test]
+fn the_batch_tool_charges_failures_too() {
+    // Each unknown command's error echoes its ~1 MiB id; with keepGoing none of them is a result.
+    let mut server = Server::new(Box::new(Headless::with_document()));
+    let id = "x".repeat(1 << 20);
+    let commands = vec![json!({"command": id}); 70];
+    let line = json!({"jsonrpc": "2.0", "id": 2, "method": "tools/call", "params": {"name": "batch", "arguments": {"commands": commands, "keepGoing": true}}});
+    let reply = server.handle_line(&line.to_string()).unwrap();
+    assert!(reply.len() < (64 << 20) + (4 << 20), "{} bytes", reply.len());
+    assert!(reply.contains("64 MiB"));
+}