diff --git a/crates/mcp/src/tools.rs b/crates/mcp/src/tools.rs index 92773e03..d04ea620 100644 --- a/crates/mcp/src/tools.rs +++ b/crates/mcp/src/tools.rs @@ -178,7 +178,9 @@ pub fn tool_definitions() -> Vec { 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)"), + "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" + ), }), &[], ), @@ -578,6 +580,14 @@ const MAX_RESPONSE: usize = 8 << 20; /// otherwise be able to push a response that only just fitted back over the limit. const MARKER_RESERVE: usize = 1 << 10; +/// Room kept for everything the text block is wrapped in: the content object and `isError`, the +/// JSON-RPC envelope (`jsonrpc`, `id`, `result`) and the newline the transport writes after it. +/// About ninety bytes of that is fixed and the rest is the request id, so `maxBytes` holds for +/// any id up to roughly 900 bytes -- which is every id in practice, a number or a UUID. An id +/// larger than that is the one case where the message can pass the caller's figure, and sizing +/// the real envelope here would mean handing the transport's request id to every tool. +const ENVELOPE_RESERVE: usize = 1 << 10; + /// 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> { @@ -602,37 +612,45 @@ fn inspect_document(b: &mut dyn Backend, a: &Map) -> Result 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 (value, wire) = inspect_at(b, depth, all)?; - if wire <= max_bytes || depth == 0 { + if wire <= budget { 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 { + let (used, mut value, used_wire) = if depth == 0 { + // Depth 0 is the shallowest tree there is, so there is nothing smaller to send: it goes + // out over the limit with the marker saying so, the same as a bisection that reaches + // depth 0. An error would leave the caller with nothing at all. + (0, value, wire) + } else { + // 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 <= budget { + best = Some((mid, v, w)); + if mid == hi { + break; + } + lo = mid + 1; + } else if mid == 0 { + best = Some((0, v, w)); break; + } else { + hi = mid - 1; } - 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)); + match best { + Some(found) => found, + None => return Ok(ToolResult::json(&value)), + } }; if let Some(obj) = value.as_object_mut() { obj.insert( diff --git a/crates/mcp/tests/inspect_response_limit.rs b/crates/mcp/tests/inspect_response_limit.rs index 0d9f9e6d..5e36b5a9 100644 --- a/crates/mcp/tests/inspect_response_limit.rs +++ b/crates/mcp/tests/inspect_response_limit.rs @@ -6,6 +6,10 @@ mod common; use common::Client; use serde_json::{Value, json}; +/// The allowance `inspect_document` keeps for the content object, the JSON-RPC envelope and the +/// request id (`ENVELOPE_RESERVE`). +const ENVELOPE: usize = 1024; + /// A document whose tree is far larger than the limits used below. fn busy() -> Client { let mut c = Client::new(); @@ -49,8 +53,29 @@ fn a_tree_too_large_for_the_client_is_reported_shallower_and_says_so() { 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"]); + // maxBytes is the whole message, so the tree fits under it with room for the envelope, not + // at it: a text block of exactly maxBytes would be sent inside a larger message than that. + let (whole, _) = inspect(&mut c, json!({"depth": 6, "allPages": true, "maxBytes": full_wire + ENVELOPE})); + assert!(whole["truncated"].is_null(), "a limit that leaves room for the envelope is not a truncation: {}", whole["truncated"]); + let (tight, _) = inspect(&mut c, json!({"depth": 6, "allPages": true, "maxBytes": full_wire})); + assert_eq!(tight["truncated"]["requestedDepth"], 6, "a limit at the exact size of the tree leaves none for the envelope: {tight}"); +} + +/// Every over-limit response carries the marker, however the depth was arrived at: a caller that +/// asked for depth 0 has no other way to learn that its limit was passed. +#[test] +fn a_depth_0_response_over_the_limit_says_so_too() { + let mut c = busy(); + let (deep, _) = inspect(&mut c, json!({"depth": 0, "allPages": true})); + assert!(deep["truncated"].is_null(), "within the default limit: {}", deep["truncated"]); + let (v, wire) = inspect(&mut c, json!({"depth": 0, "allPages": true, "maxBytes": 1024})); + let t = &v["truncated"]; + assert_eq!((t["requestedDepth"].as_u64(), t["depth"].as_u64()), (Some(0), Some(0)), "{t}"); + // `bytes` is the tree, measured before the marker was added to it, so the message that goes + // out is larger still -- there being nothing smaller to send. + assert_eq!(t["bytes"], t["requestedDepthBytes"], "{t}"); + assert!(wire > 1024, "the depth-0 tree is sent over the limit, with the marker saying so: {wire}"); + assert!(!v["pages"].is_null(), "the tree is still sent, there being nothing smaller: {v}"); } #[test]