address review: charge batch error entries against the byte budget

This commit is contained in:
Dylan M. Taylor
2026-11-04 04:33:40 +00:00
parent 867a89e112
commit b5801bb93c
2 changed files with 26 additions and 17 deletions
+14 -17
View File
@@ -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))
+12
View File
@@ -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"));
}