From b1f0eafbf7e226b29ecb211b2fe8cee1efce00e5 Mon Sep 17 00:00:00 2001 From: Hamza Date: Tue, 20 Oct 2026 16:29:32 +0500 Subject: [PATCH] MCP: inspect_document answers within what a client can read MCP stdio is newline-delimited JSON, so one tool result is one message and has to fit the client's read buffer whole; there is no streaming of part of one. The TypeScript SDK's stdio transport defaults to 10 MiB and hosts rarely raise it, so a document whose tree went past that killed the server: the client stopped reading, the pipe broke, and every later call in the session failed. The tree is now reported at the deepest depth that fits, found by bisection, with a `truncated` object giving the depth asked for, the depth used, both sizes, the limit and a pointer to get_node for reading one subtree deeply. `maxBytes` lets a caller that knows its own buffer say so; the default is 8 MiB, leaving the envelope 2 MiB of the usual 10. The size is measured as the transport sends it -- pretty-printed and then escaped inside the text block, which adds about a tenth (342,136 bytes of tree became 379,954 on the wire in the fixture). A kilobyte is reserved for the marker itself, which is added after the measurement. A limit below even the shallowest tree still answers, at depth 0, with the marker saying so: there is nothing smaller to send, and an error would leave the caller with nothing at all. The tool description now says that each level of depth multiplies the response and points at get_node, which the report notes would have been enough on its own to avoid the crash. Fixes #912. --- crates/mcp/src/tools.rs | 87 +++++++++++++++++++++- crates/mcp/tests/inspect_response_limit.rs | 72 ++++++++++++++++++ 2 files changed, 156 insertions(+), 3 deletions(-) create mode 100644 crates/mcp/tests/inspect_response_limit.rs diff --git a/crates/mcp/src/tools.rs b/crates/mcp/src/tools.rs index 7714da5a..92773e03 100644 --- a/crates/mcp/src/tools.rs +++ b/crates/mcp/src/tools.rs @@ -173,9 +173,13 @@ pub fn tool_definitions() -> Vec { tool( "inspect_document", "Inspect document", - "The active document as JSON: name, path, dirty flag, current page, selection, the current page's layer tree (id, name, type, x/y relative to the parent frame (groups don't count), width, height, absoluteBounds in page coordinates, fills, strokes, effects, cornerRadius, layout, sizing, text properties, component links, overrides, bound variables), other pages summarised, styles, variable count and the undo history. Use it to verify every change.", + "The active document as JSON: name, path, dirty flag, current page, selection, the current page's layer tree (id, name, type, x/y relative to the parent frame (groups don't count), width, height, absoluteBounds in page coordinates, fills, strokes, effects, cornerRadius, layout, sizing, text properties, component links, overrides, bound variables), other pages summarised, styles, variable count and the undo history. Use it to verify every change. Each level of `depth` multiplies the response, and `allPages` multiplies it again: on a large document the tree is reported at a shallower depth than asked for, with a `truncated` object saying so. Use get_node to read one subtree deeply instead.", obj( - json!({"depth": int("Tree depth for the current page (default 4)"), "allPages": boolean("Expand every page, not only the current one")}), + json!({ + "depth": int("Tree depth for the current page (default 4)"), + "allPages": boolean("Expand every page, not only the current one"), + "maxBytes": int("Largest response to send, in bytes on the wire (default 8388608, which fits the usual 10 MiB client buffer)"), + }), &[], ), true, @@ -562,6 +566,27 @@ fn set_properties(b: &mut dyn Backend, a: &Map) -> Result Result<(Value, usize), String> { + let value = b.call("document.inspect", json!({"depth": depth, "allPages": all}))?; + let text = serde_json::to_string_pretty(&value).unwrap_or_default(); + let wire = serde_json::to_string(&text).map_or(text.len(), |s| s.len()); + Ok((value, wire)) +} + fn inspect_document(b: &mut dyn Backend, a: &Map) -> Result { let depth = n_arg(a, "depth")?.unwrap_or(4.0).clamp(0.0, 64.0) as u64; let all = match a.get("allPages") { @@ -569,7 +594,63 @@ fn inspect_document(b: &mut dyn Backend, a: &Map) -> Result *x, Some(_) => return Err("`allPages` must be a boolean".into()), }; - b.call("document.inspect", json!({"depth": depth, "allPages": all})).map(|v| ToolResult::json(&v)) + // A caller that knows its own buffer can say so; 1 KiB is below the shallowest tree, so the + // floor only keeps the arithmetic sane. + let max_bytes = match n_arg(a, "maxBytes")? { + Some(v) if v < 1024.0 => return Err("`maxBytes` must be at least 1024".into()), + Some(v) => v.min(usize::MAX as f64) as usize, + None => MAX_RESPONSE, + }; + + let (value, wire) = inspect_at(b, depth, all)?; + if wire <= max_bytes || depth == 0 { + return Ok(ToolResult::json(&value)); + } + + // The deepest tree that fits, by bisection: a deeper tree holds every node of a shallower one, + // so the size does not fall as the depth grows. Bisecting keeps this to about six more walks + // of the document rather than one per level. + let (mut lo, mut hi) = (0u64, depth - 1); + let mut best: Option<(u64, Value, usize)> = None; + while lo <= hi { + let mid = lo + (hi - lo) / 2; + let (v, w) = inspect_at(b, mid, all)?; + if w + MARKER_RESERVE <= max_bytes { + best = Some((mid, v, w)); + if mid == hi { + break; + } + lo = mid + 1; + } else if mid == 0 { + // Not even the shallowest tree fits. Send it anyway and say so: there is nothing + // smaller to send, and an error would leave the caller with nothing at all. + best = Some((0, v, w)); + break; + } else { + hi = mid - 1; + } + } + + let Some((used, mut value, used_wire)) = best else { + return Ok(ToolResult::json(&value)); + }; + if let Some(obj) = value.as_object_mut() { + obj.insert( + "truncated".into(), + json!({ + "requestedDepth": depth, + "depth": used, + "requestedDepthBytes": wire, + "bytes": used_wire, + "maxBytes": max_bytes, + "reason": format!( + "the tree at depth {depth} would be {wire} bytes on the wire, over the {max_bytes}-byte limit a client can read in one message; it is reported at depth {used} instead" + ), + "hint": "get_node({id, depth}) walks one subtree at a time, so a deep tree can be read in parts without a size limit", + }), + ); + } + Ok(ToolResult::json(&value)) } fn get_node(b: &mut dyn Backend, a: &Map) -> Result { diff --git a/crates/mcp/tests/inspect_response_limit.rs b/crates/mcp/tests/inspect_response_limit.rs new file mode 100644 index 00000000..0d9f9e6d --- /dev/null +++ b/crates/mcp/tests/inspect_response_limit.rs @@ -0,0 +1,72 @@ +//! `inspect_document` must not emit a message larger than a client can read in one piece: MCP +//! stdio is newline-delimited JSON, so one result is one message, and a client whose read buffer +//! it exceeds closes the pipe and kills the server for the rest of the session (#912). +#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)] +mod common; +use common::Client; +use serde_json::{Value, json}; + +/// A document whose tree is far larger than the limits used below. +fn busy() -> Client { + let mut c = Client::new(); + for i in 0..40 { + let frame = c.create(json!({"type": "frame", "x": i * 10, "y": 0, "width": 200, "height": 200})); + for j in 0..12 { + c.create(json!({"type": "rectangle", "parent": frame, "x": j * 5, "y": 5, "width": 20, "height": 20})); + } + } + c +} + +/// The response and what its text block takes on the wire, escaped as the transport sends it. +fn inspect(c: &mut Client, args: Value) -> (Value, usize) { + let r = c.call("inspect_document", args.clone()); + assert_eq!(r["isError"], false, "{args}: {}", r["content"][0]["text"]); + let text = r["content"][0]["text"].as_str().expect("a text block"); + let wire = serde_json::to_string(&text).unwrap().len(); + (serde_json::from_str(text).expect("the block is JSON"), wire) +} + +#[test] +fn a_tree_too_large_for_the_client_is_reported_shallower_and_says_so() { + let mut c = busy(); + let (full, full_wire) = inspect(&mut c, json!({"depth": 6, "allPages": true})); + assert!(full["truncated"].is_null(), "a tree within the limit is not truncated: {}", full["truncated"]); + assert!(full_wire > 30_000, "the fixture has to outgrow the limits below: {full_wire}"); + + // Asked for more than the caller can read: a shallower tree, inside the limit, and a marker. + let (cut, cut_wire) = inspect(&mut c, json!({"depth": 6, "allPages": true, "maxBytes": 30_000})); + assert!(cut_wire <= 30_000, "the response still exceeds maxBytes: {cut_wire}"); + let t = &cut["truncated"]; + assert_eq!(t["requestedDepth"], 6, "{t}"); + assert!(t["depth"].as_u64().unwrap() < 6, "the depth was reduced: {t}"); + assert_eq!(t["maxBytes"], 30_000, "{t}"); + assert_eq!(t["requestedDepthBytes"], full_wire, "the marker says what the full tree would have taken: {t}"); + // The marker points at the way to read the rest. + assert!(t["hint"].as_str().unwrap_or("").contains("get_node"), "{t}"); + // Still a usable document, not just a marker: the page and its top-level layers are there. + assert!(!cut["currentPage"].is_null(), "{cut}"); + assert!(cut["pages"][0]["children"].as_array().is_some_and(|c| !c.is_empty()), "{cut}"); + + // The deepest tree that fits is the one sent: a limit that admits the whole tree is not cut. + let (whole, _) = inspect(&mut c, json!({"depth": 6, "allPages": true, "maxBytes": full_wire})); + assert!(whole["truncated"].is_null(), "a limit at the exact size is not a truncation: {}", whole["truncated"]); +} + +#[test] +fn a_limit_below_the_shallowest_tree_still_answers() { + // Nothing smaller than depth 0 exists, so it is sent with the marker rather than an error: + // an error would leave the caller with nothing at all. + let mut c = busy(); + let (v, _) = inspect(&mut c, json!({"depth": 6, "allPages": true, "maxBytes": 1024})); + assert_eq!(v["truncated"]["depth"], 0, "{}", v["truncated"]); + assert_eq!(v["truncated"]["requestedDepth"], 6, "{}", v["truncated"]); +} + +#[test] +fn maxbytes_is_checked() { + let mut c = Client::new(); + let r = c.call("inspect_document", json!({"depth": 1, "maxBytes": 10})); + assert_eq!(r["isError"], true, "{r}"); + assert!(r["content"][0]["text"].as_str().unwrap_or("").contains("at least 1024"), "{r}"); +}