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#[cfg(feature = "apr-tools")]
49#[derive(Debug)]
50pub struct McpToolEntry {
51    /// MCP tool name advertised in `tools/list` and matched in
52    /// `tools/call`. Examples: `"apr.version"`, `"apr.run"`.
53    pub name: &'static str,
54    /// Returns the full [`ToolDefinition`] (name + description + input
55    /// schema). Called once per `tools/list` request.
56    pub definition_fn: fn() -> ToolDefinition,
57    /// Returns the [`ToolCallResult`] for a `tools/call` request.
58    pub dispatch_fn: DispatchFn,
59}
60
61#[cfg(feature = "apr-tools")]
62inventory::collect!(McpToolEntry);
63
64/// One-time index built from the global inventory at startup. Holds
65/// definitions in registration order (so `tools/list` output is
66/// deterministic across builds — `inventory::iter` itself orders entries
67/// by submission link order, but we sort by name to remove that
68/// dependency for tests) and a name → dispatch map for `tools/call`.
69#[derive(Debug)]
70pub struct ToolIndex {
71    definitions: Vec<ToolDefinition>,
72    /// Keyed by an OWNED name (PMAT-3954): a caller-supplied tool's name is a `String`
73    /// in its [`ToolDefinition`], never `'static`.
74    dispatch: BTreeMap<String, DispatchFn>,
75}
76
77impl ToolIndex {
78    /// Build an index from a caller-supplied tool set (PMAT-3954) — a downstream server
79    /// serving ITS OWN tools on this crate, with no `inventory` and no apr tool linked in.
80    ///
81    /// Definitions are sorted by name, as [`Self::from_inventory`] sorts them.
82    ///
83    /// # Errors
84    ///
85    /// Two entries with one name are refused, naming the tool — the same rule as
86    /// FALSIFY-INVENTORY-002, returned rather than panicked because the caller built it.
87    pub fn from_entries(entries: Vec<(ToolDefinition, DispatchFn)>) -> Result<Self, String> {
88        let mut definitions = Vec::with_capacity(entries.len());
89        let mut dispatch: BTreeMap<String, DispatchFn> = BTreeMap::new();
90        for (definition, dispatch_fn) in entries {
91            if dispatch
92                .insert(definition.name.clone(), dispatch_fn)
93                .is_some()
94            {
95                return Err(format!(
96                    "duplicate MCP tool name {:?}: two entries advertise the same name",
97                    definition.name
98                ));
99            }
100            definitions.push(definition);
101        }
102        definitions.sort_by(|a, b| a.name.cmp(&b.name));
103        Ok(Self {
104            definitions,
105            dispatch,
106        })
107    }
108
109    /// Build the index from `inventory::iter::<McpToolEntry>()`.
110    ///
111    /// FALSIFY-INVENTORY-002: panics with a clear diagnostic if two
112    /// entries share the same `name`. The panic fires the first time
113    /// any test or production path constructs an `AprMcpServer`, so a
114    /// duplicate-name regression cannot escape CI.
115    #[cfg(feature = "apr-tools")]
116    #[must_use]
117    pub fn from_inventory() -> Self {
118        let mut by_name: BTreeMap<&'static str, &'static McpToolEntry> = BTreeMap::new();
119        for entry in inventory::iter::<McpToolEntry> {
120            if let Some(prior) = by_name.insert(entry.name, entry) {
121                // Restore the prior entry so the panic message lists both.
122                by_name.insert(prior.name, prior);
123                panic!(
124                    "FALSIFY-INVENTORY-002: duplicate MCP tool name {:?} registered twice in the \
125                     inventory. Two `register_mcp_tool!` invocations advertise the same name; \
126                     pick one. Existing: {:p}, duplicate: {:p}.",
127                    entry.name, prior, entry,
128                );
129            }
130        }
131
132        // Deterministic order — sorted by name so test goldens are stable
133        // across linker reorderings of `inventory::submit!`.
134        let mut definitions: Vec<ToolDefinition> =
135            by_name.values().map(|e| (e.definition_fn)()).collect();
136        definitions.sort_by(|a, b| a.name.cmp(&b.name));
137
138        let dispatch: BTreeMap<String, DispatchFn> = by_name
139            .iter()
140            .map(|(name, entry)| ((*name).to_string(), entry.dispatch_fn))
141            .collect();
142
143        Self {
144            definitions,
145            dispatch,
146        }
147    }
148
149    /// All tool definitions advertised by `tools/list`, sorted by name.
150    #[must_use]
151    pub fn definitions(&self) -> &[ToolDefinition] {
152        &self.definitions
153    }
154
155    /// Look up the dispatch function for a `tools/call` request.
156    /// Returns `None` if no tool with that name is registered — the
157    /// caller is responsible for emitting the appropriate `isError`
158    /// envelope.
159    #[must_use]
160    pub fn dispatch_for(&self, name: &str) -> Option<&DispatchFn> {
161        self.dispatch.get(name)
162    }
163
164    /// Names of all registered tools, alphabetically sorted. Used by
165    /// FALSIFY-INVENTORY-001 to assert the migrated set matches the
166    /// pre-migration golden list.
167    #[must_use]
168    pub fn names(&self) -> Vec<&str> {
169        self.dispatch.keys().map(String::as_str).collect()
170    }
171}
172
173/// Submit one [`McpToolEntry`] to the global inventory.
174///
175/// Usage (one invocation per `tools/<tool>.rs`):
176///
177/// ```ignore
178/// register_mcp_tool!(
179///     name: "apr.foo",
180///     definition: foo_tool_definition,
181///     dispatch: dispatch,
182/// );
183/// ```
184///
185/// The `dispatch` argument MUST point at a function with signature
186/// [`DispatchFn`] — typically a thin shim that calls the tool's
187/// existing `call` / `call_with_sink`.
188#[cfg(feature = "apr-tools")]
189#[macro_export]
190macro_rules! register_mcp_tool {
191    (name: $name:expr, definition: $def:path, dispatch: $dispatch:path $(,)?) => {
192        ::inventory::submit! {
193            $crate::tools::registry::McpToolEntry {
194                name: $name,
195                definition_fn: $def,
196                dispatch_fn: $dispatch,
197            }
198        }
199    };
200}
201
202#[cfg(all(test, feature = "apr-tools"))] // from_inventory exists only with apr's tools
203mod tests {
204    use super::*;
205
206    /// Sanity: building the index from the live inventory does not panic
207    /// and yields the 9 Phase-1 tool names in alphabetical order.
208    #[test]
209    fn live_inventory_yields_phase_one_tool_set() {
210        let index = ToolIndex::from_inventory();
211        let names = index.names();
212        let expected = [
213            "apr.bench",
214            "apr.finetune",
215            "apr.qa",
216            "apr.run",
217            "apr.serve",
218            "apr.tensors",
219            "apr.trace",
220            "apr.validate",
221            "apr.version",
222        ];
223        assert_eq!(names, expected);
224        assert_eq!(index.definitions().len(), 9);
225    }
226
227    #[test]
228    fn dispatch_for_known_tool_returns_some() {
229        let index = ToolIndex::from_inventory();
230        assert!(index.dispatch_for("apr.version").is_some());
231        assert!(index.dispatch_for("apr.qa").is_some());
232    }
233
234    #[test]
235    fn dispatch_for_unknown_tool_returns_none() {
236        let index = ToolIndex::from_inventory();
237        assert!(index.dispatch_for("apr.never").is_none());
238    }
239}
aprender_mcp/tools/registry.rs

aprender_mcp/tools/
registry.rs
