aprender_mcp/tools/
bench.rs

1//! `apr.bench` — M2 tool. Benchmark inference throughput (tok/s, latency percentiles).
2//!
3//! Wraps `apr bench <model> --json [--iterations N] [--max-tokens N] [--prompt X]`.
4
5#![allow(clippy::disallowed_methods)] // serde_json::json! macro expands to .unwrap() internally
6
7use crate::tools::args::{self, try_arg};
8use crate::tools::subprocess::run_apr;
9use crate::types::{InputSchema, ToolCallResult, ToolDefinition};
10
11/// Tool name registered with MCP clients.
12pub const NAME: &str = "apr.bench";
13
14/// Return the MCP tool definition for `apr.bench`.
15///
16/// FALSIFY-MCP-008: the `inputSchema` is parsed from the build-time codegen
17/// constant `crate::schemas::APR_BENCH_SCHEMA`, which `build.rs` emits from
18/// `contracts/apr-mcp-tool-schemas-v1.yaml`. The contract is the single
19/// source of truth — the live `tools/list` response and the YAML must agree
20/// byte-for-byte after JSON canonicalization (asserted by
21/// `tests/falsify_mcp_008.rs`).
22#[must_use]
23pub fn bench_tool_definition() -> ToolDefinition {
24    let input_schema: InputSchema = serde_json::from_str(crate::schemas::APR_BENCH_SCHEMA).expect(
25        "FALSIFY-MCP-008: apr.bench codegen constant must parse as InputSchema; \
26             regenerate by editing contracts/apr-mcp-tool-schemas-v1.yaml and rebuilding",
27    );
28    ToolDefinition {
29        name: NAME.to_string(),
30        description: crate::schemas::APR_BENCH_DESCRIPTION.to_string(),
31        input_schema,
32    }
33}
34
35/// Build the `apr bench ...` argv from `tools/call` arguments.
36///
37/// # Errors
38/// Returns the client-facing message when an argument is present but not
39/// usable at its declared type.
40pub fn build_argv(args: &serde_json::Value) -> Result<Vec<String>, String> {
41    let model_path = args::required_str(args, "model_path")?;
42
43    let mut owned: Vec<String> = vec![
44        "bench".to_string(),
45        model_path.to_string(),
46        "--json".to_string(),
47    ];
48
49    if let Some(n) = args::opt_u64(args, "iterations")? {
50        owned.push("--iterations".to_string());
51        owned.push(n.to_string());
52    }
53    if let Some(n) = args::opt_u64(args, "max_tokens")? {
54        owned.push("--max-tokens".to_string());
55        owned.push(n.to_string());
56    }
57    let prompt = args::opt_str(args, "prompt")?.unwrap_or("");
58    if !prompt.is_empty() {
59        owned.push("--prompt".to_string());
60        owned.push(prompt.to_string());
61    }
62    Ok(owned)
63}
64
65/// Execute `apr.bench` by spawning `apr bench <model> --json [...flags]`.
66#[must_use]
67pub fn call(args: &serde_json::Value) -> ToolCallResult {
68    let owned = try_arg!(build_argv(args));
69    let argv: Vec<&str> = owned.iter().map(String::as_str).collect();
70    run_apr(&argv)
71}
72
73/// HELIX-IDEA-002 — unified-signature shim for the inventory dispatcher.
74pub fn dispatch(
75    args: &serde_json::Value,
76    _cancel: &std::sync::mpsc::Receiver<()>,
77    _sink: Option<&crate::server::NotificationSink>,
78    _token: Option<serde_json::Value>,
79) -> ToolCallResult {
80    call(args)
81}
82
83crate::register_mcp_tool!(
84    name: NAME,
85    definition: bench_tool_definition,
86    dispatch: dispatch,
87);
88
89#[cfg(test)]
90#[allow(clippy::disallowed_methods)]
91mod tests {
92    use super::*;
93
94    #[test]
95    fn definition_has_correct_name_and_required_field() {
96        let def = bench_tool_definition();
97        assert_eq!(def.name, "apr.bench");
98        assert_eq!(def.input_schema.schema_type, "object");
99        assert_eq!(def.input_schema.required, vec!["model_path".to_string()]);
100        for field in ["model_path", "iterations", "max_tokens", "prompt"] {
101            assert!(
102                def.input_schema.properties.contains_key(field),
103                "property {field} present"
104            );
105        }
106    }
107
108    #[test]
109    fn missing_model_path_returns_error() {
110        let result = call(&serde_json::json!({}));
111        assert_eq!(result.is_error, Some(true));
112        assert!(result.content[0].text.contains("model_path"));
113    }
114
115    /// #2403 — the audit asked for 1 iteration of 8 tokens as JSON strings and
116    /// got 5 iterations of 32 tokens (the CLI defaults), reported back as if
117    /// that is what it had asked for.
118    #[test]
119    fn string_iterations_and_max_tokens_reach_the_cli() {
120        let argv = build_argv(&serde_json::json!({
121            "model_path": "m.gguf",
122            "iterations": "1",
123            "max_tokens": "8"
124        }))
125        .expect("numeric strings are usable");
126        assert_eq!(
127            argv,
128            vec![
129                "bench",
130                "m.gguf",
131                "--json",
132                "--iterations",
133                "1",
134                "--max-tokens",
135                "8"
136            ]
137        );
138    }
139
140    #[test]
141    fn unusable_iterations_is_an_error_not_a_dropped_flag() {
142        let result = call(&serde_json::json!({ "model_path": "m.gguf", "iterations": "lots" }));
143        assert_eq!(result.is_error, Some(true));
144        assert!(result.content[0].text.contains("iterations"));
145    }
146}
aprender_mcp/tools/bench.rs

aprender_mcp/tools/
bench.rs
