lean_ctx_server/tools/registered/
ctx_shape.rs

1// SPDX-License-Identifier: Apache-2.0
2use rmcp::ErrorData;
3use rmcp::model::Tool;
4use serde_json::{Map, Value, json};
5
6use crate::server::tool_trait::{McpTool, ToolContext, ToolOutput};
7use lean_ctx::core::output_shapers::OutputShapers;
8use lean_ctx::tool_defs::tool_def;
9
10/// Shapes the output of a host-native tool *after* it ran (concept K2,
11/// "shape, don't redirect"): a host mod (Claude Code `tool.call`) forwards the
12/// native Bash stdout here instead of denying the call and forcing a round
13/// trip to `ctx_shell`. The text it rewrites gets lean-ctx's secret redaction,
14/// and the result flows through the same dispatch pipeline as every `ctx_*`
15/// output (sensitivity floor, policy-pack redaction, input filters).
16///
17/// Host-only surface: a host extension calls it through `ctx_call` (`$.mcp.call`
18/// in Claude Code 2.1.296+ only reaches listed tools, #2089); it is never
19/// advertised, and every gate refusal on it is an error so the host keeps the
20/// native output — see `dynamic_tools::INTERNAL_HOST_TOOLS`.
21pub struct CtxShapeTool;
22
23/// Below this many **bytes** the shaping gain cannot pay for the call. Bytes,
24/// not chars: the Claude Code mod measures JavaScript string length (UTF-16
25/// units), and a text is never fewer bytes than UTF-16 units — so anything the
26/// mod sends clears this floor instead of slipping back unredacted.
27const MIN_SHAPE_BYTES: usize = 2_000;
28/// Above this the input is returned unchanged: bounds memory and CPU for a
29/// single call (native outputs that large are persisted by the host anyway).
30const MAX_SHAPE_BYTES: usize = 8 * 1024 * 1024;
31
32impl McpTool for CtxShapeTool {
33    fn name(&self) -> &'static str {
34        "ctx_shape"
35    }
36
37    /// The body replaces a native tool's output verbatim, so the pipeline must
38    /// hand it back without hints, nudges or checkpoints — after its redaction,
39    /// sensitivity and filter passes, which machine-readable calls keep.
40    fn produces_machine_readable(&self, _args: Option<&Map<String, Value>>) -> bool {
41        true
42    }
43
44    fn tool_def(&self) -> Tool {
45        tool_def(
46            "ctx_shape",
47            "Internal host hook: shape a native tool's output after it ran. Not for agents.",
48            json!({
49                "type": "object",
50                "properties": {
51                    "tool": { "type": "string", "description": "Native tool name, e.g. Bash" },
52                    "command": { "type": "string", "description": "Exact command line (Bash)" },
53                    "exit_code": { "type": "integer", "description": "The command's exit code, when known" },
54                    "output": { "type": "string", "description": "The native tool's output as the model would read it" }
55                },
56                "required": ["tool", "output"]
57            }),
58        )
59    }
60
61    fn handle(
62        &self,
63        args: &Map<String, Value>,
64        _ctx: &ToolContext,
65    ) -> Result<ToolOutput, ErrorData> {
66        let tool = args
67            .get("tool")
68            .and_then(Value::as_str)
69            .ok_or_else(|| ErrorData::invalid_params("tool is required", None))?;
70        let output = args
71            .get("output")
72            .and_then(Value::as_str)
73            .ok_or_else(|| ErrorData::invalid_params("output is required", None))?;
74        let command = args.get("command").and_then(Value::as_str);
75        let exit_code = args
76            .get("exit_code")
77            .and_then(Value::as_i64)
78            .and_then(|c| i32::try_from(c).ok());
79
80        let shaped = shape(tool, command, exit_code, output);
81        let original_tokens = lean_ctx::core::tokens::count_tokens(output);
82        let shaped_tokens = lean_ctx::core::tokens::count_tokens(&shaped);
83        Ok(ToolOutput {
84            text: shaped,
85            original_tokens,
86            saved_tokens: original_tokens.saturating_sub(shaped_tokens),
87            mode: Some("shape".to_string()),
88            path: None,
89            changed: false,
90            shell_outcome: None,
91            content_blocks: None,
92        })
93    }
94}
95
96/// Pure shaping decision, separated from MCP plumbing for testing. Shapes
97/// with the registered `core::output_shapers`.
98pub fn shape(tool: &str, command: Option<&str>, exit_code: Option<i32>, output: &str) -> String {
99    shape_with(
100        lean_ctx::core::output_shapers::output_shapers(),
101        tool,
102        command,
103        exit_code,
104        output,
105    )
106}
107
108/// `shape` with explicit shapers; `None` (nothing registered) hands the
109/// redacted output back unshaped.
110fn shape_with(
111    shapers: Option<&OutputShapers>,
112    tool: &str,
113    command: Option<&str>,
114    exit_code: Option<i32>,
115    output: &str,
116) -> String {
117    if output.len() < MIN_SHAPE_BYTES || output.len() > MAX_SHAPE_BYTES {
118        return output.to_string();
119    }
120    let redacted = super::ctx_shell_background::redact_shell_output_secrets(output);
121    let Some(shapers) = shapers else {
122        return redacted;
123    };
124    let shaped = match (tool, command) {
125        ("Bash", Some(cmd)) if !cmd.trim().is_empty() => {
126            (shapers.command)(cmd, &redacted, exit_code)
127        }
128        _ => (shapers.tool)(&redacted, Some(tool)),
129    };
130    // Never hand back more than the host already has.
131    if shaped.len() >= redacted.len() {
132        redacted
133    } else {
134        shaped
135    }
136}
137
138#[cfg(test)]
139mod tests {
140    use super::*;
141
142    #[test]
143    fn outside_the_size_window_output_is_returned_verbatim() {
144        assert_eq!(shape("Bash", Some("ls"), Some(0), "a\nb\n"), "a\nb\n");
145        let huge = "x\n".repeat(MAX_SHAPE_BYTES / 2 + 1);
146        assert_eq!(shape("Bash", Some("cat big"), Some(0), &huge), huge);
147    }
148
149    /// Without registered shapers (library use, tests) the redacted output
150    /// comes back unshaped, never a panic. The shaping path itself is driven
151    /// with the proxy's shapers in `proxy::shape_route_tests`.
152    #[test]
153    fn without_registered_shapers_the_redacted_output_is_returned() {
154        use std::fmt::Write as _;
155        let mut noisy = String::new();
156        for i in 0..400 {
157            let _ = writeln!(noisy, "   Compiling crate-{i} v0.1.{i} (/tmp/x)");
158        }
159        let expected = super::super::ctx_shell_background::redact_shell_output_secrets(&noisy);
160        assert_eq!(
161            shape_with(None, "Bash", Some("cargo build"), Some(0), &noisy),
162            expected
163        );
164        assert_eq!(shape_with(None, "Grep", None, None, &noisy), expected);
165    }
166}
lean_ctx_server/tools/registered/ctx_shape.rs

lean_ctx_server/tools/registered/
ctx_shape.rs
