aprender_mcp/tools/
registry.rs

1//! HELIX-IDEA-002 — link-time tool registration via `inventory`.
2//!
3//! Contract: `contracts/apr-mcp-tool-inventory-v1.yaml`.
4//!
5//! Each `tools/<tool>.rs` submits one [`McpToolEntry`] to the global
6//! inventory using the [`register_mcp_tool!`] macro. The dispatcher
7//! ([`AprMcpServer`](crate::server::AprMcpServer)) iterates
8//! `inventory::iter::<McpToolEntry>` once at startup and builds:
9//!
10//! * a `Vec<ToolDefinition>` (replaces the hardcoded vec at
11//!   `server.rs:221-233`),
12//! * a `BTreeMap<&str, DispatchFn>` (replaces the hardcoded match arms
13//!   at `server.rs:461-483`).
14//!
15//! Adding a new tool is therefore a one-file edit under `tools/`: write
16//! the dispatcher, the definition factory, and one
17//! `register_mcp_tool!` invocation. No more edits to `server.rs`.
18//!
19//! `inputSchema` remains contracts-driven (FALSIFY-MCP-008); this
20//! module only owns the *registration* path.
21
22use crate::server::NotificationSink;
23use crate::types::{ToolCallResult, ToolDefinition};
24use std::collections::BTreeMap;
25use std::sync::mpsc;
26
27/// Unified dispatch signature shared by every MCP tool. Tools that do
28/// not need cancel / sink simply ignore those parameters.
29///
30/// * `args` — the `params.arguments` JSON value from the MCP request.
31/// * `cancel_rx` — an mpsc channel signalled by `notifications/cancelled`.
32///   Honoured today only by `apr.run`; other tools accept it for ABI
33///   uniformity.
34/// * `sink` — `Some` when the client supplied `params._meta.progressToken`
35///   AND the tool supports streaming. Honoured today by `apr.finetune`
36///   and `apr.run`; other tools accept it for ABI uniformity.
37/// * `progress_token` — the JSON value of `params._meta.progressToken`
38///   passed back verbatim in each `notifications/progress`.
39pub type DispatchFn = fn(
40    &serde_json::Value,
41    &mpsc::Receiver<()>,
42    Option<&NotificationSink>,
43    Option<serde_json::Value>,
44) -> ToolCallResult;
45
46/// Submitted by every tool module via [`register_mcp_tool!`]. The
47/// dispatcher reads these out of [`inventory`] at startup.
48#[derive(Debug)]
49pub struct McpToolEntry {
50    /// MCP tool name advertised in `tools/list` and matched in
51    /// `tools/call`. Examples: `"apr.version"`, `"apr.run"`.
52    pub name: &'static str,
53    /// Returns the full [`ToolDefinition`] (name + description + input
54    /// schema). Called once per `tools/list` request.
55    pub definition_fn: fn() -> ToolDefinition,
56    /// Returns the [`ToolCallResult`] for a `tools/call` request.
57    pub dispatch_fn: DispatchFn,
58}
59
60inventory::collect!(McpToolEntry);
61
62/// One-time index built from the global inventory at startup. Holds
63/// definitions in registration order (so `tools/list` output is
64/// deterministic across builds — `inventory::iter` itself orders entries
65/// by submission link order, but we sort by name to remove that
66/// dependency for tests) and a name → dispatch map for `tools/call`.
67#[derive(Debug)]
68pub struct ToolIndex {
69    definitions: Vec<ToolDefinition>,
70    dispatch: BTreeMap<&'static str, DispatchFn>,
71}
72
73impl ToolIndex {
74    /// Build the index from `inventory::iter::<McpToolEntry>()`.
75    ///
76    /// FALSIFY-INVENTORY-002: panics with a clear diagnostic if two
77    /// entries share the same `name`. The panic fires the first time
78    /// any test or production path constructs an `AprMcpServer`, so a
79    /// duplicate-name regression cannot escape CI.
80    #[must_use]
81    pub fn from_inventory() -> Self {
82        let mut by_name: BTreeMap<&'static str, &'static McpToolEntry> = BTreeMap::new();
83        for entry in inventory::iter::<McpToolEntry> {
84            if let Some(prior) = by_name.insert(entry.name, entry) {
85                // Restore the prior entry so the panic message lists both.
86                by_name.insert(prior.name, prior);
87                panic!(
88                    "FALSIFY-INVENTORY-002: duplicate MCP tool name {:?} registered twice in the \
89                     inventory. Two `register_mcp_tool!` invocations advertise the same name; \
90                     pick one. Existing: {:p}, duplicate: {:p}.",
91                    entry.name, prior, entry,
92                );
93            }
94        }
95
96        // Deterministic order — sorted by name so test goldens are stable
97        // across linker reorderings of `inventory::submit!`.
98        let mut definitions: Vec<ToolDefinition> =
99            by_name.values().map(|e| (e.definition_fn)()).collect();
100        definitions.sort_by(|a, b| a.name.cmp(&b.name));
101
102        let dispatch: BTreeMap<&'static str, DispatchFn> = by_name
103            .iter()
104            .map(|(name, entry)| (*name, entry.dispatch_fn))
105            .collect();
106
107        Self {
108            definitions,
109            dispatch,
110        }
111    }
112
113    /// All tool definitions advertised by `tools/list`, sorted by name.
114    #[must_use]
115    pub fn definitions(&self) -> &[ToolDefinition] {
116        &self.definitions
117    }
118
119    /// Look up the dispatch function for a `tools/call` request.
120    /// Returns `None` if no tool with that name is registered — the
121    /// caller is responsible for emitting the appropriate `isError`
122    /// envelope.
123    #[must_use]
124    pub fn dispatch_for(&self, name: &str) -> Option<&DispatchFn> {
125        self.dispatch.get(name)
126    }
127
128    /// Names of all registered tools, alphabetically sorted. Used by
129    /// FALSIFY-INVENTORY-001 to assert the migrated set matches the
130    /// pre-migration golden list.
131    #[must_use]
132    pub fn names(&self) -> Vec<&'static str> {
133        self.dispatch.keys().copied().collect()
134    }
135}
136
137/// Submit one [`McpToolEntry`] to the global inventory.
138///
139/// Usage (one invocation per `tools/<tool>.rs`):
140///
141/// ```ignore
142/// register_mcp_tool!(
143///     name: "apr.foo",
144///     definition: foo_tool_definition,
145///     dispatch: dispatch,
146/// );
147/// ```
148///
149/// The `dispatch` argument MUST point at a function with signature
150/// [`DispatchFn`] — typically a thin shim that calls the tool's
151/// existing `call` / `call_with_sink`.
152#[macro_export]
153macro_rules! register_mcp_tool {
154    (name: $name:expr, definition: $def:path, dispatch: $dispatch:path $(,)?) => {
155        ::inventory::submit! {
156            $crate::tools::registry::McpToolEntry {
157                name: $name,
158                definition_fn: $def,
159                dispatch_fn: $dispatch,
160            }
161        }
162    };
163}
164
165#[cfg(test)]
166mod tests {
167    use super::*;
168
169    /// Sanity: building the index from the live inventory does not panic
170    /// and yields the 9 Phase-1 tool names in alphabetical order.
171    #[test]
172    fn live_inventory_yields_phase_one_tool_set() {
173        let index = ToolIndex::from_inventory();
174        let names = index.names();
175        let expected = [
176            "apr.bench",
177            "apr.finetune",
178            "apr.qa",
179            "apr.run",
180            "apr.serve",
181            "apr.tensors",
182            "apr.trace",
183            "apr.validate",
184            "apr.version",
185        ];
186        assert_eq!(names, expected);
187        assert_eq!(index.definitions().len(), 9);
188    }
189
190    #[test]
191    fn dispatch_for_known_tool_returns_some() {
192        let index = ToolIndex::from_inventory();
193        assert!(index.dispatch_for("apr.version").is_some());
194        assert!(index.dispatch_for("apr.qa").is_some());
195    }
196
197    #[test]
198    fn dispatch_for_unknown_tool_returns_none() {
199        let index = ToolIndex::from_inventory();
200        assert!(index.dispatch_for("apr.never").is_none());
201    }
202}
aprender_mcp/tools/registry.rs

aprender_mcp/tools/
registry.rs
