Skip to main content

trusty_memory/tools/
mod.rs

1//! MCP tool surface for trusty-memory.
2//!
3//! Why: Concentrates the public tool contract in one file so changes are
4//! auditable and the MCP schema stays in sync with the implementation.
5//! What: Defines `MemoryMcpServer`, `tool_definitions()` (the MCP
6//! `tools/list` payload), and the in-process tool dispatcher wired to the
7//! real `PalaceRegistry` + retrieval / KG APIs.
8//! Test: `cargo test -p trusty-memory-mcp` validates the schema and dispatch.
9//!
10//! Tools exposed:
11//! - `memory_remember(palace, text, room?, tags?)` -> drawer_id
12//! - `memory_recall(palace, query, top_k?)`        -> Vec<Drawer> (L0+L1+L2)
13//! - `memory_recall_deep(palace, query, top_k?)`   -> Vec<Drawer> (L3 deep)
14//! - `memory_list(palace, room?, tag?, limit?)`    -> Vec<Drawer>
15//! - `memory_forget(palace, drawer_id)`            -> status: deleted|not_found
16//! - `palace_create(name, description?)`           -> PalaceId
17//! - `palace_list()`                                -> Vec<PalaceId>
18//! - `palace_info(palace)`                          -> palace metadata + stats
19//! - `room_list(palace)`                            -> Vec<RoomSummary>
20//! - `room_create(palace, label, description?)`     -> room_id (idempotent)
21//! - `room_rename(palace, room, new_label)`         -> renamed room
22//! - `kg_assert(palace, subject, predicate, object, confidence?, provenance?)` -> ()
23//! - `kg_retract_triple(palace, subject, predicate, object)` -> closed count
24//! - `kg_query(palace, subject)`                    -> Vec<Triple>
25//! - `kg_list_subjects(palace, limit?, with_counts?)` -> subjects (#4776)
26//! - `wing_list(palace)`                            -> Vec<WingSummary>
27//! - `wing_create(palace, label)`                   -> wing_id (idempotent)
28//! - `wing_rename(palace, wing, new_label)`         -> WingSummary
29
30pub mod bm25;
31// #7493: the serialized-response byte ceiling every result-returning tool
32// folds to, applied once in `dispatch_tool`.
33mod byte_cap;
34mod chat_assets;
35pub mod chat_definitions;
36pub mod chat_ops;
37pub mod definitions;
38pub mod dream_ops;
39// #5000 / #4786: answer "is this findable?" per id and per estate.
40pub mod embed_audit;
41pub mod embed_audit_definitions;
42pub mod helpers;
43pub mod kg_ops;
44pub mod memory_ops;
45// #6318: the no-palace fallback every palace-scoped READ tool shares.
46pub mod palace_index;
47pub mod palace_ops;
48// Owner ruling 2026-09-14: the recall read family, split out of `memory_ops`
49// at the 500-SLOC cap, plus the projection that shapes what a hit shows.
50pub mod recall_ops;
51pub mod recall_projection;
52pub mod room_definitions;
53pub mod room_ops;
54pub mod task_definitions;
55pub mod task_ops;
56// ADR-0027 T9 (#4809): the wing surface ships WITH the wing entity — a level
57// nobody reads is the defect the ADR exists to correct.
58pub mod wing_definitions;
59pub mod wing_ops;
60
61// Re-export the public + cross-module surface so external call sites
62// (`crate::tools::X`) and the `super::*` glob in `tools::tests` keep
63// resolving exactly as they did against the former monolithic module.
64pub use bm25::{spawn_bm25_index_worker, Bm25IndexRequest, BM25_INDEX_QUEUE_CAPACITY};
65pub use definitions::{tool_definitions, tool_definitions_with, MemoryMcpServer};
66pub(crate) use helpers::{auto_extract_and_assert, room_label};
67
68// Re-exports used only by the in-crate test module (`super::*`).
69#[cfg(test)]
70pub(crate) use bm25::{bm25_hits_to_recall_results, bm25_index_enqueue};
71#[cfg(test)]
72pub(crate) use helpers::{blocklist_gate, content_gate, dedup_gate, open_palace_handle};
73
74use crate::AppState;
75use anyhow::Result;
76use serde_json::Value;
77
78use chat_ops::{
79    handle_chat_session_add_turn, handle_chat_session_create, handle_chat_session_delete,
80    handle_chat_session_get, handle_chat_session_list, handle_chat_session_recall,
81    handle_chat_turn_append,
82};
83use dream_ops::{handle_dream_consolidate_room, handle_palace_dream};
84use kg_ops::{
85    handle_add_alias, handle_discover_aliases, handle_get_prompt_context, handle_kg_assert,
86    handle_kg_bootstrap, handle_kg_gaps, handle_kg_list_subjects, handle_kg_query,
87    handle_kg_retract_triple, handle_list_prompt_facts, handle_remove_prompt_fact,
88    handle_upgrade_tool,
89};
90use memory_ops::{
91    handle_memory_forget, handle_memory_list, handle_memory_note, handle_memory_remember,
92    handle_memory_send_message,
93};
94use palace_ops::{
95    handle_palace_compact, handle_palace_create, handle_palace_delete, handle_palace_info,
96    handle_palace_list, handle_palace_reembed, handle_palace_unalias, handle_palace_update,
97};
98use recall_ops::{handle_memory_recall, handle_memory_recall_all, handle_memory_recall_deep};
99use room_ops::{handle_room_create, handle_room_list, handle_room_rename};
100use task_ops::{handle_task_add, handle_task_complete, handle_task_list};
101use wing_ops::{handle_wing_create, handle_wing_list, handle_wing_rename};
102
103/// Dispatch a tool call by name to its real handler.
104///
105/// Why: Centralises the name → handler mapping; every handler now performs a
106/// real read/write against the live `PalaceRegistry` instead of returning a
107/// stub. After issue #227 the body is a thin router — every tool's logic
108/// lives in its own `handle_*` function above so the dispatcher itself is
109/// auditable at a glance.
110/// What: Returns `Ok(Value)` on success, `Err` on unknown tool / bad args /
111/// underlying failure.
112/// Test: `dispatch_palace_create_persists`, `dispatch_remember_then_recall`,
113/// `dispatch_kg_assert_then_query`, `dispatch_unknown_tool_errors`.
114///
115/// #7493: the serialized-response byte ceiling is enforced HERE, after the
116/// handler and before the caller, so a result-returning tool cannot ship
117/// unbounded by forgetting to fold its own body. [`byte_cap::apply`] is a
118/// no-op for every uncapped tool; a measurement failure propagates as a tool
119/// error rather than returning an unmeasured body.
120pub async fn dispatch_tool(state: &AppState, name: &str, args: Value) -> Result<Value> {
121    // #6424: a successful recall, remember or note is what the console's Last
122    // Used column means by "used". The args are cloned only for those four
123    // tools, since the dispatch below consumes them.
124    let stamp_args = USE_STAMPING_TOOLS.contains(&name).then(|| args.clone());
125    // #7493: the fold reads `max_bytes`/`full` off the caller's arguments,
126    // which the dispatch below also consumes.
127    let cap_args = byte_cap::is_capped(name).then(|| args.clone());
128    let result = dispatch_tool_inner(state, name, args).await;
129    if let (Ok(_), Some(args)) = (&result, &stamp_args) {
130        stamp_palace_use(state, args, name);
131    }
132    let mut value = result?;
133    if let Some(args) = cap_args {
134        byte_cap::apply(name, &args, &mut value)?;
135    }
136    Ok(value)
137}
138
139/// The tools whose success counts as using a palace (#6424).
140///
141/// Why: the column answers "when did anyone last put something in or take
142/// something out of this palace". Housekeeping — `palace_info`, `console_metrics`,
143/// the embed-audit sweeps that open every palace on disk — is not use, and
144/// stamping on it would make every palace look equally fresh forever.
145/// What: the read and write verbs a caller reaches for deliberately.
146/// `memory_recall_all` is absent on purpose: it spans every palace, so it says
147/// nothing about any one of them.
148/// Test: `dispatch_remember_and_recall_stamp_last_used`,
149/// `dispatch_palace_info_does_not_stamp_last_used` in `tools::tests`.
150const USE_STAMPING_TOOLS: &[&str] = &[
151    "memory_remember",
152    "memory_note",
153    "memory_recall",
154    "memory_recall_deep",
155];
156
157/// Record the palace a just-completed tool call used (#6424).
158///
159/// Why: one place, so no handler can drift into its own cadence — the throttle
160/// and its cost live in [`crate::palace_last_used`].
161/// What: resolves the same palace id the handler resolved, follows a palace
162/// alias to its target the way the handler's own open did, takes the data
163/// directory off the already-resident handle via `peek` (no open, no eviction,
164/// no I/O), and hands both to the throttled stamp. Every step is best-effort:
165/// an unresolvable palace, a handle the registry has since evicted, or a failed
166/// write all leave the column stale rather than fail the caller's operation,
167/// which has already succeeded.
168///
169/// Dropping the alias step breaks the column outright (#6424 review), which is
170/// why it is here. `resolve_palace` returns the caller's raw `palace` argument,
171/// but
172/// `PalaceRegistry::open_palace_bounded` registers the handle under
173/// `resolve_palace_alias`'s CANONICAL id, and `peek` is a bare LRU lookup that
174/// resolves nothing. Keying on the raw string therefore missed on every
175/// alias-addressed call, and the column never advanced for as long as callers
176/// used the alias.
177/// Test: `dispatch_remember_and_recall_stamp_last_used`,
178/// `an_alias_addressed_recall_stamps_the_canonical_palace`.
179fn stamp_palace_use(state: &AppState, args: &Value, tool: &str) {
180    let Ok(palace) = helpers::resolve_palace(state, args, tool) else {
181        return;
182    };
183    // The same rule `PalaceRegistry::resolve_palace_alias` applies, from the
184    // one place that owns it — a second spelling here is how the two drift.
185    let palace = trusty_common::palace_alias::alias_target_if_absent(&state.data_root, &palace)
186        .unwrap_or(palace);
187    let id = trusty_common::memory_core::PalaceId::new(&palace);
188    let Some(data_dir) = state.registry.peek(&id).and_then(|h| h.data_dir.clone()) else {
189        return;
190    };
191    crate::palace_last_used::stamp(
192        &state.palace_last_used,
193        &data_dir,
194        &palace,
195        crate::palace_last_used::now_unix(),
196    );
197}
198
199/// The name -> handler match, unwrapped from [`dispatch_tool`]'s stamping.
200async fn dispatch_tool_inner(state: &AppState, name: &str, args: Value) -> Result<Value> {
201    match name {
202        "memory_remember" => handle_memory_remember(state, args).await,
203        "memory_note" => handle_memory_note(state, args).await,
204        "memory_recall" => handle_memory_recall(state, args).await,
205        "memory_recall_deep" => handle_memory_recall_deep(state, args).await,
206        "palace_create" => handle_palace_create(state, args).await,
207        "palace_list" => handle_palace_list(state, args).await,
208        "palace_delete" => handle_palace_delete(state, args).await,
209        "palace_update" => handle_palace_update(state, args).await,
210        "kg_assert" => handle_kg_assert(state, args).await,
211        // The inverse of `kg_assert`: closes one (subject, predicate, object)
212        // and leaves the pair's other objects live.
213        "kg_retract_triple" => handle_kg_retract_triple(state, args).await,
214        "add_alias" => handle_add_alias(state, args).await,
215        "list_prompt_facts" => handle_list_prompt_facts(state, args).await,
216        "remove_prompt_fact" => handle_remove_prompt_fact(state, args).await,
217        "kg_query" => handle_kg_query(state, args).await,
218        // #4776: subject discovery — the read that makes `kg_query` usable
219        // without already knowing a subject.
220        "kg_list_subjects" => handle_kg_list_subjects(state, args).await,
221        "memory_list" => handle_memory_list(state, args).await,
222        "memory_forget" => handle_memory_forget(state, args).await,
223        "palace_info" => handle_palace_info(state, args).await,
224        "palace_compact" => handle_palace_compact(state, args).await,
225        // #4906: report / repair drawers that have no vector.
226        "palace_reembed" => handle_palace_reembed(state, args).await,
227        // #5005: free drawers destroyed by a vector-id collision.
228        "palace_unalias" => handle_palace_unalias(state, args).await,
229        // #5000: verify a caller's OWN drawer ids, not the whole missing set.
230        "palace_verify_embedded" => embed_audit::handle_palace_verify_embedded(state, args).await,
231        // #5000 / #4786: every palace on disk, uncapped — the console report is
232        // capped at 20 and shows an uncached palace as 0/0, i.e. healthy.
233        "palace_embed_sweep" => embed_audit::handle_palace_embed_sweep(state, args).await,
234        "kg_gaps" => handle_kg_gaps(state, args).await,
235        "memory_recall_all" => handle_memory_recall_all(state, args).await,
236        "get_prompt_context" => handle_get_prompt_context(state, args).await,
237        "discover_aliases" => handle_discover_aliases(state, args).await,
238        "kg_bootstrap" => handle_kg_bootstrap(state, args).await,
239        "memory_send_message" => handle_memory_send_message(state, args).await,
240        "upgrade" => handle_upgrade_tool(state, args).await,
241        "console_metrics" => crate::console_metrics::handle_console_metrics(state, args).await,
242        "chat_asset_capabilities" | "chat_asset_put" | "chat_asset_get" => {
243            chat_assets::handle(state, name, args).await
244        }
245        "chat_session_create" => handle_chat_session_create(state, args).await,
246        "chat_session_add_turn" => handle_chat_session_add_turn(state, args).await,
247        "chat_session_get" => handle_chat_session_get(state, args).await,
248        "chat_session_recall" => handle_chat_session_recall(state, args).await,
249        "chat_session_list" => handle_chat_session_list(state, args).await,
250        "chat_session_delete" => handle_chat_session_delete(state, args).await,
251        "chat_turn_append" => handle_chat_turn_append(state, args).await,
252        "dream_consolidate_room" => handle_dream_consolidate_room(state, args).await,
253        "palace_dream" => handle_palace_dream(state, args).await,
254        "task_add" => handle_task_add(state, args).await,
255        "task_list" => handle_task_list(state, args).await,
256        "task_complete" => handle_task_complete(state, args).await,
257        // ADR-0027 T6 (#4805): the room surface — discovery, idempotent
258        // creation, and the rename that repairs an `unresolved-*` label.
259        "room_list" => handle_room_list(state, args).await,
260        "room_create" => handle_room_create(state, args).await,
261        "room_rename" => handle_room_rename(state, args).await,
262        // ADR-0027 T9 (#4809): the wing surface — the scope axis over rooms.
263        "wing_list" => handle_wing_list(state, args).await,
264        "wing_create" => handle_wing_create(state, args).await,
265        "wing_rename" => handle_wing_rename(state, args).await,
266        other => anyhow::bail!("unknown tool: {other}"),
267    }
268}
269
270#[cfg(test)]
271mod tests;
272// #7493: the serialized-response byte ceiling, its truncation notice, and the
273// descriptor/router consistency the enforcement table owns.
274#[cfg(test)]
275mod byte_cap_tests;