address review: drop the legacy-batch reply that crosses the output budget; budget batch results as escaped text

crates/layout/src/boxes.rs

@@ -519,9 +519,7 @@ pub(crate) const BOX_INSET_Y: f32 = 3.6;
 /// The first line always stays, so even a tiny box shows the start of its text.
 pub(crate) fn fit_box(items: Vec<Placed>, height: f32) -> Vec<Placed> {
     let limit = height + 0.5;
-    // Sized to what is kept, not to `items`: a box showing a few lines of a long story keeps
-    // no room for the rest of it.
-    let mut out = Vec::new();
+    let mut out = Vec::with_capacity(items.len());
     let mut kept_line = false;
     for it in items {
         if let Some(r) = it.turned_bounds() {

crates/layout/src/tests/budget.rs

@@ -294,17 +294,3 @@ fn a_kept_outline_holds_on_to_its_pictures_allocation() {
     drop(disc);
     assert!(weak.upgrade().is_none(), "the bytes are freed with the document; only their allocation is kept");
 }
-
-#[test]
-fn a_box_keeps_room_only_for_what_it_shows() {
-    // A linked box showing a few lines of a long story kept room for all of it (each of up to
-    // 256 boxes), and copied the whole shared layout before cutting it.
-    let color = wordcraft_doc::Rgb(0, 0, 0);
-    let items: Vec<Placed> = (0..10_000).map(|i| Placed::Fill { rect: Rect::new(0.0, i as f32 * 10.0, 5.0, 10.0), color }).collect();
-    let kept = boxes::fit_box(items.clone(), 25.0);
-    assert_eq!(kept.len(), 3);
-    assert!(kept.capacity() < 100, "{}", kept.capacity());
-    let below = textbox::from_y(&items, 99_950.0);
-    assert_eq!(below.len(), 6);
-    assert!(below.capacity() < 100, "{}", below.capacity());
-}

crates/layout/src/textbox.rs

@@ -252,35 +252,22 @@ fn flow(ctx: &mut Ctx, chain: &[u32], depth: usize) -> Vec<Vec<Placed>> {
         // Past the cap the box takes the nearest layout in its direction: the widest line length
         // that fits it, else the narrowest. The first box in a direction is laid out (one more
         // layout for each of the few directions at most).
-        let near = || {
-            layouts
-                .iter()
-                .enumerate()
-                .filter(|(_, ((_, v), _))| *v == key.1)
-                .min_by_key(|(_, ((k, _), _))| (*k > key.0, k.abs_diff(key.0)))
-                .map(|(i, _)| i)
-        };
-        let kept = layouts.iter().position(|(k, _)| *k == key).or_else(|| if layouts.len() >= MAX_FLOW_WIDTHS { near() } else { None });
-        let at = match kept {
-            Some(at) => at,
+        let near = || layouts.iter().filter(|((_, v), _)| *v == key.1).min_by_key(|((k, _), _)| (*k > key.0, k.abs_diff(key.0)));
+        let kept = layouts.iter().find(|(k, _)| *k == key).or_else(|| if layouts.len() >= MAX_FLOW_WIDTHS { near() } else { None });
+        let items = match kept {
+            Some((_, items)) => items.clone(),
             None => {
                 (ctx.counters, ctx.eq_count) = (snap.0.clone(), snap.1);
                 let items = lay(ctx, head, &blocks, len, body.vert, depth);
                 after.get_or_insert_with(|| (ctx.counters.clone(), ctx.eq_count));
-                layouts.push((key, items));
-                layouts.len() - 1
+                layouts.push((key, items.clone()));
+                items
             }
         };
-        // The shared layout is read in place: a box copies only the items from its own top
-        // down, not the whole story first.
-        let Some((_, shared)) = layouts.get(at) else {
-            out.push(Vec::new());
-            continue;
-        };
         let items = match &cut {
-            None => shared.clone(),
-            Some(cut) => match cut_y(shared, cut) {
-                Some(top) => from_y(shared, top),
+            None => items,
+            Some(cut) => match cut_y(&items, cut) {
+                Some(top) => from_y(items, top),
                 None => {
                     out.push(Vec::new());
                     continue;
@@ -311,28 +298,28 @@ fn cut_y(items: &[Placed], cut: &Cut) -> Option<f32> {
 }
 
 /// The items from `top` down, moved up to start at 0.
-pub(crate) fn from_y(items: &[Placed], top: f32) -> Vec<Placed> {
+fn from_y(items: Vec<Placed>, top: f32) -> Vec<Placed> {
     let lim = top - 0.5;
-    let mut out = Vec::new();
+    let mut out = Vec::with_capacity(items.len());
     for it in items {
         if let Some(r) = it.turned_bounds() {
             if r.y >= lim {
-                out.push(it.clone());
+                out.push(it);
             }
             continue;
         }
         match it {
             Placed::Lines { story, path, para, l0, l1, x, y, turn } => {
-                let Some(first) = para.lines.get(*l0) else { continue };
-                let Some(k) = (*l0..*l1).find(|k| para.lines.get(*k).is_some_and(|l| y + l.top - first.top >= lim)) else { continue };
-                let ky = para.lines.get(k).map_or(*y, |l| y + l.top - first.top);
-                out.push(Placed::Lines { story: *story, path: path.clone(), para: para.clone(), l0: k, l1: *l1, x: *x, y: ky, turn: *turn });
+                let Some(first) = para.lines.get(l0) else { continue };
+                let Some(k) = (l0..l1).find(|k| para.lines.get(*k).is_some_and(|l| y + l.top - first.top >= lim)) else { continue };
+                let ky = para.lines.get(k).map_or(y, |l| y + l.top - first.top);
+                out.push(Placed::Lines { story, path, para, l0: k, l1, x, y: ky, turn });
             }
             Placed::Fill { rect, color } if rect.bottom() > lim => {
                 let y0 = rect.y.max(top);
-                out.push(Placed::Fill { rect: Rect::new(rect.x, y0, rect.w, rect.bottom() - y0), color: *color });
+                out.push(Placed::Fill { rect: Rect::new(rect.x, y0, rect.w, rect.bottom() - y0), color });
             }
-            Placed::Rule { y0, y1, .. } if y0.max(*y1) >= lim => out.push(it.clone()),
+            Placed::Rule { y0, y1, .. } if y0.max(y1) >= lim => out.push(it),
             Placed::Image { rect, .. }
             | Placed::Shape { rect, .. }
             | Placed::Graphic { rect, .. }
@@ -341,7 +328,7 @@ pub(crate) fn from_y(items: &[Placed], top: f32) -> Vec<Placed> {
             | Placed::Object { rect, .. }
                 if rect.y >= lim =>
             {
-                out.push(it.clone())
+                out.push(it)
             }
             _ => {}
         }

crates/mcp/src/server.rs

@@ -132,7 +132,34 @@ 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;
+                        }
+                        let id = m.get("id").cloned();
+                        if let Some(r) = self.handle(m) {
+                            size = size.saturating_add(crate::tools::json_len(&r));
+                            if size > crate::tools::MAX_BATCH_OUTPUT {
+                                // This reply is the one that crosses the limit: it is not kept.
+                                let id = id.unwrap_or(Value::Null);
+                                replies.push(error(
+                                    id,
+                                    INTERNAL_ERROR,
+                                    "the batch's replies passed the output limit; this reply and the rest are left out",
+                                ));
+                            } else {
+                                replies.push(r);
+                            }
+                        }
+                    }
                     (!replies.is_empty()).then_some(Value::Array(replies))
                 }
             }

crates/mcp/src/tests.rs

@@ -367,3 +367,32 @@ fn render_page_rejects_a_page_number_that_does_not_fit() {
     let r = call(&mut s, 31, "tools/call", json!({"name": "render_page", "arguments": {"page": 1, "scale": 0.5}}));
     assert_eq!(r["result"]["isError"], false, "page 1 still renders: {r}");
 }
+
+/// A backend whose every result is 1 MiB of text.
+struct Big;
+
+impl crate::Backend for Big {
+    fn call(&mut self, _: &str, _: Value) -> Result<Value, String> {
+        Ok(json!("x".repeat(1 << 20)))
+    }
+    fn has_ui(&self) -> bool {
+        false
+    }
+    fn describe(&self) -> String {
+        "big".into()
+    }
+}
+
+#[test]
+fn a_legacy_batch_does_not_keep_the_reply_that_crosses_the_output_budget() {
+    let n = (crate::tools::MAX_BATCH_OUTPUT >> 20) + 5;
+    let one =
+        |i: usize| json!({"jsonrpc": "2.0", "id": i, "method": "tools/call", "params": {"name": "execute", "arguments": {"command": "edit.copy"}}});
+    let line = Value::Array((0..n).map(one).collect()).to_string();
+    let mut s = Server::new(Box::new(Big));
+    let out = s.handle_line(&line).expect("reply");
+    assert!(out.len() <= crate::tools::MAX_BATCH_OUTPUT + (1 << 16), "{}", out.len());
+    let replies: Vec<Value> = serde_json::from_str(&out).unwrap();
+    assert_eq!(replies.len(), n);
+    assert!(replies.iter().any(|r| r["error"]["message"].as_str().is_some_and(|m| m.contains("left out"))));
+}

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,27 @@ 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 {
+                };
+                // Counted as the reply stores it: the array is text content, so its quotes and
+                // backslashes are escaped again.
+                size = size.saturating_add(json_len(&Value::String(entry.to_string())));
+                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 +416,68 @@ 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()
+        }
+    }
+
+    /// A backend whose every result is 1 MiB of quotes, which the reply escapes twice.
+    struct Quoted {
+        calls: usize,
+    }
+
+    impl Backend for Quoted {
+        fn call(&mut self, _: &str, _: Value) -> Result<Value, String> {
+            self.calls += 1;
+            Ok(json!("\"".repeat(1 << 20)))
+        }
+        fn has_ui(&self) -> bool {
+            false
+        }
+        fn describe(&self) -> String {
+            "quoted".into()
+        }
+    }
+
+    #[test]
+    fn a_batch_budgets_the_escaped_text_it_returns() {
+        // Each result is 2 MiB as JSON and 4 MiB once the array is stored as text.
+        let commands: Vec<Value> = (0..40).map(|_| json!({"command": "edit.copy"})).collect();
+        let mut b = Quoted { calls: 0 };
+        let r = call_tool(&mut b, "batch", &json!({"commands": commands, "keepGoing": true}));
+        assert!(r.is_error);
+        assert!(b.calls <= (MAX_BATCH_OUTPUT >> 22) + 1, "{} commands ran", b.calls);
+    }
+
+    #[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() {