aprender_mcp/tools/
trace.rs

1//! `apr.trace` — M2 tool. Layer-by-layer tensor trace for debugging a model.
2//!
3//! Wraps `apr trace <model> --json [--layer <pat>]`.
4//!
5//! #2407: the tool used to advertise a `reference` argument in its
6//! `inputSchema` and forward it as `apr trace --reference <path>`, which is a
7//! stub: it printed `{"comparison": "reference comparison not yet
8//! implemented"}` and exited 0. The wrapper drops stderr on success, so an
9//! MCP client saw a plain success result for a comparison that never
10//! happened. `reference` is no longer advertised, and supplying it now
11//! returns `isError`.
12
13#![allow(clippy::disallowed_methods)] // serde_json::json! macro expands to .unwrap() internally
14
15use crate::tools::args::{self, try_arg};
16use crate::tools::subprocess::run_apr;
17use crate::types::{InputSchema, ToolCallResult, ToolDefinition};
18
19/// Tool name registered with MCP clients.
20pub const NAME: &str = "apr.trace";
21
22/// Return the MCP tool definition for `apr.trace`.
23///
24/// FALSIFY-MCP-008: the `inputSchema` is parsed from the build-time codegen
25/// constant `crate::schemas::APR_TRACE_SCHEMA`, which `build.rs` emits from
26/// `contracts/apr-mcp-tool-schemas-v1.yaml`. The contract is the single
27/// source of truth — the live `tools/list` response and the YAML must agree
28/// byte-for-byte after JSON canonicalization (asserted by
29/// `tests/falsify_mcp_008.rs`).
30#[must_use]
31pub fn trace_tool_definition() -> ToolDefinition {
32    let input_schema: InputSchema = serde_json::from_str(crate::schemas::APR_TRACE_SCHEMA).expect(
33        "FALSIFY-MCP-008: apr.trace codegen constant must parse as InputSchema; \
34             regenerate by editing contracts/apr-mcp-tool-schemas-v1.yaml and rebuilding",
35    );
36    ToolDefinition {
37        name: NAME.to_string(),
38        description: crate::schemas::APR_TRACE_DESCRIPTION.to_string(),
39        input_schema,
40    }
41}
42
43/// Build the `apr trace ...` argv from `tools/call` arguments.
44///
45/// # Errors
46/// Returns the client-facing message when an argument is present but not
47/// usable at its declared type.
48pub fn build_argv(args: &serde_json::Value) -> Result<Vec<String>, String> {
49    let model_path = args::required_str(args, "model_path")?;
50
51    let mut owned: Vec<String> = vec![
52        "trace".to_string(),
53        model_path.to_string(),
54        "--json".to_string(),
55    ];
56
57    if let Some(pat) = args::opt_str(args, "layer")? {
58        if !pat.is_empty() {
59            owned.push("--layer".to_string());
60            owned.push(pat.to_string());
61        }
62    }
63    if let Some(ref_path) = args::opt_str(args, "reference")? {
64        if !ref_path.is_empty() {
65            owned.push("--reference".to_string());
66            owned.push(ref_path.to_string());
67        }
68    }
69    Ok(owned)
70}
71
72/// Execute `apr.trace` by spawning `apr trace <model> --json [...flags]`.
73#[must_use]
74pub fn call(args: &serde_json::Value) -> ToolCallResult {
75    let owned = try_arg!(build_argv(args));
76
77    let argv: Vec<&str> = owned.iter().map(String::as_str).collect();
78    run_apr(&argv)
79}
80
81/// HELIX-IDEA-002 — unified-signature shim for the inventory dispatcher.
82pub fn dispatch(
83    args: &serde_json::Value,
84    _cancel: &std::sync::mpsc::Receiver<()>,
85    _sink: Option<&crate::server::NotificationSink>,
86    _token: Option<serde_json::Value>,
87) -> ToolCallResult {
88    call(args)
89}
90
91crate::register_mcp_tool!(
92    name: NAME,
93    definition: trace_tool_definition,
94    dispatch: dispatch,
95);
96
97#[cfg(test)]
98#[allow(clippy::disallowed_methods)]
99mod tests {
100    use super::*;
101
102    /// #2407: this test used to require `reference` to be an advertised
103    /// property, which is what made the unimplemented option discoverable in
104    /// the first place. It now asserts the opposite: a client must not be
105    /// invited to pass an argument the tool cannot honour.
106    #[test]
107    fn definition_has_correct_name_and_required_field() {
108        let def = trace_tool_definition();
109        assert_eq!(def.name, "apr.trace");
110        assert_eq!(def.input_schema.schema_type, "object");
111        assert_eq!(def.input_schema.required, vec!["model_path".to_string()]);
112        for field in ["model_path", "layer"] {
113            assert!(
114                def.input_schema.properties.contains_key(field),
115                "property {field} present"
116            );
117        }
118        assert!(
119            !def.input_schema.properties.contains_key("reference"),
120            "`reference` is not implemented and must not be advertised"
121        );
122    }
123
124    #[test]
125    fn missing_model_path_returns_error() {
126        let result = call(&serde_json::json!({}));
127        assert_eq!(result.is_error, Some(true));
128        assert!(result.content[0].text.contains("model_path"));
129    }
130
131    /// #2403 — `{"layer": 7, "reference": 42}` dropped BOTH flags and traced
132    /// the whole model against nothing, reported as a success.
133    #[test]
134    fn integer_layer_and_reference_are_errors_not_dropped_flags() {
135        let result = call(&serde_json::json!({ "model_path": "m.gguf", "layer": 7 }));
136        assert_eq!(result.is_error, Some(true));
137        assert!(result.content[0].text.contains("layer"));
138
139        let result = call(&serde_json::json!({ "model_path": "m.gguf", "reference": 42 }));
140        assert_eq!(result.is_error, Some(true));
141        assert!(result.content[0].text.contains("reference"));
142    }
143
144    #[test]
145    fn string_layer_and_reference_reach_the_cli() {
146        let argv = build_argv(&serde_json::json!({
147            "model_path": "m.gguf",
148            "layer": "blk.7",
149            "reference": "ref.gguf"
150        }))
151        .expect("strings are usable");
152        assert_eq!(
153            argv,
154            vec![
155                "trace",
156                "m.gguf",
157                "--json",
158                "--layer",
159                "blk.7",
160                "--reference",
161                "ref.gguf"
162            ]
163        );
164    }
165
166    #[test]
167    fn non_string_layer_is_rejected() {
168        let result = call(&serde_json::json!({
169            "model_path": "/nonexistent/model.gguf",
170            "layer": 3,
171        }));
172        assert_eq!(result.is_error, Some(true));
173        assert!(result.content[0].text.contains("Invalid layer"));
174    }
175}
aprender_mcp/tools/trace.rs

aprender_mcp/tools/
trace.rs
