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 fn value(&self) -> &Value {
32        &self.value
33    }
34
35    pub(super) fn visited_nodes(&self) -> u64 {
36        self.visited_nodes
37    }
38}
39
40impl serde::Serialize for McpResponse {
41    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
42        self.value.serialize(serializer)
43    }
44}
45impl<'de> serde::Deserialize<'de> for McpResponse {
46    fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
47        let mut value = Value::deserialize(deserializer)?;
48        Self::take(&mut value)
49            .map_err(|_| serde::de::Error::custom("recorded MCP response exceeds its bounds"))
50    }
51}
52
53// Bound retained tree data independently of wire JSON escaping. Each node/key
54// contributes a fixed allowance, and string bytes are added without rescanning.
55struct TreeBudget {
56    bytes: usize,
57    nodes: u64,
58}
59
60impl TreeBudget {
61    fn visit(&mut self, value: &Value, depth: usize) -> bool {
62        if depth >= MAX_DEPTH || !self.node() {
63            return false;
64        }
65        match value {
66            Value::Null | Value::Bool(_) | Value::Number(_) => true,
67            Value::String(value) => self.bytes(value.len()),
68            Value::Array(values) => values.iter().all(|value| self.visit(value, depth + 1)),
69            Value::Object(fields) => fields.iter().all(|(key, value)| {
70                self.node() && self.bytes(key.len()) && self.visit(value, depth + 1)
71            }),
72        }
73    }
74
75    fn node(&mut self) -> bool {
76        self.nodes = self.nodes.saturating_add(1);
77        self.nodes <= MAX_NODES && self.bytes(16)
78    }
79
80    fn bytes(&mut self, bytes: usize) -> bool {
81        self.bytes = self.bytes.saturating_add(bytes);
82        self.bytes <= MCP_MAX_RESPONSE_BYTES
83    }
84}
85
86#[cfg(test)]
87mod tests {
88    use super::*;
89
90    #[test]
91    fn admission_moves_the_tree_and_counts_payload_and_nodes() {
92        let mut value = serde_json::json!({ "x": ["a\nb", null] });
93        let mut budget = TreeBudget { bytes: 0, nodes: 0 };
94        assert!(budget.visit(&value, 0));
95        assert_eq!(budget.bytes, 5 * 16 + 1 + 3);
96        assert_eq!(budget.nodes, 5);
97        let response = McpResponse::take(&mut value).unwrap();
98        assert!(value.is_null());
99        assert_eq!(response.value()["x"][0], "a\nb");
100    }
101
102    #[test]
103    fn admission_rejects_depth_and_node_limits_without_taking_ownership() {
104        let mut nested = Value::Null;
105        for _ in 0..MAX_DEPTH {
106            nested = Value::Array(vec![nested]);
107        }
108        assert!(matches!(
109            McpResponse::take(&mut nested),
110            Err(McpCallError::ResponseTooLarge)
111        ));
112        assert!(nested.is_array());
113        let mut wide = Value::Array(vec![Value::Null; MAX_NODES as usize]);
114        assert!(McpResponse::take(&mut wide).is_err());
115    }
116}
submilli_engine/runtime/mcp/response.rs

submilli_engine/runtime/mcp/
response.rs
