From 1aa1d61bd4aeb2aa93bfc04e2a3156227631bab4 Mon Sep 17 00:00:00 2001 From: Hamza Date: Wed, 21 Oct 2026 00:29:14 +0500 Subject: [PATCH] MCP: get_node keeps to the limit too, and windows a wide node inspect_document sends a shallower tree when the document's is too large and points at get_node for the rest, so get_node was the way out of that limit -- and it had none of its own. Measured against a document with a frame of 15,000 children, through figcraft-cli mcp: inspect_document depth 64, allPages 14.29 MiB (8.00 MiB cap not applied before #915) get_node depth 64 10.51 MiB Both are over the 10 MiB buffer the report is about, and for get_node no depth makes it smaller: the width is at depth 1, so depth 1 is already 9.12 MiB of text and depth 0 has no children at all. get_node now takes the same maxBytes, cuts the subtree by depth as the document's tree is cut, and then takes a window over the children when even depth 1 is too large. The marker says how many children were shown of how many, and the offset to ask for next; `offset` is a new argument for reading the rest. After this: get_node depth 64 8.00 MiB, children 11427 of 15000 from 0, next 11427 The hint in inspect_document's marker no longer promises that get_node has no size limit, since that was the hole: it now says a wide node answers with a window and an offset. The test pages through a windowed node and asserts every child is reported exactly once, in order, with every page inside the limit. Unbounded, the first page is 23,003 bytes against a 7,667-byte limit. --- crates/mcp/src/tools.rs | 185 +++++++++++++++++++-- crates/mcp/tests/inspect_response_limit.rs | 57 +++++++ 2 files changed, 224 insertions(+), 18 deletions(-) diff --git a/crates/mcp/src/tools.rs b/crates/mcp/src/tools.rs index d04ea620..19f2be35 100644 --- a/crates/mcp/src/tools.rs +++ b/crates/mcp/src/tools.rs @@ -189,8 +189,18 @@ pub fn tool_definitions() -> Vec { tool( "get_node", "Get layer", - "One layer as JSON (same fields as inspect_document) with its children to `depth`.", - obj(json!({"id": id_schema("Layer id"), "depth": int("Child depth (default 1)")}), &["id"]), + "One layer as JSON (same fields as inspect_document) with its children to `depth`. A subtree too large for one message is reported at a shallower depth, or with a window of the children, and a `truncated` object says so and gives the `offset` to ask for next.", + obj( + json!({ + "id": id_schema("Layer id"), + "depth": int("Child depth (default 1)"), + "offset": int("First child to report (default 0), for reading a node with more children than fit in one message"), + "maxBytes": int( + "Largest response to send, in bytes of the whole JSON-RPC message (default 8388608, which fits the usual 10 MiB client buffer); 1024 of it is kept for the envelope and the request id" + ), + }), + &["id"], + ), true, ), tool( @@ -588,12 +598,38 @@ const MARKER_RESERVE: usize = 1 << 10; /// the real envelope here would mean handing the transport's request id to every tool. const ENVELOPE_RESERVE: usize = 1 << 10; +/// What the caller says it can read, as the bytes its *tree* may take: the figure is for the whole +/// message, so the envelope's share comes off it here. +fn budget(a: &Map) -> Result<(usize, usize), String> { + // 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, + }; + Ok((max_bytes, max_bytes.saturating_sub(ENVELOPE_RESERVE))) +} + +/// What a value's text block would take on the wire. It is sent as a text block, so it is measured +/// pretty-printed *and* escaped -- escaping adds about a tenth. +fn wire_size(value: &Value) -> usize { + let text = serde_json::to_string_pretty(value).unwrap_or_default(); + serde_json::to_string(&text).map_or(text.len(), |s| s.len()) +} + /// The document at `depth`, and what its text block would take on the wire. The tree is sent as a /// text block, so it is measured pretty-printed *and* escaped -- escaping adds about a tenth. fn inspect_at(b: &mut dyn Backend, depth: u64, all: bool) -> 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()); + let wire = wire_size(&value); + Ok((value, wire)) +} + +/// One layer's subtree at `depth`, and what its text block would take on the wire. +fn node_at(b: &mut dyn Backend, id: u64, depth: u64) -> Result<(Value, usize), String> { + let value = b.call("node.get", json!({"id": id, "depth": depth}))?; + let wire = wire_size(&value); Ok((value, wire)) } @@ -604,16 +640,7 @@ fn inspect_document(b: &mut dyn Backend, a: &Map) -> Result *x, Some(_) => return Err("`allPages` must be a boolean".into()), }; - // 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, - }; - - // What the tree itself may take: the caller's figure is for the whole message. - let budget = max_bytes.saturating_sub(ENVELOPE_RESERVE); + let (max_bytes, budget) = budget(a)?; let (value, wire) = inspect_at(b, depth, all)?; if wire <= budget { @@ -664,21 +691,143 @@ fn inspect_document(b: &mut dyn Backend, a: &Map) -> Result) -> Result { let id = id_arg(a.get("id").ok_or("`id` is required")?)?; let depth = n_arg(a, "depth")?.unwrap_or(1.0).clamp(0.0, 64.0) as u64; - let v = b.call("node.get", json!({"id": id, "depth": depth}))?; - if v.is_null() { + let offset = n_arg(a, "offset")?.unwrap_or(0.0).max(0.0).min(u32::MAX as f64) as usize; + let (max_bytes, budget) = budget(a)?; + + let (value, wire) = node_at(b, id, depth)?; + if value.is_null() { return Err(format!("no such layer {id}")); } - Ok(ToolResult::json(&v)) + if wire <= budget && offset == 0 { + return Ok(ToolResult::json(&value)); + } + + // The deepest subtree that fits, by bisection, as the document's tree is cut. A node asked for + // at a depth of its own children only (1) cannot be made shallower without losing them all, + // so that is where the window takes over. + let (mut used, mut value, mut used_wire) = (depth, value, wire); + if depth > 1 && offset == 0 { + let (mut lo, mut hi) = (1u64, depth - 1); + while lo <= hi { + let mid = lo + (hi - lo) / 2; + let (v, w) = node_at(b, id, mid)?; + if w + MARKER_RESERVE <= budget { + (used, value, used_wire) = (mid, v, w); + if mid == hi { + break; + } + lo = mid + 1; + } else if mid == 1 { + (used, value, used_wire) = (1, v, w); + break; + } else { + hi = mid - 1; + } + } + if used_wire + MARKER_RESERVE <= budget { + return Ok(ToolResult::json(&with_marker(value, depth, used, wire, used_wire, max_bytes, None))); + } + } + + // Still over, or a window was asked for: the children from `offset` that fit. + let total = value.get("children").and_then(Value::as_array).map_or(0, Vec::len); + let window = children_window(&mut value, offset, budget.saturating_sub(MARKER_RESERVE)); + let shown = value.get("children").and_then(Value::as_array).map_or(0, Vec::len); + let next = (offset + shown < total).then_some(offset + shown); + let used_wire = wire_size(&value); + let mut out = with_marker( + value, + depth, + used, + wire, + used_wire, + max_bytes, + Some(json!({ + "from": offset, + "shown": shown, + "total": total, + "nextOffset": next, + })), + ); + if window && next.is_none() && offset == 0 && used == depth { + // Nothing was left out after all (an offset of zero that fitted whole). + out.as_object_mut().map(|o| o.remove("truncated")); + } + Ok(ToolResult::json(&out)) +} + +/// Keep the children from `offset` that fit in `budget`, and say whether any were left out. The +/// window is found by bisection: a node with thousands of children would otherwise be serialised +/// once per child. +fn children_window(value: &mut Value, offset: usize, budget: usize) -> bool { + let Some(all) = value.get("children").and_then(Value::as_array).cloned() else { return false }; + let rest: Vec = all.iter().skip(offset).cloned().collect(); + let fits = |value: &mut Value, n: usize| { + if let Some(obj) = value.as_object_mut() { + obj.insert("children".into(), Value::Array(rest.iter().take(n).cloned().collect())); + } + wire_size(value) <= budget + }; + if fits(value, rest.len()) { + return rest.len() != all.len(); + } + // The largest window that fits, between none and all of the rest. + let (mut lo, mut hi) = (0usize, rest.len()); + while lo < hi { + let mid = lo + (hi - lo).div_ceil(2); + if fits(value, mid) { lo = mid } else { hi = mid - 1 } + } + fits(value, lo); + true +} + +/// `value` with the `truncated` object saying how it was cut. `children`: the window, when the +/// children were cut as well as (or instead of) the depth. +fn with_marker(mut value: Value, asked: u64, used: u64, asked_wire: usize, used_wire: usize, max_bytes: usize, children: Option) -> Value { + if let Some(obj) = value.as_object_mut() { + let mut marker = json!({ + "requestedDepth": asked, + "depth": used, + "requestedDepthBytes": asked_wire, + "bytes": used_wire, + "maxBytes": max_bytes, + }); + match &children { + Some(window) => { + marker["children"] = window.clone(); + marker["reason"] = json!(format!( + "the subtree at depth {asked} would be {asked_wire} bytes on the wire, over the {max_bytes}-byte limit a client can read in one message; its children are reported from {} instead", + window.get("from").and_then(Value::as_u64).unwrap_or(0) + )); + marker["hint"] = json!("ask again with the `offset` in `children.nextOffset` for the rest, or get_node({id, depth}) on one child"); + } + None => { + marker["reason"] = json!(format!( + "the subtree at depth {asked} would be {asked_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" + )); + marker["hint"] = json!("get_node({id, depth}) on a child reads deeper, one subtree at a time"); + } + } + obj.insert("truncated".into(), marker); + } + value } fn select(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 index 5e36b5a9..b6d8b92e 100644 --- a/crates/mcp/tests/inspect_response_limit.rs +++ b/crates/mcp/tests/inspect_response_limit.rs @@ -95,3 +95,60 @@ fn maxbytes_is_checked() { assert_eq!(r["isError"], true, "{r}"); assert!(r["content"][0]["text"].as_str().unwrap_or("").contains("at least 1024"), "{r}"); } + +/// `inspect_document` sends a shallower tree and points at `get_node` for the rest, so `get_node` +/// must keep to the same limit -- it was unbounded. Measured on `main`: a frame with 15,000 +/// children answers with a 10.5 MiB message at depth 1, over the 10 MiB buffer the report is +/// about, and no depth makes that smaller because the width is at depth 1. The children are +/// windowed instead, with the offset to ask for next. +#[test] +fn a_node_with_more_children_than_fit_is_windowed() { + let mut c = Client::new(); + let frame = c.create(json!({"type": "frame", "x": 0, "y": 0, "width": 400, "height": 400})); + for i in 0..40 { + c.create(json!({"type": "rectangle", "parent": frame, "x": i * 5, "y": 5, "width": 20, "height": 20})); + } + + // Whole, to measure against: every child, no marker. + let (whole, whole_wire) = get_node(&mut c, json!({"id": frame, "depth": 1})); + let total = whole["children"].as_array().expect("children").len(); + assert_eq!(total, 40, "{whole}"); + assert!(whole["truncated"].is_null(), "within the default limit: {}", whole["truncated"]); + + // A limit that holds only some of them: a window, and the marker says how many and what next. + let limit = whole_wire / 3; + assert!(limit >= 1024, "the fixture has to outgrow the floor: {whole_wire}"); + let (cut, cut_wire) = get_node(&mut c, json!({"id": frame, "depth": 1, "maxBytes": limit})); + assert!(cut_wire <= limit, "the response still exceeds maxBytes: {cut_wire} > {limit}"); + let window = &cut["truncated"]["children"]; + let shown = cut["children"].as_array().expect("children").len(); + assert_eq!(window["from"], json!(0), "{window}"); + assert_eq!(window["shown"], json!(shown), "{window}"); + assert_eq!(window["total"], json!(total), "{window}"); + assert!(shown > 0 && shown < total, "some children, not all: {shown} of {total}"); + assert_eq!(window["nextOffset"], json!(shown), "the offset to ask for next: {window}"); + // The layer itself is still there, not just a marker. + assert_eq!(cut["id"], whole["id"]); + + // Paging through with that offset reaches every child exactly once, in order. + let mut seen: Vec = Vec::new(); + let mut offset = Some(0u64); + while let Some(at) = offset { + let (page, wire) = get_node(&mut c, json!({"id": frame, "depth": 1, "maxBytes": limit, "offset": at})); + assert!(wire <= limit, "page at {at} exceeds maxBytes: {wire}"); + seen.extend(page["children"].as_array().cloned().unwrap_or_default()); + offset = page["truncated"]["children"]["nextOffset"].as_u64(); + assert!(seen.len() <= total, "a page repeated a child: {} of {total}", seen.len()); + } + assert_eq!(seen.len(), total, "every child was reported once"); + assert_eq!(seen, whole["children"].as_array().cloned().unwrap_or_default(), "in the same order"); +} + +/// The response and what its text block takes on the wire, for one `get_node` call. +fn get_node(c: &mut Client, args: Value) -> (Value, usize) { + let r = c.call("get_node", 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() + ENVELOPE; + (serde_json::from_str(text).expect("the block is JSON"), wire) +}