harness/tools/
bounded.rs

1//! Cross-cutting wrapper applied to every tool runtime.
2//!
3//! Mirrors the concerns MiMoCode bundles into its `Tool.wrap()`:
4//!
5//!   1. **Repair** — schema-guided input repair is the single source of
6//!      truth here (`repair_invocation`). `agent_loop` calls it BEFORE
7//!      pushing history / emitting `ToolCall` events so the recorded args
8//!      match what the inner runtime executes; the same call inside
9//!      `invoke_cancellable` is then idempotent (already-repaired input
10//!      yields `None`) and also covers bypass callers (`runner`, `mcp`).
11//!   2. **Validate** — lightweight schema validation; a violation returns a
12//!      teaching [`crate::tools::ToolFailure`] (with an `Expected shape` example) WITHOUT
13//!      reaching the inner runtime.
14//!   3. **Span** — one `tracing` span per invocation.
15//!   4. **Safety-net bound** — a success output whose serialized form blows
16//!      past a hard ceiling is clipped (error-aware head+tail), covering
17//!      tools that don't self-bound their output (MCP, custom plugins).
18
19use async_trait::async_trait;
20use serde_json::{json, Value};
21use tracing::Instrument;
22
23use crate::tool_repair::{self, ToolInputRepair};
24use crate::tools::{
25    clip_overflow, invalid_input_failure, ToolInvocation, ToolOutcome, ToolRuntime,
26    ToolRuntimeError, ToolSpec, MAX_OUTPUT_BYTES,
27};
28
29/// Hard ceiling on a single tool's serialized success output before the
30/// catch-all clip fires. Set well above [`MAX_OUTPUT_BYTES`] so that the
31/// per-tool field bounding (which already caps stdout/stderr/content near
32/// `MAX_OUTPUT_BYTES`) is never clobbered — this only rescues genuinely
33/// unbounded outputs from tools that forgot to self-limit.
34const CATCH_ALL_CEILING: usize = 4 * MAX_OUTPUT_BYTES;
35
36/// Decorates any [`ToolRuntime`] with repair + validation + tracing + a
37/// safety-net output cap. Construct via [`BoundedToolRuntime::new`].
38#[derive(Clone)]
39pub struct BoundedToolRuntime<R> {
40    inner: R,
41    /// Specs captured at construction, used only for schema lookup during
42    /// repair / validation. A tool's schema is stable for the life of the
43    /// runtime, so caching avoids re-running `inner.specs()` on the hot path.
44    specs: Vec<ToolSpec>,
45}
46
47impl<R: ToolRuntime> BoundedToolRuntime<R> {
48    pub fn new(inner: R) -> Self {
49        let specs = inner.specs();
50        Self { inner, specs }
51    }
52
53    /// Borrow the inner runtime (e.g. for downcasting / direct access in
54    /// call sites that need the concrete type).
55    pub fn inner(&self) -> &R {
56        &self.inner
57    }
58
59    fn schema_for(&self, name: &str) -> Option<&Value> {
60        self.specs
61            .iter()
62            .find(|s| s.name == name)
63            .map(|s| &s.input_schema)
64    }
65}
66
67#[async_trait]
68impl<R: ToolRuntime> ToolRuntime for BoundedToolRuntime<R> {
69    fn specs(&self) -> Vec<ToolSpec> {
70        self.inner.specs()
71    }
72
73    /// Schema-guided repair, the single source of truth. Idempotent:
74    /// re-running on already-clean input returns `None`.
75    fn repair_invocation(&self, inv: &mut ToolInvocation) -> Option<Vec<ToolInputRepair>> {
76        let schema = self.schema_for(&inv.name)?;
77        let (fixed, repairs) = tool_repair::repair_tool_input_for_spec(schema, &inv.input)?;
78        inv.input = fixed;
79        Some(repairs)
80    }
81
82    async fn invoke(&self, inv: ToolInvocation) -> Result<ToolOutcome, ToolRuntimeError> {
83        self.invoke_cancellable(inv, None).await
84    }
85
86    async fn invoke_cancellable(
87        &self,
88        mut inv: ToolInvocation,
89        cancel: Option<&tokio_util::sync::CancellationToken>,
90    ) -> Result<ToolOutcome, ToolRuntimeError> {
91        let span = tracing::info_span!("tool.invoke", tool = %inv.name, id = %inv.id);
92        async move {
93            // 1. Repair (idempotent — a no-op when agent_loop already ran it).
94            if let Some(repairs) = self.repair_invocation(&mut inv) {
95                tracing::warn!(
96                    target: "harness::tool_repair",
97                    tool = %inv.name,
98                    id = %inv.id,
99                    repairs = ?repairs,
100                    "schema-guided tool input repair applied"
101                );
102            }
103
104            // 2. Validate. A violation is a model-observable teaching failure
105            //    that never reaches the inner runtime.
106            if let Some(schema) = self.schema_for(&inv.name) {
107                if let Err(detail) = tool_repair::validate_against_schema(schema, &inv.input) {
108                    return Ok(ToolOutcome {
109                        output: Err(invalid_input_failure(
110                            &inv.name,
111                            detail,
112                            &inv.input,
113                            Some(schema),
114                        )),
115                        attachments: vec![],
116                    });
117                }
118            }
119
120            // 3. Dispatch to the inner runtime.
121            let call_id = inv.id.clone();
122            let mut outcome = self.inner.invoke_cancellable(inv, cancel).await?;
123
124            // 4. Safety-net bound for tools that don't self-limit.
125            if let Ok(value) = &outcome.output {
126                let serialized = value.to_string();
127                if serialized.len() > CATCH_ALL_CEILING {
128                    tracing::warn!(
129                        target: "harness::tool_bound",
130                        id = %call_id,
131                        bytes = serialized.len(),
132                        "tool output exceeded catch-all ceiling; clipping"
133                    );
134                    outcome.output = Ok(json!({
135                        "tool_output_clipped": true,
136                        "preview": clip_overflow(&serialized),
137                    }));
138                }
139            }
140            Ok(outcome)
141        }
142        .instrument(span)
143        .await
144    }
145}
146
147#[cfg(test)]
148mod tests {
149    use super::*;
150    use crate::tools::{ToolFailure, ToolFailureKind};
151
152    /// Inner runtime returning a fixed oversized success output to exercise
153    /// the catch-all clip.
154    struct BigOutput;
155
156    #[async_trait]
157    impl ToolRuntime for BigOutput {
158        fn specs(&self) -> Vec<ToolSpec> {
159            vec![ToolSpec {
160                name: "big".into(),
161                description: "returns a huge blob".into(),
162                input_schema: json!({
163                    "type": "object",
164                    "properties": { "n": { "type": "integer" } },
165                    "required": ["n"],
166                    "additionalProperties": false
167                }),
168            }]
169        }
170
171        async fn invoke(&self, _inv: ToolInvocation) -> Result<ToolOutcome, ToolRuntimeError> {
172            let blob = "x".repeat(CATCH_ALL_CEILING + 1_000);
173            Ok(ToolOutcome {
174                output: Ok(json!({ "blob": blob })),
175                attachments: vec![],
176            })
177        }
178    }
179
180    fn inv(name: &str, input: Value) -> ToolInvocation {
181        ToolInvocation {
182            id: "tc_1".into(),
183            name: name.into(),
184            input,
185            raw_emitted_args: None,
186        }
187    }
188
189    #[tokio::test]
190    async fn invalid_input_returns_teaching_failure_with_example() {
191        let rt = BoundedToolRuntime::new(BigOutput);
192        // Missing required `n`.
193        let out = rt.invoke(inv("big", json!({}))).await.unwrap();
194        let ToolFailure { kind, message } = out.output.unwrap_err();
195        assert_eq!(kind, ToolFailureKind::InvalidInput);
196        assert!(message.contains("Expected shape"), "msg: {message}");
197        assert!(message.contains("\"n\""), "msg: {message}");
198    }
199
200    #[tokio::test]
201    async fn oversized_success_output_is_clipped() {
202        let rt = BoundedToolRuntime::new(BigOutput);
203        let out = rt.invoke(inv("big", json!({ "n": 1 }))).await.unwrap();
204        let value = out.output.unwrap();
205        assert_eq!(value["tool_output_clipped"], true);
206        let preview = value["preview"].as_str().unwrap();
207        assert!(preview.len() < CATCH_ALL_CEILING, "preview not clipped");
208        assert!(preview.contains("output clipped"));
209    }
210
211    #[tokio::test]
212    async fn repair_invocation_is_idempotent() {
213        let rt = BoundedToolRuntime::new(BigOutput);
214        // `n` arrives as a stringified integer — repairable to a number.
215        let mut i = inv("big", json!({ "n": "5" }));
216        let first = rt.repair_invocation(&mut i);
217        assert!(first.is_some(), "expected a repair on first pass");
218        assert_eq!(i.input["n"], json!(5));
219        // Second pass over the now-clean input is a no-op.
220        assert!(rt.repair_invocation(&mut i).is_none());
221    }
222}
harness/tools/bounded.rs

harness/tools/
bounded.rs
