Skip to main content

mnemo_engine/
mcp.rs

1//! Hand-rolled Model Context Protocol (MCP) stdio server for MNemo.
2//!
3//! Exposes a `.mnemo` file as an MCP server so any MCP-compatible agent
4//! (Claude Desktop, Cursor, mcp-server, custom clients) can drive it as a
5//! tool. Framed as **JSON-RPC 2.0 with newline-delimited JSON** over
6//! stdin/stdout, per the MCP spec. Stderr is used for human-readable
7//! diagnostics only — never for protocol traffic.
8//!
9//! ## Why hand-rolled
10//!
11//! The official [`rmcp`](https://crates.io/crates/rmcp) SDK declares
12//! `edition = "2024"` on its workspace, which requires Rust 1.85+; our
13//! MSRV is 1.75 and locked (transitive deps pin us). `rmcp` also brings
14//! in tokio + async, a big architectural shift for a sync-only core.
15//! Hand-rolling stdio JSON-RPC is ~250 LoC and adds zero dependencies
16//! beyond `serde_json` (already a dep). See CHANGELOG's Phase 4 entry
17//! for the decision record.
18//!
19//! ## Protocol version
20//!
21//! We advertise MCP protocol `2025-06-18` on `initialize`. If a client
22//! requests a different version we echo their version back — the standard
23//! compatibility dance. The seven tools we expose are stable across
24//! recent protocol revisions.
25//!
26//! ## Tools
27//!
28//! - `about` — read-only briefing (manifest + onboarding memories)
29//! - `remember` — insert a memory (BYO vector; embedder integration deferred to Phase 3)
30//! - `recall` — vector recall (BYO query vector)
31//! - `forget` — delete by ULID
32//! - `list` — enumerate live memories
33//! - `snapshot_list` — enumerate PITR snapshots
34//! - `stats` — database statistics
35//!
36//! Every mutation calls `Mnemo::flush()` before returning success — MCP
37//! clients treat tool calls as atomic and expect on-disk durability.
38//!
39//! ## Passphrase
40//!
41//! Read from the `MNEMO_PASSPHRASE` environment variable only.
42//! `serve_stdio` refuses to start without it. There is no CLI-flag
43//! fallback here — that would leave the passphrase in shell history
44//! and process listings.
45
46use std::io::{BufRead, BufReader, Read, Write};
47use std::path::Path;
48
49use serde::{Deserialize, Serialize};
50use serde_json::{json, Map, Value};
51
52use crate::error::{MnemoError, Result};
53use crate::memory::{Memory, MemoryType, Metric, Scope};
54use crate::result_cache::{CachePutOpts, SemanticCachePutOpts, DEFAULT_SEMANTIC_THRESHOLD};
55use crate::store::{Mnemo, RecallRequest};
56
57/// MCP protocol version we advertise on `initialize`.
58const PROTOCOL_VERSION: &str = "2025-06-18";
59
60/// Environment variable that supplies the database passphrase.
61const PASSPHRASE_ENV: &str = "MNEMO_PASSPHRASE";
62
63/// Fuzz-target helper: parse a single JSON-RPC line as a `Request` and
64/// discard. Panic-free — malformed input returns `Err(String)`, valid
65/// requests return `Ok(())`. Exposed only via `crate::__fuzz`; not part
66/// of the stable public API.
67#[doc(hidden)]
68pub fn __fuzz_parse_request_line(line: &str) -> std::result::Result<(), String> {
69    serde_json::from_str::<Request>(line)
70        .map(|_| ())
71        .map_err(|e| e.to_string())
72}
73
74// --- JSON-RPC 2.0 wire types ---------------------------------------------
75
76/// Incoming JSON-RPC 2.0 request or notification (a notification is a
77/// request with no `id`).
78#[derive(Debug, Deserialize)]
79struct Request {
80    #[serde(default)]
81    jsonrpc: String,
82    #[serde(default)]
83    id: Option<Value>,
84    method: String,
85    #[serde(default)]
86    params: Value,
87}
88
89/// Outgoing JSON-RPC 2.0 response envelope.
90#[derive(Debug, Serialize)]
91struct Response {
92    jsonrpc: &'static str,
93    id: Value,
94    #[serde(skip_serializing_if = "Option::is_none")]
95    result: Option<Value>,
96    #[serde(skip_serializing_if = "Option::is_none")]
97    error: Option<RpcError>,
98}
99
100/// JSON-RPC 2.0 error object.
101#[derive(Debug, Serialize)]
102struct RpcError {
103    code: i32,
104    message: String,
105    #[serde(skip_serializing_if = "Option::is_none")]
106    data: Option<Value>,
107}
108
109impl RpcError {
110    fn new(code: i32, message: impl Into<String>) -> Self {
111        Self {
112            code,
113            message: message.into(),
114            data: None,
115        }
116    }
117}
118
119// Standard JSON-RPC error codes.
120const PARSE_ERROR: i32 = -32700;
121const INVALID_REQUEST: i32 = -32600;
122const METHOD_NOT_FOUND: i32 = -32601;
123const INVALID_PARAMS: i32 = -32602;
124const INTERNAL_ERROR: i32 = -32603;
125
126// --- Entry point ---------------------------------------------------------
127
128/// Run the MCP stdio server over the provided `.mnemo` file.
129///
130/// Blocks the calling thread — reads JSON-RPC requests line-by-line from
131/// stdin, dispatches to the tool implementations, writes responses to
132/// stdout, and logs diagnostics to stderr. Returns when stdin closes
133/// (typical MCP client shutdown) or on a fatal I/O error.
134///
135/// The passphrase is read from the `MNEMO_PASSPHRASE` environment
136/// variable. If unset, returns [`MnemoError::Invalid`] without opening
137/// the file.
138pub fn serve_stdio(path: &Path) -> Result<()> {
139    serve_stdio_with_embedder(path, None)
140}
141
142/// Variant that optionally registers an [`crate::EmbedCmd`] embedder on
143/// the handle at server start (Phase 3 PR 10b). The MCP `remember`,
144/// `recall`, `cache_put_semantic` and `cache_get_semantic` tools then
145/// accept a `text` param as an alternative to the explicit vector —
146/// the embedder runs once per request against the persistent child
147/// process. Pass `None` for the pre-Phase-3 vector-only behavior.
148///
149/// `embed_spec` tuple fields: `(command, name, normalize)`. See
150/// [`crate::EmbedCmd::new`] for the command format.
151pub fn serve_stdio_with_embedder(
152    path: &Path,
153    embed_spec: Option<(String, String, bool, std::time::Duration)>,
154) -> Result<()> {
155    let stdin = std::io::stdin();
156    let stdout = std::io::stdout();
157    serve_with_streams(path, stdin.lock(), stdout.lock(), embed_spec)
158}
159
160/// Run the MCP loop against explicit reader/writer streams. Extracted so
161/// integration tests can pipe scripted input without touching the real
162/// stdin/stdout. `embed_spec` is as in [`serve_stdio_with_embedder`].
163pub fn serve_with_streams<R: Read, W: Write>(
164    path: &Path,
165    reader: R,
166    mut writer: W,
167    embed_spec: Option<(String, String, bool, std::time::Duration)>,
168) -> Result<()> {
169    let passphrase = std::env::var(PASSPHRASE_ENV).map_err(|_| {
170        MnemoError::Invalid(format!(
171            "{PASSPHRASE_ENV} is required to serve — set it in the shell before starting"
172        ))
173    })?;
174
175    eprintln!(
176        "mnemo mcp: opening {} (protocol {PROTOCOL_VERSION})",
177        path.display()
178    );
179    let mut db = Mnemo::open(path, &passphrase)?;
180    if let Some((cmd, name, normalize, timeout)) = embed_spec {
181        let identity = crate::EmbedderIdentity {
182            name,
183            dimensions: db.dimensions(),
184            normalize,
185        };
186        let embedder = crate::EmbedCmd::new(&cmd, identity)?.with_timeout(timeout);
187        db.set_embedder(embedder)?;
188        eprintln!(
189            "mnemo mcp: embedder ready (text params accepted, timeout {}s)",
190            timeout.as_secs()
191        );
192    }
193    eprintln!("mnemo mcp: ready ({} live memories)", db.len());
194
195    let reader = BufReader::new(reader);
196    for line in reader.lines() {
197        let line = line.map_err(MnemoError::Io)?;
198        let trimmed = line.trim();
199        if trimmed.is_empty() {
200            continue;
201        }
202
203        let (id, out) = match serde_json::from_str::<Request>(trimmed) {
204            Ok(req) => {
205                let id = req.id.clone();
206                let is_notification = id.is_none();
207                let result = handle(&mut db, &req);
208                if is_notification {
209                    // Per JSON-RPC 2.0: notifications get no response.
210                    continue;
211                }
212                (id.unwrap_or(Value::Null), result)
213            }
214            Err(e) => (
215                Value::Null,
216                Err(RpcError::new(PARSE_ERROR, format!("parse error: {e}"))),
217            ),
218        };
219
220        let resp = match out {
221            Ok(value) => Response {
222                jsonrpc: "2.0",
223                id,
224                result: Some(value),
225                error: None,
226            },
227            Err(err) => Response {
228                jsonrpc: "2.0",
229                id,
230                result: None,
231                error: Some(err),
232            },
233        };
234
235        let line =
236            serde_json::to_string(&resp).map_err(|e| MnemoError::Serialize(e.to_string()))?;
237        writeln!(writer, "{line}").map_err(MnemoError::Io)?;
238        writer.flush().map_err(MnemoError::Io)?;
239    }
240    // Scoring v2 (v0.7.0) flipped `tool_recall` to `track_access(true)`:
241    // every recall bumps `accessed_at` / `access_count` in the catalog
242    // in memory (no per-call fsync). The previous `db.flush()` only
243    // fired on each mutating tool call's exit path, so a session that
244    // issued only `recall` / `about` / `stats` / `list` /
245    // `cache_get_semantic`-on-hit calls (read-shaped) would exit with
246    // every one of its access-stat bumps still in RAM, discarded on
247    // shutdown. Flush at shutdown so a read-only session persists its
248    // access stats too — the write is cheap (catalog page only; the
249    // record bodies aren't dirty) and guarantees the usage-saturation
250    // signal survives across sessions.
251    if let Err(e) = db.flush() {
252        eprintln!("mnemo mcp: flush on shutdown failed: {e}");
253    }
254    eprintln!("mnemo mcp: stdin closed, shutting down");
255    Ok(())
256}
257
258// --- Dispatch ------------------------------------------------------------
259
260/// Route a request to a method handler. Returns either the `result`
261/// payload for a JSON-RPC success or an [`RpcError`] for a failure.
262fn handle(db: &mut Mnemo, req: &Request) -> std::result::Result<Value, RpcError> {
263    if req.jsonrpc != "2.0" {
264        return Err(RpcError::new(
265            INVALID_REQUEST,
266            format!("expected jsonrpc \"2.0\", got {:?}", req.jsonrpc),
267        ));
268    }
269    match req.method.as_str() {
270        "initialize" => Ok(initialize_result(&req.params)),
271        "notifications/initialized" | "notifications/cancelled" => Ok(Value::Null),
272        "ping" => Ok(json!({})),
273        "tools/list" => Ok(tools_list()),
274        "tools/call" => tools_call(db, &req.params),
275        other => Err(RpcError::new(
276            METHOD_NOT_FOUND,
277            format!("method '{other}' not implemented"),
278        )),
279    }
280}
281
282/// Echo the client's requested protocol version if present; fall back to
283/// ours. Always advertise the same `serverInfo` and `capabilities` block.
284fn initialize_result(params: &Value) -> Value {
285    let client_version = params
286        .get("protocolVersion")
287        .and_then(|v| v.as_str())
288        .unwrap_or(PROTOCOL_VERSION)
289        .to_string();
290    json!({
291        "protocolVersion": client_version,
292        "capabilities": {
293            "tools": { "listChanged": false },
294        },
295        "serverInfo": {
296            "name": "mnemo",
297            "version": env!("CARGO_PKG_VERSION"),
298        },
299    })
300}
301
302// --- Tools registry ------------------------------------------------------
303
304/// The seven tools this server exposes. Kept as a hand-authored constant
305/// (rather than macro-derived) so the JSON schemas stay auditable —
306/// every field a client sees is visible right here.
307fn tools_list() -> Value {
308    json!({
309        "tools": [
310            {
311                "name": "about",
312                "description": "Read-only briefing for this .mnemo file: manifest scaffold plus any memories tagged as onboarding. Run this first to learn how the file expects to be used.",
313                "inputSchema": {
314                    "type": "object",
315                    "properties": {},
316                    "additionalProperties": false
317                }
318            },
319            {
320                "name": "remember",
321                "description": "Insert or overwrite a memory. Pass EITHER `vector` (explicit embedding) OR `text` (embedded by the server if it was started with --embed-cmd / --embed-name). Persists to disk before returning. Returns `{id}`.",
322                "inputSchema": {
323                    "type": "object",
324                    "required": ["content", "memory_type"],
325                    "properties": {
326                        "content": { "type": "string" },
327                        "memory_type": { "type": "string", "enum": ["episodic", "semantic", "procedural", "working"] },
328                        "vector": { "type": "array", "items": { "type": "number" }, "description": "Explicit embedding; mutually exclusive with `text`." },
329                        "text": { "type": "string", "description": "Content to embed via the server's registered embedder; mutually exclusive with `vector`. Requires the server to have been started with --embed-cmd." },
330                        "importance": { "type": "number", "minimum": 0, "maximum": 1 },
331                        "agent_id": { "type": "string" },
332                        "session_id": { "type": "string" },
333                        "ttl_secs": { "type": "integer" },
334                        "shared": { "type": "boolean" },
335                        "metadata": { "type": "object" }
336                    },
337                    "additionalProperties": false
338                }
339            },
340            {
341                "name": "recall",
342                "description": "Score every live memory (or the ANN-narrowed candidate set) against a query and return the top-k. Pass EITHER `query` (explicit vector) OR `text` (embedded by the server; requires --embed-cmd at serve start). Does not track access on this handle.",
343                "inputSchema": {
344                    "type": "object",
345                    "properties": {
346                        "query": { "type": "array", "items": { "type": "number" }, "description": "Explicit query vector; mutually exclusive with `text`." },
347                        "text": { "type": "string", "description": "Query text; embedded by the server's registered embedder. Mutually exclusive with `query`." },
348                        "top_k": { "type": "integer", "default": 10 },
349                        "types": { "type": "array", "items": { "type": "string" } },
350                        "agent": { "type": "string" },
351                        "metric": { "type": "string", "enum": ["cosine", "l2", "dot"], "default": "cosine" }
352                    },
353                    "additionalProperties": false
354                }
355            },
356            {
357                "name": "forget",
358                "description": "Soft-delete a memory by ULID. Space is reclaimed on the next `mnemo compact`. Persists to disk before returning.",
359                "inputSchema": {
360                    "type": "object",
361                    "required": ["id"],
362                    "properties": {
363                        "id": { "type": "string", "description": "ULID of the memory to delete" }
364                    },
365                    "additionalProperties": false
366                }
367            },
368            {
369                "name": "list",
370                "description": "Return every live memory as JSON. Sorted by creation time, oldest first. Vectors are omitted by default to keep payloads small.",
371                "inputSchema": {
372                    "type": "object",
373                    "properties": {
374                        "include_vectors": { "type": "boolean", "default": false }
375                    },
376                    "additionalProperties": false
377                }
378            },
379            {
380                "name": "snapshot_list",
381                "description": "List every restorable snapshot (one per committed flush). Each entry gives the transaction ID, creation timestamp, and memory count.",
382                "inputSchema": {
383                    "type": "object",
384                    "properties": {},
385                    "additionalProperties": false
386                }
387            },
388            {
389                "name": "stats",
390                "description": "Summary statistics: memory count, deleted count, agents present, snapshot count, page cache occupancy.",
391                "inputSchema": {
392                    "type": "object",
393                    "properties": {},
394                    "additionalProperties": false
395                }
396            },
397            {
398                "name": "cache_put",
399                "description": "Insert or overwrite an exact-key cache entry. Value is a UTF-8 string; use `content_type` to label the payload. Persists to disk before returning. Namespace and key are caller-chosen; the key is SHA-256 hashed by the engine.",
400                "inputSchema": {
401                    "type": "object",
402                    "required": ["namespace", "key", "value"],
403                    "properties": {
404                        "namespace": { "type": "string" },
405                        "key": { "type": "string" },
406                        "value": { "type": "string" },
407                        "content_type": { "type": "string", "default": "text" },
408                        "ttl_secs": { "type": "integer", "minimum": 0 }
409                    },
410                    "additionalProperties": false
411                }
412            },
413            {
414                "name": "cache_get",
415                "description": "Look up an exact-key cache entry. Returns `{value, content_type, hit}` on a hit, `{hit: false}` on miss (missing or TTL-expired). Bumps access stats catalog-only.",
416                "inputSchema": {
417                    "type": "object",
418                    "required": ["namespace", "key"],
419                    "properties": {
420                        "namespace": { "type": "string" },
421                        "key": { "type": "string" }
422                    },
423                    "additionalProperties": false
424                }
425            },
426            {
427                "name": "cache_stats",
428                "description": "Cache statistics for a namespace (omit for the whole cache). Returns entries, bytes, hits, misses, hit_rate, evictions.",
429                "inputSchema": {
430                    "type": "object",
431                    "properties": {
432                        "namespace": { "type": "string" }
433                    },
434                    "additionalProperties": false
435                }
436            },
437            {
438                "name": "cache_put_semantic",
439                "description": "Insert an entry with an embedding vector for semantic recall (Phase 10.2). Pass EITHER `vector` + `model` (explicit) OR `text` (embedded by the server; `model` defaults to the server embedder's identity name). Later cache_get_semantic calls under a different `model` won't see this entry (a different-model hit would be a semantics bug).",
440                "inputSchema": {
441                    "type": "object",
442                    "required": ["namespace", "key", "value"],
443                    "properties": {
444                        "namespace": { "type": "string" },
445                        "key": { "type": "string" },
446                        "vector": { "type": "array", "items": { "type": "number" }, "description": "Explicit embedding; mutually exclusive with `text`. Requires `model`." },
447                        "text": { "type": "string", "description": "Content to embed via the server's registered embedder; mutually exclusive with `vector`." },
448                        "value": { "type": "string" },
449                        "model": { "type": "string", "description": "Required with `vector`; optional with `text` (defaults to the embedder's identity name)." },
450                        "content_type": { "type": "string", "default": "text" },
451                        "ttl_secs": { "type": "integer", "minimum": 0 }
452                    },
453                    "additionalProperties": false
454                }
455            },
456            {
457                "name": "cache_get_semantic",
458                "description": "Semantic recall over the cache: top-1 cosine over live vectored entries in the namespace whose `model` matches. Pass EITHER `query` + `model` OR `text` (which embeds via the server and defaults the model to the embedder's identity). Returns `{hit: true, value, similarity, ...}` iff similarity >= threshold, else `{hit: false}`. Threshold defaults to 0.97.",
459                "inputSchema": {
460                    "type": "object",
461                    "required": ["namespace"],
462                    "properties": {
463                        "namespace": { "type": "string" },
464                        "query": { "type": "array", "items": { "type": "number" }, "description": "Explicit query vector; mutually exclusive with `text`. Requires `model`." },
465                        "text": { "type": "string", "description": "Query text; embedded by the server. Mutually exclusive with `query`." },
466                        "model": { "type": "string", "description": "Required with `query`; optional with `text`." },
467                        "threshold": { "type": "number", "minimum": -1.0, "maximum": 1.0 }
468                    },
469                    "additionalProperties": false
470                }
471            }
472        ]
473    })
474}
475
476// --- Tool dispatch -------------------------------------------------------
477
478/// The MCP `tools/call` method: dispatches on the caller-supplied `name`
479/// and returns the tool's `content` array wrapped in an MCP result.
480fn tools_call(db: &mut Mnemo, params: &Value) -> std::result::Result<Value, RpcError> {
481    let name = params
482        .get("name")
483        .and_then(|v| v.as_str())
484        .ok_or_else(|| RpcError::new(INVALID_PARAMS, "missing tool name"))?;
485    let empty = json!({});
486    let args = params.get("arguments").unwrap_or(&empty);
487
488    let text = match name {
489        "about" => tool_about(db)?,
490        "remember" => tool_remember(db, args)?,
491        "recall" => tool_recall(db, args)?,
492        "forget" => tool_forget(db, args)?,
493        "list" => tool_list(db, args)?,
494        "snapshot_list" => tool_snapshot_list(db)?,
495        "stats" => tool_stats(db)?,
496        "cache_put" => tool_cache_put(db, args)?,
497        "cache_get" => tool_cache_get(db, args)?,
498        "cache_stats" => tool_cache_stats(db, args)?,
499        "cache_put_semantic" => tool_cache_put_semantic(db, args)?,
500        "cache_get_semantic" => tool_cache_get_semantic(db, args)?,
501        other => {
502            return Err(RpcError::new(
503                METHOD_NOT_FOUND,
504                format!("unknown tool '{other}'"),
505            ))
506        }
507    };
508
509    Ok(json!({
510        "content": [ { "type": "text", "text": text } ],
511        "isError": false,
512    }))
513}
514
515/// Convert a [`MnemoError`] to an [`RpcError`]. Every error becomes an
516/// `INTERNAL_ERROR` with the human-readable message — MCP clients
517/// display `error.message` verbatim to the user.
518/// Resolve a text argument into a vector via the server's registered
519/// embedder (Phase 3 PR 10b). Called by the four MCP tools that
520/// accept `text` as an alternative to the explicit vector: `remember`,
521/// `recall`, `cache_put_semantic`, `cache_get_semantic`. If no
522/// embedder was configured at server start (`mnemo serve --mcp
523/// --embed-cmd ...`) the error message tells the caller how to turn
524/// it on.
525fn embed_text_via_server(
526    db: &mut Mnemo,
527    tool: &str,
528    text: &str,
529) -> std::result::Result<Vec<f32>, RpcError> {
530    let embedder = db.embedder().ok_or_else(|| {
531        RpcError::new(
532            INVALID_PARAMS,
533            format!(
534                "{tool}: 'text' requires an embedder on the server. \
535                 Start with: mnemo serve --mcp <file> --embed-cmd '<cmd>' --embed-name <name>"
536            ),
537        )
538    })?;
539    let mut vs = embedder.embed(&[text]).map_err(mnemo_err)?;
540    vs.pop().ok_or_else(|| {
541        RpcError::new(
542            INTERNAL_ERROR,
543            format!("{tool}: embedder returned no vector for the input text"),
544        )
545    })
546}
547
548fn mnemo_err(e: MnemoError) -> RpcError {
549    RpcError::new(INTERNAL_ERROR, e.to_string())
550}
551
552fn tool_about(db: &mut Mnemo) -> std::result::Result<String, RpcError> {
553    let mems = db.about().map_err(mnemo_err)?;
554    let arr: Vec<Value> = mems.iter().map(memory_to_json).collect();
555    Ok(serde_json::to_string_pretty(&arr).unwrap_or_else(|_| "[]".into()))
556}
557
558fn tool_remember(db: &mut Mnemo, args: &Value) -> std::result::Result<String, RpcError> {
559    let content = args
560        .get("content")
561        .and_then(|v| v.as_str())
562        .ok_or_else(|| RpcError::new(INVALID_PARAMS, "remember: missing 'content'"))?;
563    let mtype_str = args
564        .get("memory_type")
565        .and_then(|v| v.as_str())
566        .ok_or_else(|| RpcError::new(INVALID_PARAMS, "remember: missing 'memory_type'"))?;
567    let mtype = MemoryType::parse(mtype_str).ok_or_else(|| {
568        RpcError::new(
569            INVALID_PARAMS,
570            format!("remember: unknown memory_type '{mtype_str}'"),
571        )
572    })?;
573    // Accept either `vector` (explicit) or `text` (embedded via the
574    // server's registered embedder, Phase 3). Exactly one required.
575    let vector = match (args.get("vector"), args.get("text")) {
576        (Some(_), Some(_)) => {
577            return Err(RpcError::new(
578                INVALID_PARAMS,
579                "remember: pass EITHER 'vector' or 'text', not both",
580            ));
581        }
582        (Some(v), None) => v
583            .as_array()
584            .ok_or_else(|| {
585                RpcError::new(INVALID_PARAMS, "remember: 'vector' must be a JSON array")
586            })?
587            .iter()
588            .map(|v| v.as_f64().map(|f| f as f32))
589            .collect::<Option<Vec<f32>>>()
590            .ok_or_else(|| {
591                RpcError::new(
592                    INVALID_PARAMS,
593                    "remember: 'vector' must be an array of numbers",
594                )
595            })?,
596        (None, Some(t)) => {
597            let text = t.as_str().ok_or_else(|| {
598                RpcError::new(INVALID_PARAMS, "remember: 'text' must be a string")
599            })?;
600            embed_text_via_server(db, "remember", text)?
601        }
602        (None, None) => {
603            return Err(RpcError::new(
604                INVALID_PARAMS,
605                "remember: missing 'vector' or 'text'",
606            ))
607        }
608    };
609
610    let mut mem = Memory::new(content, mtype, vector);
611    if let Some(agent) = args.get("agent_id").and_then(|v| v.as_str()) {
612        mem = mem.with_agent(agent);
613    }
614    if let Some(sess) = args.get("session_id").and_then(|v| v.as_str()) {
615        mem = mem.with_session(sess);
616    }
617    if let Some(imp) = args.get("importance").and_then(|v| v.as_f64()) {
618        mem = mem.with_importance(imp as f32);
619    }
620    if let Some(ttl) = args.get("ttl_secs").and_then(|v| v.as_i64()) {
621        mem = mem.with_ttl(ttl);
622    }
623    if let Some(shared) = args.get("shared").and_then(|v| v.as_bool()) {
624        if shared {
625            mem = mem.with_scope(Scope::Shared);
626        }
627    }
628    if let Some(meta) = args.get("metadata").and_then(|v| v.as_object()) {
629        for (k, v) in meta {
630            mem.metadata.insert(k.clone(), v.clone());
631        }
632    }
633
634    let id = db.remember(mem).map_err(mnemo_err)?;
635    db.flush().map_err(mnemo_err)?;
636    Ok(json!({ "id": id.to_string() }).to_string())
637}
638
639fn tool_recall(db: &mut Mnemo, args: &Value) -> std::result::Result<String, RpcError> {
640    // Accept either `query` (explicit) or `text` (embedded).
641    let query = match (args.get("query"), args.get("text")) {
642        (Some(_), Some(_)) => {
643            return Err(RpcError::new(
644                INVALID_PARAMS,
645                "recall: pass EITHER 'query' or 'text', not both",
646            ))
647        }
648        (Some(v), None) => v
649            .as_array()
650            .ok_or_else(|| RpcError::new(INVALID_PARAMS, "recall: 'query' must be a JSON array"))?
651            .iter()
652            .map(|v| v.as_f64().map(|f| f as f32))
653            .collect::<Option<Vec<f32>>>()
654            .ok_or_else(|| {
655                RpcError::new(
656                    INVALID_PARAMS,
657                    "recall: 'query' must be an array of numbers",
658                )
659            })?,
660        (None, Some(t)) => {
661            let text = t
662                .as_str()
663                .ok_or_else(|| RpcError::new(INVALID_PARAMS, "recall: 'text' must be a string"))?;
664            embed_text_via_server(db, "recall", text)?
665        }
666        (None, None) => {
667            return Err(RpcError::new(
668                INVALID_PARAMS,
669                "recall: missing 'query' or 'text'",
670            ))
671        }
672    };
673    let top_k = args.get("top_k").and_then(|v| v.as_u64()).unwrap_or(10) as usize;
674    let metric = match args
675        .get("metric")
676        .and_then(|v| v.as_str())
677        .unwrap_or("cosine")
678    {
679        "cosine" => Metric::Cosine,
680        "l2" => Metric::L2,
681        "dot" => Metric::Dot,
682        other => {
683            return Err(RpcError::new(
684                INVALID_PARAMS,
685                format!("recall: unknown metric '{other}'"),
686            ))
687        }
688    };
689    // Scoring v2 (v0.7.0) flipped the default: MCP recalls DO track
690    // access, so an agent's own retrievals contribute to the usage
691    // saturation term on later queries. The bump happens in-memory
692    // (`accessed_at` / `access_count` + `dirty_catalog`), and persists
693    // through the next flush — recall itself does not fsync, so there
694    // is no per-tool-call write-amplification cost. Access stats reach
695    // disk on either (a) the next mutating tool call's exit path (every
696    // one of `remember` / `forget` / `cache_put_semantic` /
697    // `cache_get_semantic`-on-miss ends in `db.flush()`), or (b) the
698    // server's own shutdown flush when stdin closes (see
699    // `serve_with_streams`), so a read-only session that issued only
700    // recalls still persists its stats.
701    let mut req = RecallRequest::new(query)
702        .top_k(top_k)
703        .metric(metric)
704        .track_access(true);
705    if let Some(types) = args.get("types").and_then(|v| v.as_array()) {
706        let parsed: Vec<MemoryType> = types
707            .iter()
708            .filter_map(|t| t.as_str().and_then(MemoryType::parse))
709            .collect();
710        if !parsed.is_empty() {
711            req = req.types(parsed);
712        }
713    }
714    if let Some(agent) = args.get("agent").and_then(|v| v.as_str()) {
715        req = req.agent(agent.to_string());
716    }
717
718    let hits = db.recall(&req).map_err(mnemo_err)?;
719    let out: Vec<Value> = hits
720        .into_iter()
721        .map(|h| {
722            let mut obj = memory_to_json(&h.memory);
723            if let Some(map) = obj.as_object_mut() {
724                map.insert("score".into(), json!(h.score));
725                map.insert("similarity".into(), json!(h.similarity));
726            }
727            obj
728        })
729        .collect();
730    Ok(serde_json::to_string_pretty(&out).unwrap_or_else(|_| "[]".into()))
731}
732
733fn tool_forget(db: &mut Mnemo, args: &Value) -> std::result::Result<String, RpcError> {
734    let id_str = args
735        .get("id")
736        .and_then(|v| v.as_str())
737        .ok_or_else(|| RpcError::new(INVALID_PARAMS, "forget: missing 'id'"))?;
738    let ulid = crate::Ulid::from_string(id_str)
739        .map_err(|_| RpcError::new(INVALID_PARAMS, format!("forget: invalid ULID '{id_str}'")))?;
740    db.delete(&ulid).map_err(mnemo_err)?;
741    db.flush().map_err(mnemo_err)?;
742    Ok(json!({ "deleted": id_str }).to_string())
743}
744
745fn tool_list(db: &mut Mnemo, args: &Value) -> std::result::Result<String, RpcError> {
746    let include_vectors = args
747        .get("include_vectors")
748        .and_then(|v| v.as_bool())
749        .unwrap_or(false);
750    let mems = db.memories().map_err(mnemo_err)?;
751    let out: Vec<Value> = mems
752        .iter()
753        .map(|m| {
754            let mut v = memory_to_json(m);
755            if !include_vectors {
756                if let Some(o) = v.as_object_mut() {
757                    o.remove("vector");
758                }
759            }
760            v
761        })
762        .collect();
763    Ok(serde_json::to_string_pretty(&out).unwrap_or_else(|_| "[]".into()))
764}
765
766fn tool_snapshot_list(db: &mut Mnemo) -> std::result::Result<String, RpcError> {
767    let snaps = db.snapshots();
768    let out: Vec<Value> = snaps
769        .into_iter()
770        .map(|s| {
771            json!({
772                "txn_id": s.txn_id,
773                "created_at": s.created_at,
774                "memory_count": s.memory_count,
775            })
776        })
777        .collect();
778    Ok(serde_json::to_string_pretty(&out).unwrap_or_else(|_| "[]".into()))
779}
780
781fn tool_stats(db: &mut Mnemo) -> std::result::Result<String, RpcError> {
782    let s = db.stats().map_err(mnemo_err)?;
783    let (cache_pages, cache_capacity) = db.page_cache_stats();
784    let snapshot_count = db.snapshots().len();
785    // Field names mirror `Stats` in `store.rs` verbatim so MCP clients
786    // that reference the Rust doc for the same DB see the same shape.
787    Ok(json!({
788        "memories": s.memories,
789        "deleted": s.deleted,
790        "dimensions": s.dimensions,
791        "file_bytes": s.file_bytes,
792        "agents": s.agents,
793        "encrypted": s.encrypted,
794        "created_at": s.created_at,
795        "wal_pages": s.wal_pages,
796        "index": s.index,
797        "snapshot_count": snapshot_count,
798        "cache_pages": cache_pages,
799        "cache_capacity": cache_capacity,
800    })
801    .to_string())
802}
803
804// --- Result-cache tool handlers (Phase 10.4) -----------------------------
805
806fn tool_cache_put(db: &mut Mnemo, args: &Value) -> std::result::Result<String, RpcError> {
807    let ns = args
808        .get("namespace")
809        .and_then(|v| v.as_str())
810        .ok_or_else(|| RpcError::new(INVALID_PARAMS, "cache_put: missing 'namespace'"))?;
811    let key = args
812        .get("key")
813        .and_then(|v| v.as_str())
814        .ok_or_else(|| RpcError::new(INVALID_PARAMS, "cache_put: missing 'key'"))?;
815    let value = args
816        .get("value")
817        .and_then(|v| v.as_str())
818        .ok_or_else(|| RpcError::new(INVALID_PARAMS, "cache_put: missing 'value'"))?;
819    let content_type = args
820        .get("content_type")
821        .and_then(|v| v.as_str())
822        .unwrap_or("text")
823        .to_string();
824    let ttl_secs = args.get("ttl_secs").and_then(|v| v.as_u64());
825    let opts = CachePutOpts {
826        content_type,
827        ttl_secs,
828        cost_hint_ms: None,
829    };
830    db.cache_put(ns, key, value.as_bytes(), opts)
831        .map_err(mnemo_err)?;
832    db.flush().map_err(mnemo_err)?;
833    Ok(json!({ "cached": true, "namespace": ns, "key": key, "bytes": value.len() }).to_string())
834}
835
836fn tool_cache_get(db: &mut Mnemo, args: &Value) -> std::result::Result<String, RpcError> {
837    let ns = args
838        .get("namespace")
839        .and_then(|v| v.as_str())
840        .ok_or_else(|| RpcError::new(INVALID_PARAMS, "cache_get: missing 'namespace'"))?;
841    let key = args
842        .get("key")
843        .and_then(|v| v.as_str())
844        .ok_or_else(|| RpcError::new(INVALID_PARAMS, "cache_get: missing 'key'"))?;
845    match db.cache_get(ns, key).map_err(mnemo_err)? {
846        Some(hit) => Ok(json!({
847            "hit": true,
848            "value": String::from_utf8_lossy(&hit.value),
849            "content_type": hit.content_type,
850            "created_at": hit.created_at,
851            "accessed_at": hit.accessed_at,
852            "access_count": hit.access_count,
853        })
854        .to_string()),
855        None => Ok(json!({ "hit": false }).to_string()),
856    }
857}
858
859fn tool_cache_stats(db: &mut Mnemo, args: &Value) -> std::result::Result<String, RpcError> {
860    let ns = args.get("namespace").and_then(|v| v.as_str());
861    let s = db.cache_stats(ns);
862    Ok(json!({
863        "namespace": ns,
864        "entries": s.entries,
865        "bytes": s.bytes,
866        "hits": s.hits,
867        "misses": s.misses,
868        "hit_rate": s.hit_rate,
869        "evictions": s.evictions,
870    })
871    .to_string())
872}
873
874fn tool_cache_put_semantic(db: &mut Mnemo, args: &Value) -> std::result::Result<String, RpcError> {
875    let ns = args
876        .get("namespace")
877        .and_then(|v| v.as_str())
878        .ok_or_else(|| RpcError::new(INVALID_PARAMS, "cache_put_semantic: missing 'namespace'"))?;
879    let key = args
880        .get("key")
881        .and_then(|v| v.as_str())
882        .ok_or_else(|| RpcError::new(INVALID_PARAMS, "cache_put_semantic: missing 'key'"))?;
883    let value = args
884        .get("value")
885        .and_then(|v| v.as_str())
886        .ok_or_else(|| RpcError::new(INVALID_PARAMS, "cache_put_semantic: missing 'value'"))?;
887    // When `text` is used, derive the model name from the server's
888    // registered embedder (identity). Explicit `model` wins otherwise.
889    let (vector, model) = match (args.get("vector"), args.get("text")) {
890        (Some(_), Some(_)) => {
891            return Err(RpcError::new(
892                INVALID_PARAMS,
893                "cache_put_semantic: pass EITHER 'vector' or 'text', not both",
894            ));
895        }
896        (Some(v), None) => {
897            let vec = v
898                .as_array()
899                .ok_or_else(|| {
900                    RpcError::new(
901                        INVALID_PARAMS,
902                        "cache_put_semantic: 'vector' must be a JSON array",
903                    )
904                })?
905                .iter()
906                .map(|v| v.as_f64().map(|f| f as f32))
907                .collect::<Option<Vec<f32>>>()
908                .ok_or_else(|| {
909                    RpcError::new(
910                        INVALID_PARAMS,
911                        "cache_put_semantic: 'vector' must be an array of numbers",
912                    )
913                })?;
914            let model = args
915                .get("model")
916                .and_then(|v| v.as_str())
917                .ok_or_else(|| {
918                    RpcError::new(INVALID_PARAMS, "cache_put_semantic: missing 'model'")
919                })?
920                .to_string();
921            (vec, model)
922        }
923        (None, Some(t)) => {
924            let text = t.as_str().ok_or_else(|| {
925                RpcError::new(
926                    INVALID_PARAMS,
927                    "cache_put_semantic: 'text' must be a string",
928                )
929            })?;
930            let vec = embed_text_via_server(db, "cache_put_semantic", text)?;
931            // Allow explicit model override; else derive from embedder.
932            let model = match args.get("model").and_then(|v| v.as_str()) {
933                Some(m) => m.to_string(),
934                None => {
935                    db.embedder()
936                        .expect("embed_text_via_server would have errored otherwise")
937                        .identity()
938                        .name
939                }
940            };
941            (vec, model)
942        }
943        (None, None) => {
944            return Err(RpcError::new(
945                INVALID_PARAMS,
946                "cache_put_semantic: missing 'vector' or 'text'",
947            ))
948        }
949    };
950    let content_type = args
951        .get("content_type")
952        .and_then(|v| v.as_str())
953        .unwrap_or("text")
954        .to_string();
955    let ttl_secs = args.get("ttl_secs").and_then(|v| v.as_u64());
956    let opts = SemanticCachePutOpts {
957        content_type,
958        ttl_secs,
959        cost_hint_ms: None,
960        model,
961    };
962    db.cache_put_semantic(ns, key, vector, value.as_bytes(), opts)
963        .map_err(mnemo_err)?;
964    db.flush().map_err(mnemo_err)?;
965    Ok(json!({ "cached": true, "namespace": ns, "key": key, "bytes": value.len() }).to_string())
966}
967
968fn tool_cache_get_semantic(db: &mut Mnemo, args: &Value) -> std::result::Result<String, RpcError> {
969    let ns = args
970        .get("namespace")
971        .and_then(|v| v.as_str())
972        .ok_or_else(|| RpcError::new(INVALID_PARAMS, "cache_get_semantic: missing 'namespace'"))?;
973    // When `text` is used, derive the model name from the server's
974    // registered embedder (identity). Explicit `model` wins otherwise.
975    let (query, model_owned) = match (args.get("query"), args.get("text")) {
976        (Some(_), Some(_)) => {
977            return Err(RpcError::new(
978                INVALID_PARAMS,
979                "cache_get_semantic: pass EITHER 'query' or 'text', not both",
980            ));
981        }
982        (Some(q), None) => {
983            let vec = q
984                .as_array()
985                .ok_or_else(|| {
986                    RpcError::new(
987                        INVALID_PARAMS,
988                        "cache_get_semantic: 'query' must be a JSON array",
989                    )
990                })?
991                .iter()
992                .map(|v| v.as_f64().map(|f| f as f32))
993                .collect::<Option<Vec<f32>>>()
994                .ok_or_else(|| {
995                    RpcError::new(
996                        INVALID_PARAMS,
997                        "cache_get_semantic: 'query' must be an array of numbers",
998                    )
999                })?;
1000            let model = args
1001                .get("model")
1002                .and_then(|v| v.as_str())
1003                .ok_or_else(|| {
1004                    RpcError::new(INVALID_PARAMS, "cache_get_semantic: missing 'model'")
1005                })?
1006                .to_string();
1007            (vec, model)
1008        }
1009        (None, Some(t)) => {
1010            let text = t.as_str().ok_or_else(|| {
1011                RpcError::new(
1012                    INVALID_PARAMS,
1013                    "cache_get_semantic: 'text' must be a string",
1014                )
1015            })?;
1016            let vec = embed_text_via_server(db, "cache_get_semantic", text)?;
1017            let model = match args.get("model").and_then(|v| v.as_str()) {
1018                Some(m) => m.to_string(),
1019                None => {
1020                    db.embedder()
1021                        .expect("embed_text_via_server would have errored otherwise")
1022                        .identity()
1023                        .name
1024                }
1025            };
1026            (vec, model)
1027        }
1028        (None, None) => {
1029            return Err(RpcError::new(
1030                INVALID_PARAMS,
1031                "cache_get_semantic: missing 'query' or 'text'",
1032            ))
1033        }
1034    };
1035    let model = model_owned.as_str();
1036    let threshold = args
1037        .get("threshold")
1038        .and_then(|v| v.as_f64())
1039        .map(|f| f as f32)
1040        .unwrap_or(DEFAULT_SEMANTIC_THRESHOLD);
1041    match db
1042        .cache_get_semantic(ns, &query, threshold, model)
1043        .map_err(mnemo_err)?
1044    {
1045        Some((hit, similarity)) => Ok(json!({
1046            "hit": true,
1047            "similarity": similarity,
1048            "value": String::from_utf8_lossy(&hit.value),
1049            "content_type": hit.content_type,
1050            "created_at": hit.created_at,
1051            "accessed_at": hit.accessed_at,
1052            "access_count": hit.access_count,
1053        })
1054        .to_string()),
1055        None => Ok(json!({ "hit": false }).to_string()),
1056    }
1057}
1058
1059/// Serialize a `Memory` into the JSON shape MCP tool responses use.
1060/// Serde on `Memory` produces the canonical shape; we mirror what the
1061/// CLI's `memory_to_json` does but without the CLI-only vector-stripping
1062/// switch (that lives in the individual tools that care).
1063fn memory_to_json(m: &Memory) -> Value {
1064    serde_json::to_value(m).unwrap_or_else(|_| {
1065        let mut map = Map::new();
1066        map.insert("id".into(), json!(m.id.to_string()));
1067        Value::Object(map)
1068    })
1069}