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>] [--reference <path>]`.
4
5#![allow(clippy::disallowed_methods)] // serde_json::json! macro expands to .unwrap() internally
6
7use crate::tools::subprocess::run_apr;
8use crate::types::{InputSchema, ToolCallResult, ToolDefinition};
9
10/// Tool name registered with MCP clients.
11pub const NAME: &str = "apr.trace";
12
13/// Return the MCP tool definition for `apr.trace`.
14///
15/// FALSIFY-MCP-008: the `inputSchema` is parsed from the build-time codegen
16/// constant `crate::schemas::APR_TRACE_SCHEMA`, which `build.rs` emits from
17/// `contracts/apr-mcp-tool-schemas-v1.yaml`. The contract is the single
18/// source of truth — the live `tools/list` response and the YAML must agree
19/// byte-for-byte after JSON canonicalization (asserted by
20/// `tests/falsify_mcp_008.rs`).
21#[must_use]
22pub fn trace_tool_definition() -> ToolDefinition {
23    let input_schema: InputSchema = serde_json::from_str(crate::schemas::APR_TRACE_SCHEMA).expect(
24        "FALSIFY-MCP-008: apr.trace codegen constant must parse as InputSchema; \
25             regenerate by editing contracts/apr-mcp-tool-schemas-v1.yaml and rebuilding",
26    );
27    ToolDefinition {
28        name: NAME.to_string(),
29        description: crate::schemas::APR_TRACE_DESCRIPTION.to_string(),
30        input_schema,
31    }
32}
33
34/// Execute `apr.trace` by spawning `apr trace <model> --json [...flags]`.
35#[must_use]
36pub fn call(args: &serde_json::Value) -> ToolCallResult {
37    let Some(model_path) = args.get("model_path").and_then(|v| v.as_str()) else {
38        return ToolCallResult::error("Missing required argument: model_path");
39    };
40
41    let mut owned: Vec<String> = vec![
42        "trace".to_string(),
43        model_path.to_string(),
44        "--json".to_string(),
45    ];
46
47    if let Some(pat) = args.get("layer").and_then(|v| v.as_str()) {
48        if !pat.is_empty() {
49            owned.push("--layer".to_string());
50            owned.push(pat.to_string());
51        }
52    }
53    if let Some(ref_path) = args.get("reference").and_then(|v| v.as_str()) {
54        if !ref_path.is_empty() {
55            owned.push("--reference".to_string());
56            owned.push(ref_path.to_string());
57        }
58    }
59
60    let argv: Vec<&str> = owned.iter().map(String::as_str).collect();
61    run_apr(&argv)
62}
63
64/// HELIX-IDEA-002 — unified-signature shim for the inventory dispatcher.
65pub fn dispatch(
66    args: &serde_json::Value,
67    _cancel: &std::sync::mpsc::Receiver<()>,
68    _sink: Option<&crate::server::NotificationSink>,
69    _token: Option<serde_json::Value>,
70) -> ToolCallResult {
71    call(args)
72}
73
74crate::register_mcp_tool!(
75    name: NAME,
76    definition: trace_tool_definition,
77    dispatch: dispatch,
78);
79
80#[cfg(test)]
81#[allow(clippy::disallowed_methods)]
82mod tests {
83    use super::*;
84
85    #[test]
86    fn definition_has_correct_name_and_required_field() {
87        let def = trace_tool_definition();
88        assert_eq!(def.name, "apr.trace");
89        assert_eq!(def.input_schema.schema_type, "object");
90        assert_eq!(def.input_schema.required, vec!["model_path".to_string()]);
91        for field in ["model_path", "layer", "reference"] {
92            assert!(
93                def.input_schema.properties.contains_key(field),
94                "property {field} present"
95            );
96        }
97    }
98
99    #[test]
100    fn missing_model_path_returns_error() {
101        let result = call(&serde_json::json!({}));
102        assert_eq!(result.is_error, Some(true));
103        assert!(result.content[0].text.contains("model_path"));
104    }
105}
aprender_mcp/tools/trace.rs

aprender_mcp/tools/
trace.rs
