submilli_engine/runtime/mcp/
response.rs

1use serde_json::Value;
2
3use super::{MCP_MAX_RESPONSE_BYTES, McpCallError};
4
5const MAX_DEPTH: usize = 128;
6const MAX_NODES: u64 = 100_000;
7
8/// A JSON tree admitted before it crosses the transport/runtime boundary.
9/// Private ownership keeps both guest conversion and recursive destruction bounded.
10#[derive(Debug)]
11pub struct McpResponse {
12    value: Value,
13    visited_nodes: u64,
14}
15
16impl McpResponse {
17    /// Validate a transport-owned tree, then move it without cloning or encoding.
18    /// On rejection the transport retains ownership. Its decoder must enforce
19    /// nesting limits while parsing, including for values it never returns.
20    pub fn take(value: &mut Value) -> Result<Self, McpCallError> {
21        let mut budget = TreeBudget { bytes: 0, nodes: 0 };
22        if !budget.visit(value, 0) {
23            return Err(McpCallError::ResponseTooLarge);
24        }
25        Ok(Self {
26            value: value.take(),
27            visited_nodes: budget.nodes,
28        })
29    }
30
31    pub(super) fn value(&self) -> &Value {
32        &self.value
33    }
34
35    pub(super) fn visited_nodes(&self) -> u64 {
36        self.visited_nodes
37    }
38}
39
40// Bound retained tree data independently of wire JSON escaping. Each node/key
41// contributes a fixed allowance, and string bytes are added without rescanning.
42struct TreeBudget {
43    bytes: usize,
44    nodes: u64,
45}
46
47impl TreeBudget {
48    fn visit(&mut self, value: &Value, depth: usize) -> bool {
49        if depth >= MAX_DEPTH || !self.node() {
50            return false;
51        }
52        match value {
53            Value::Null | Value::Bool(_) | Value::Number(_) => true,
54            Value::String(value) => self.bytes(value.len()),
55            Value::Array(values) => values.iter().all(|value| self.visit(value, depth + 1)),
56            Value::Object(fields) => fields.iter().all(|(key, value)| {
57                self.node() && self.bytes(key.len()) && self.visit(value, depth + 1)
58            }),
59        }
60    }
61
62    fn node(&mut self) -> bool {
63        self.nodes = self.nodes.saturating_add(1);
64        self.nodes <= MAX_NODES && self.bytes(16)
65    }
66
67    fn bytes(&mut self, bytes: usize) -> bool {
68        self.bytes = self.bytes.saturating_add(bytes);
69        self.bytes <= MCP_MAX_RESPONSE_BYTES
70    }
71}
72
73#[cfg(test)]
74mod tests {
75    use super::*;
76
77    #[test]
78    fn admission_moves_the_tree_and_counts_payload_and_nodes() {
79        let mut value = serde_json::json!({ "x": ["a\nb", null] });
80        let mut budget = TreeBudget { bytes: 0, nodes: 0 };
81        assert!(budget.visit(&value, 0));
82        assert_eq!(budget.bytes, 5 * 16 + 1 + 3);
83        assert_eq!(budget.nodes, 5);
84        let response = McpResponse::take(&mut value).unwrap();
85        assert!(value.is_null());
86        assert_eq!(response.value()["x"][0], "a\nb");
87    }
88
89    #[test]
90    fn admission_rejects_depth_and_node_limits_without_taking_ownership() {
91        let mut nested = Value::Null;
92        for _ in 0..MAX_DEPTH {
93            nested = Value::Array(vec![nested]);
94        }
95        assert!(matches!(
96            McpResponse::take(&mut nested),
97            Err(McpCallError::ResponseTooLarge)
98        ));
99        assert!(nested.is_array());
100        let mut wide = Value::Array(vec![Value::Null; MAX_NODES as usize]);
101        assert!(McpResponse::take(&mut wide).is_err());
102    }
103}
submilli_engine/runtime/mcp/response.rs

submilli_engine/runtime/mcp/
response.rs
