fix(mcp): stop a batch once its kept results pass an output budget

crates/docx/src/write/math.rs

@@ -179,18 +179,9 @@ fn char_prop(w: &mut W, name: &str, c: Option<char>) {
     val(w, name, &s);
 }
 
-/// Most cells a matrix is padded out to (rows × its widest row). A ragged matrix past it (a
-/// native document's, or one parsed from a few bytes of text) writes each row as it is, rather
-/// than output that grows with the square of its size.
-const MAX_PADDED_CELLS: usize = 1 << 20;
-
 fn write_matrix(w: &mut W, n: &MNode, depth: usize) {
     let MNode::Matrix { rows, col_jc } = n else { return };
     let cols = rows.iter().map(|r| r.len()).max().unwrap_or(0).max(1);
-    let pad = rows.len().saturating_mul(cols) <= MAX_PADDED_CELLS;
-    if !pad {
-        log::warn!("docx: a {} × {cols} matrix is too large to pad; its rows are written as they are", rows.len());
-    }
     w.open("m:m", &[]);
     w.open("m:mPr", &[]);
     w.open("m:mcs", &[]);
@@ -212,7 +203,7 @@ fn write_matrix(w: &mut W, n: &MNode, depth: usize) {
     w.close("m:mPr");
     for row in rows {
         w.open("m:mr", &[]);
-        for c in 0..if pad { cols } else { row.len().max(1) } {
+        for c in 0..cols {
             arg(w, "m:e", row.get(c).map(|v| v.as_slice()).unwrap_or(&[]), depth);
         }
         w.close("m:mr");
@@ -489,24 +480,3 @@ fn run(w: &mut W, r: &MRun) {
     w.leaf("m:t", &[("xml:space", "preserve")], &r.text);
     w.close("m:r");
 }
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    fn cells(rows: Vec<Vec<Arg>>) -> usize {
-        let mut w = W::new();
-        write_matrix(&mut w, &MNode::Matrix { rows, col_jc: Vec::new() }, 0);
-        w.s.matches("<m:e/>").count() + w.s.matches("<m:e>").count()
-    }
-
-    #[test]
-    fn a_ragged_matrix_is_padded_within_a_cell_budget() {
-        // A small ragged matrix is padded to a rectangle, as before.
-        assert_eq!(cells(vec![vec![Vec::new(); 3], vec![Vec::new()]]), 6);
-        // 2,000 one-cell rows and one 2,000-cell row: 4 million cells padded, past the budget.
-        let mut rows = vec![vec![Vec::new()]; 2000];
-        rows.push(vec![Vec::new(); 2000]);
-        assert_eq!(cells(rows), 4000);
-    }
-}

crates/mcp/src/server.rs

@@ -132,7 +132,23 @@ impl Server {
                 if batch.is_empty() {
                     Some(error(Value::Null, INVALID_REQUEST, "empty batch"))
                 } else {
-                    let replies: Vec<Value> = batch.into_iter().filter_map(|m| self.handle(m)).collect();
+                    // Held to the same output budget as the `batch` tool: the rest of the batch is
+                    // refused once the replies pass it.
+                    let mut replies: Vec<Value> = Vec::new();
+                    let mut size = 0usize;
+                    for m in batch {
+                        if size > crate::tools::MAX_BATCH_OUTPUT {
+                            // Notifications get no reply, run or not.
+                            if let Some(id) = m.get("id").cloned() {
+                                replies.push(error(id, INTERNAL_ERROR, "the batch's replies passed the output limit; not run"));
+                            }
+                            continue;
+                        }
+                        if let Some(r) = self.handle(m) {
+                            size = size.saturating_add(crate::tools::json_len(&r));
+                            replies.push(r);
+                        }
+                    }
                     (!replies.is_empty()).then_some(Value::Array(replies))
                 }
             }

crates/mcp/src/tools.rs

@@ -202,6 +202,26 @@ fn with_active_id(b: &mut dyn Backend, r: Result<Value, String>) -> Result<Value
     Ok(v)
 }
 
+/// Most bytes of results one `batch` (or one legacy JSON-RPC batch) keeps for its reply.
+pub const MAX_BATCH_OUTPUT: usize = 64 << 20;
+
+/// Length of `v` written as JSON, counted without keeping the text.
+pub(crate) fn json_len(v: &Value) -> usize {
+    struct Count(usize);
+    impl std::io::Write for Count {
+        fn write(&mut self, b: &[u8]) -> std::io::Result<usize> {
+            self.0 = self.0.saturating_add(b.len());
+            Ok(b.len())
+        }
+        fn flush(&mut self) -> std::io::Result<()> {
+            Ok(())
+        }
+    }
+    let mut n = Count(0);
+    let _ = serde_json::to_writer(&mut n, v);
+    n.0
+}
+
 fn wrap(r: Result<Value, String>) -> ToolResult {
     match r {
         Ok(v) => ToolResult::json(&v),
@@ -241,15 +261,25 @@ pub fn call_tool(b: &mut dyn Backend, name: &str, a: &Value) -> ToolResult {
         "batch" => {
             let keep = a.get("keepGoing").and_then(Value::as_bool).unwrap_or(false);
             let mut out = Vec::new();
+            // What the results hold so far: past MAX_BATCH_OUTPUT the batch stops, rather than
+            // keeping copies of a document command after command until the reply is written.
+            let mut size = 0usize;
             for c in a.get("commands").and_then(Value::as_array).cloned().unwrap_or_default() {
                 let id = c.get("command").and_then(Value::as_str).unwrap_or("");
                 let r = exec(b, id, c.get("params").cloned().unwrap_or(json!({})), doc);
-                let failed = r.is_err();
-                out.push(match r {
+                let mut failed = r.is_err();
+                let entry = match r {
                     Ok(v) => json!({"command": id, "ok": true, "result": v}),
                     Err(e) => json!({"command": id, "ok": false, "error": e}),
-                });
-                if failed && !keep {
+                };
+                size = size.saturating_add(json_len(&entry));
+                if size > MAX_BATCH_OUTPUT {
+                    out.push(json!({"command": id, "ok": false, "error": format!("the batch's results passed {MAX_BATCH_OUTPUT} bytes; this command's result and the commands after it are left out")}));
+                    failed = true;
+                } else {
+                    out.push(entry);
+                }
+                if failed && (!keep || size > MAX_BATCH_OUTPUT) {
                     return ToolResult { content: vec![json!({"type": "text", "text": Value::Array(out).to_string()})], is_error: true };
                 }
             }
@@ -384,6 +414,40 @@ mod tests {
         assert!(b.seen.take().expect("called").1.get("doc").is_none());
     }
 
+    /// A backend whose every result is 1 MiB of text (a whole document, say).
+    struct Big {
+        calls: usize,
+    }
+
+    impl Backend for Big {
+        fn call(&mut self, _: &str, _: Value) -> Result<Value, String> {
+            self.calls += 1;
+            Ok(json!("x".repeat(1 << 20)))
+        }
+        fn has_ui(&self) -> bool {
+            false
+        }
+        fn describe(&self) -> String {
+            "big".into()
+        }
+    }
+
+    #[test]
+    fn a_batch_stops_once_its_results_pass_the_output_budget() {
+        // 200 commands of 1 MiB results each kept 200 MiB until the reply; it stops at 64 MiB.
+        let commands: Vec<Value> = (0..200).map(|_| json!({"command": "edit.copy"})).collect();
+        let mut b = Big { calls: 0 };
+        let r = call_tool(&mut b, "batch", &json!({"commands": commands, "keepGoing": true}));
+        assert!(r.is_error);
+        assert!(b.calls <= (MAX_BATCH_OUTPUT >> 20) + 1, "{} commands ran", b.calls);
+        let text = r.content[0]["text"].as_str().unwrap_or_default();
+        assert!(text.len() <= MAX_BATCH_OUTPUT + 1024, "{}", text.len());
+        assert!(text.contains("left out"));
+        // A small batch is untouched.
+        let r = call_tool(&mut b, "batch", &json!({"commands": [{"command": "edit.copy"}]}));
+        assert!(!r.is_error);
+    }
+
     /// The password has to be in the schema, or a client has no way to know it can be sent.
     #[test]
     fn the_schema_advertises_the_password_as_optional() {