Skip to main content

server/
mcp.rs

1//! MCP server: JSON-RPC 2.0 over newline-delimited stdio.
2//!
3//! Framing is **one JSON object per line**. LSP-style `Content-Length`
4//! headers are not accepted — a header line is a parse error (`-32700`)
5//! and the loop continues. Blank lines are skipped.
6//!
7//! # Methods
8//!
9//! - `initialize` — `protocolVersion` `"2024-11-05"`, `capabilities.tools`,
10//!   `serverInfo.name` `"mushroomdb"`, `serverInfo.version` (crate version)
11//! - `notifications/initialized` — ignored
12//! - `tools/list` — the default listing follows the store the server opened
13//!   (see [`Surface`]): a store a repository was ingested into lists three —
14//!   `explore`, `query`, `stats` — and any other store lists the fifteen of
15//!   [`ASSOCIATION_TOOLS`], the tools that answer a question about an entity
16//!   graph, in that order. Graph-tool descriptions carry the
17//!   prefix `Advanced: ` so a host ranking tools by description puts the task
18//!   tools in front. `mushroomdb mcp --all-tools` lists all twenty-seven; the
19//!   rest are callable either way, just not advertised
20//! - `tools/call` — dispatch; success for a graph tool is
21//!   `{content:[{type:"text", text:<json string>}]}`, and for a task tool one
22//!   text block holding the rendered digest — or, with `json: true`, the
23//!   serialised report. No task tool returns `structuredContent`
24//!
25//! Unknown methods on a **request** (has `id`) → `-32601`. A notification
26//! (no `id` member) never writes a response, including unknown methods.
27//!
28//! # Error split
29//!
30//! Protocol errors are JSON-RPC `error` objects:
31//! - `-32700` parse — unparseable line (invalid JSON / invalid UTF-8)
32//! - `-32600` invalid request — parsed JSON that is not an object, or a
33//!   request with missing / non-string `method`
34//! - `-32601` method — unknown `method` on a request
35//! - `-32602` params — `tools/call` envelope invalid: `params` not an object,
36//!   missing / non-string `name`, `arguments` present but not an object,
37//!   or unknown tool name
38//!
39//! Tool-level failures are JSON-RPC **results** with `isError: true` and a
40//! text message: missing or wrong-typed fields inside a known tool's
41//! `arguments`, and every [`GraphError`] from core-api.
42//!
43//! # Deadlock
44//!
45//! [`SharedDb::read`] / [`SharedDb::write`] guards are held only for the
46//! public core-api call, then dropped before serializing or writing. Do not
47//! nest a second lock on the same handle (the `RwLock` is not re-entrant).
48//!
49//! EOF on `reader` returns `Ok(())`. Read/write I/O errors propagate.
50
51use crate::json::{
52    edge_history_result_json, node_history_json, node_info_json, params_from_json,
53    parse_ingest_edges, result_set_json, rule_def_from_json,
54};
55use core_api::{
56    json_to_rows, json_to_value, AutoFk, GraphError, IngestOptions, MaskMode, NodeMask, SharedDb,
57    Value,
58};
59use serde_json::{json, Value as Js};
60use std::collections::BTreeMap;
61use std::io::{self, BufRead, Write};
62use std::path::{Path, PathBuf};
63
64/// Run the MCP loop until `reader` hits EOF.
65///
66/// `db_dir` is where the store lives on disk. `mushroomdb mcp <db>` passes it;
67/// a caller that has only a handle passes `None`, and the one tool that needs a
68/// path — `sync`, which re-runs this binary against the store — reports that it
69/// cannot run rather than guessing one.
70pub fn run_mcp_stdio(
71    db: SharedDb,
72    db_dir: Option<PathBuf>,
73    reader: impl BufRead,
74    writer: impl Write,
75) -> io::Result<()> {
76    run_mcp_stdio_with(db, db_dir, false, reader, writer)
77}
78
79/// [`run_mcp_stdio`], with the tool list chosen by the caller.
80///
81/// `all_tools` false lists what the store's [`Surface`] names — three on a
82/// code graph, fifteen on a memory store; true lists all twenty-seven. Either
83/// way every tool remains callable — the flag decides what is advertised, not
84/// what is served.
85///
86/// The surface is read once, here, rather than per `tools/list`: a store does
87/// not become a code graph half way through a session, and a listing that
88/// changed under a host that caches it would be worse than one that is merely
89/// stale.
90pub fn run_mcp_stdio_with(
91    db: SharedDb,
92    db_dir: Option<PathBuf>,
93    all_tools: bool,
94    mut reader: impl BufRead,
95    mut writer: impl Write,
96) -> io::Result<()> {
97    let surface = surface_of(&db);
98    let mut buf = Vec::new();
99    loop {
100        buf.clear();
101        let n = reader.read_until(b'\n', &mut buf)?;
102        if n == 0 {
103            return Ok(());
104        }
105        match std::str::from_utf8(&buf) {
106            Ok(s) if s.trim().is_empty() => continue,
107            Ok(s) => handle_line(
108                &db,
109                db_dir.as_deref(),
110                all_tools,
111                surface,
112                s.trim(),
113                &mut writer,
114            )?,
115            Err(_) => write_error(&mut writer, None, -32700, "Parse error")?,
116        }
117    }
118}
119
120fn handle_line(
121    db: &SharedDb,
122    db_dir: Option<&Path>,
123    all_tools: bool,
124    surface: Surface,
125    line: &str,
126    writer: &mut impl Write,
127) -> io::Result<()> {
128    let msg: Js = match serde_json::from_str(line) {
129        Ok(v) => v,
130        Err(_) => return write_error(writer, None, -32700, "Parse error"),
131    };
132    let Some(obj) = msg.as_object() else {
133        return write_error(writer, None, -32600, "Invalid Request");
134    };
135    let is_request = obj.contains_key("id");
136    let id = obj.get("id").cloned();
137    let method = match obj.get("method").and_then(Js::as_str) {
138        Some(m) => m,
139        None => {
140            if is_request {
141                write_error(writer, id, -32600, "Invalid Request")?;
142            }
143            return Ok(());
144        }
145    };
146    match method {
147        "initialize" => {
148            if is_request {
149                write_result(writer, id, initialize_result())?;
150            }
151        }
152        "notifications/initialized" => {
153            if is_request {
154                write_result(writer, id, json!({}))?;
155            }
156        }
157        "tools/list" => {
158            if is_request {
159                write_result(writer, id, tools_list(all_tools, surface))?;
160            }
161        }
162        "tools/call" => {
163            if is_request {
164                match dispatch_call(db, db_dir, obj.get("params")) {
165                    CallOutcome::Protocol { code, message } => {
166                        write_error(writer, id, code, &message)?;
167                    }
168                    CallOutcome::ToolOk(payload) => {
169                        write_result(writer, id, tool_ok(payload))?;
170                    }
171                    CallOutcome::TaskOk { text } => {
172                        write_result(writer, id, task_ok(&text))?;
173                    }
174                    CallOutcome::ToolErr(message) => {
175                        write_result(writer, id, tool_err(&message))?;
176                    }
177                }
178            }
179        }
180        _ => {
181            if is_request {
182                write_error(writer, id, -32601, "Method not found")?;
183            }
184        }
185    }
186    Ok(())
187}
188
189pub(crate) enum CallOutcome {
190    Protocol {
191        code: i64,
192        message: String,
193    },
194    /// A graph tool's JSON payload, returned as a JSON string in `content`.
195    ToolOk(Js),
196    /// A task tool's answer, as text and nothing else: the rendered digest, or
197    /// the serialised report when the call passed `json: true`.
198    TaskOk {
199        text: String,
200    },
201    ToolErr(String),
202}
203
204fn dispatch_call(db: &SharedDb, db_dir: Option<&Path>, params: Option<&Js>) -> CallOutcome {
205    let Some(params) = params.and_then(Js::as_object) else {
206        return protocol_invalid();
207    };
208    let Some(name) = params.get("name").and_then(Js::as_str) else {
209        return protocol_invalid();
210    };
211    let empty = json!({});
212    let args = match params.get("arguments") {
213        None => &empty,
214        Some(a) if a.is_object() => a,
215        Some(_) => return protocol_invalid(),
216    };
217    // The repository task tools first, in the order `tools/list` advertises.
218    if let Some(outcome) = crate::mcp_tasks::dispatch(db, db_dir, name, args) {
219        return outcome;
220    }
221    match name {
222        "query" => tool_query(db, args),
223        "ingest_json" => tool_ingest(db, args),
224        "create_rule" => tool_create_rule(db, args),
225        "explain" => tool_explain(db, args),
226        "stats" => tool_stats(db),
227        "node_info" => tool_node_info(db, args),
228        "upsert_entity" => tool_upsert_entity(db, args),
229        "find_similar" => tool_find_similar(db, args),
230        "hybrid_search" => tool_hybrid_search(db, args),
231        "node_history" => tool_node_history(db, args),
232        "edge_history" => tool_edge_history(db, args),
233        "was_linked" => tool_was_linked(db, args),
234        "rename_node" => tool_rename_node(db, args),
235        _ => protocol_invalid(),
236    }
237}
238
239fn protocol_invalid() -> CallOutcome {
240    CallOutcome::Protocol {
241        code: -32602,
242        message: "Invalid params".into(),
243    }
244}
245
246fn tool_query(db: &SharedDb, args: &Js) -> CallOutcome {
247    let Some(cypher) = args.get("cypher").and_then(Js::as_str) else {
248        return CallOutcome::ToolErr("missing cypher".into());
249    };
250    let params = match params_from_json(args.get("params")) {
251        Ok(p) => p,
252        Err(e) => return CallOutcome::ToolErr(e),
253    };
254
255    // Two ways to ask the same restricted question: a `role` names one the
256    // store already defines, a `mask` writes the allow-list out by hand. Both
257    // route to `query_masked` (read-only). Passing both is not a merge of the
258    // two — it is a caller that has not decided which restriction applies, so
259    // it is refused rather than silently resolved one way.
260    let role = match args.get("role") {
261        None | Some(Js::Null) => None,
262        Some(Js::String(s)) if !s.is_empty() => Some(s.as_str()),
263        Some(_) => return CallOutcome::ToolErr("role must be a non-empty string".into()),
264    };
265    let mask_keys = match args.get("mask") {
266        None => None,
267        Some(v) => match mask_key_list(v) {
268            Ok(keys) => Some(keys),
269            Err(e) => return CallOutcome::ToolErr(e),
270        },
271    };
272    if role.is_some() && mask_keys.is_some() {
273        return CallOutcome::ToolErr("pass role or mask, not both".into());
274    }
275
276    if role.is_some() || mask_keys.is_some() {
277        let stub_hidden = args
278            .get("stub_hidden")
279            .and_then(|v| v.as_bool())
280            .unwrap_or(false);
281        let g = db.read();
282        let mask = match (role, &mask_keys) {
283            (Some(role), _) => match g.mask_for_role(role) {
284                Ok(m) => m,
285                // The one error a caller can fix by rereading `roles.json`,
286                // told apart from a store whose roles never loaded at all.
287                Err(GraphError::KeyNotFound { .. }) => {
288                    return CallOutcome::ToolErr(format!("unknown role '{role}'"))
289                }
290                Err(e) => return CallOutcome::ToolErr(graph_err_msg(e)),
291            },
292            (None, Some(keys)) => NodeMask::from_keys(&*g, keys.iter().map(String::as_str)),
293            (None, None) => unreachable!("one of the two is Some in this branch"),
294        };
295        let mask = if stub_hidden {
296            mask.with_mode(MaskMode::Stub)
297        } else {
298            mask
299        };
300        return match g.query_masked(cypher, &params, &mask) {
301            Ok(rs) => CallOutcome::ToolOk(result_set_json(&rs)),
302            Err(e) => CallOutcome::ToolErr(graph_err_msg(e)),
303        };
304    }
305
306    let is_write = match core_api::is_write_query(cypher) {
307        Ok(b) => b,
308        Err(e) => return CallOutcome::ToolErr(e),
309    };
310    let rs = if is_write {
311        let mut g = db.write();
312        g.query_write(cypher, &params)
313    } else {
314        let g = db.read();
315        g.query(cypher, &params)
316    };
317    match rs {
318        Ok(rs) => CallOutcome::ToolOk(result_set_json(&rs)),
319        Err(e) => CallOutcome::ToolErr(graph_err_msg(e)),
320    }
321}
322
323/// A `mask` argument as a key list. `Err` when it is anything but an array of
324/// strings — including `null`, which is a caller that meant to pass one.
325fn mask_key_list(mask: &Js) -> Result<Vec<String>, String> {
326    let arr = mask
327        .as_array()
328        .ok_or_else(|| "mask must be an array of strings".to_string())?;
329    arr.iter()
330        .map(|v| {
331            v.as_str()
332                .map(str::to_string)
333                .ok_or_else(|| "mask must be an array of strings".to_string())
334        })
335        .collect()
336}
337
338fn tool_ingest(db: &SharedDb, args: &Js) -> CallOutcome {
339    let Some(label) = args.get("label").and_then(Js::as_str) else {
340        return CallOutcome::ToolErr("missing label".into());
341    };
342    let Some(rows_json) = args.get("rows_json").and_then(Js::as_str) else {
343        return CallOutcome::ToolErr("missing rows_json".into());
344    };
345    let mut opts = IngestOptions::default();
346    if let Some(kf) = args.get("key_field") {
347        match kf.as_str() {
348            Some(s) => opts.key_field = s.to_string(),
349            None => return CallOutcome::ToolErr("key_field must be a string".into()),
350        }
351    }
352    if let Some(suf) = args.get("auto_fk_suffix") {
353        match suf.as_str() {
354            Some(s) => {
355                opts.auto_fk = AutoFk::Auto {
356                    suffix: s.to_string(),
357                }
358            }
359            None => return CallOutcome::ToolErr("auto_fk_suffix must be a string".into()),
360        }
361    }
362    let edges = match args.get("edges") {
363        None | Some(Js::Null) => Vec::new(),
364        Some(raw) => match parse_ingest_edges(raw) {
365            Ok(e) => e,
366            Err(e) => return CallOutcome::ToolErr(e),
367        },
368    };
369    let parsed: Js = match serde_json::from_str(rows_json) {
370        Ok(v) => v,
371        Err(e) => {
372            return CallOutcome::ToolErr(graph_err_msg(GraphError::IngestError {
373                detail: e.to_string(),
374            }))
375        }
376    };
377    let mut converted = match json_to_rows(&parsed) {
378        Ok(c) => c,
379        Err(e) => return CallOutcome::ToolErr(graph_err_msg(e)),
380    };
381    let taken = std::mem::take(&mut converted.rows);
382    let report = {
383        let mut g = db.write();
384        g.ingest_with_edges(label, taken, &opts, &edges)
385    };
386    match report.map(|r| converted.into_report(r)) {
387        Ok(r) => match serde_json::to_value(&r) {
388            Ok(v) => CallOutcome::ToolOk(v),
389            Err(e) => CallOutcome::ToolErr(e.to_string()),
390        },
391        Err(e) => CallOutcome::ToolErr(graph_err_msg(e)),
392    }
393}
394
395fn tool_create_rule(db: &SharedDb, args: &Js) -> CallOutcome {
396    let def = match rule_def_from_json(args.clone()) {
397        Ok(d) => d,
398        Err(e) => return CallOutcome::ToolErr(e),
399    };
400    let name = def.name.clone();
401    let res = {
402        let mut g = db.write();
403        g.create_rule(def)
404    };
405    match res {
406        Ok(()) => CallOutcome::ToolOk(json!({"ok": true, "name": name})),
407        Err(e) => CallOutcome::ToolErr(graph_err_msg(e)),
408    }
409}
410
411fn tool_explain(db: &SharedDb, args: &Js) -> CallOutcome {
412    let Some(a) = args.get("a").and_then(Js::as_str).filter(|s| !s.is_empty()) else {
413        return CallOutcome::ToolErr("missing a".into());
414    };
415    let Some(b) = args.get("b").and_then(Js::as_str).filter(|s| !s.is_empty()) else {
416        return CallOutcome::ToolErr("missing b".into());
417    };
418    let out = {
419        let g = db.read();
420        g.explain(a, b)
421    };
422    match out {
423        Ok(v) => match serde_json::to_value(&v) {
424            Ok(j) => CallOutcome::ToolOk(j),
425            Err(e) => CallOutcome::ToolErr(e.to_string()),
426        },
427        Err(e) => CallOutcome::ToolErr(graph_err_msg(e)),
428    }
429}
430
431fn tool_stats(db: &SharedDb) -> CallOutcome {
432    let snap = {
433        let g = db.read();
434        g.stats()
435    };
436    match serde_json::to_value(&snap) {
437        Ok(v) => CallOutcome::ToolOk(v),
438        Err(e) => CallOutcome::ToolErr(e.to_string()),
439    }
440}
441
442fn tool_node_info(db: &SharedDb, args: &Js) -> CallOutcome {
443    let Some(key) = args.get("key").and_then(Js::as_str) else {
444        return CallOutcome::ToolErr("missing key".into());
445    };
446    let info = {
447        let g = db.read();
448        g.node_info(key)
449    };
450    match info {
451        Some(info) => CallOutcome::ToolOk(node_info_json(&info)),
452        None => CallOutcome::ToolErr(graph_err_msg(GraphError::KeyNotFound {
453            key: key.to_string(),
454        })),
455    }
456}
457
458/// Insert a new node or update an existing node's properties, keyed by `key`.
459///
460/// If the node exists: each prop in `props` is written via `set_prop`.
461/// If the node does not exist: `label` is required; the node is ingested with
462/// `key_field = "id"` and the supplied props.
463///
464/// Returns `{ok, key, created, updated_fields?}`.
465fn tool_upsert_entity(db: &SharedDb, args: &Js) -> CallOutcome {
466    let Some(key) = args.get("key").and_then(Js::as_str) else {
467        return CallOutcome::ToolErr("missing key".into());
468    };
469    let label_opt = args.get("label").and_then(Js::as_str);
470    let Some(props_obj) = args.get("props").and_then(Js::as_object) else {
471        return CallOutcome::ToolErr("missing props".into());
472    };
473
474    let exists = {
475        let g = db.read();
476        g.has_node(key)
477    };
478
479    if exists {
480        let mut g = db.write();
481        let mut count = 0usize;
482        for (field, json_val) in props_obj {
483            match json_to_value(json_val.clone()) {
484                Some(v) => {
485                    if let Err(e) = g.set_prop(key, field, v) {
486                        return CallOutcome::ToolErr(graph_err_msg(e));
487                    }
488                    count += 1;
489                }
490                None => {
491                    return CallOutcome::ToolErr(format!(
492                        "prop {field} is not a supported value type"
493                    ))
494                }
495            }
496        }
497        CallOutcome::ToolOk(json!({
498            "ok": true,
499            "key": key,
500            "created": false,
501            "updated_fields": count
502        }))
503    } else {
504        let Some(label) = label_opt else {
505            return CallOutcome::ToolErr("label required when creating a new entity".into());
506        };
507        let mut row: BTreeMap<String, Value> = BTreeMap::new();
508        row.insert("id".to_string(), Value::Str(key.to_string()));
509        for (field, json_val) in props_obj {
510            if field == "id" {
511                continue;
512            }
513            match json_to_value(json_val.clone()) {
514                Some(v) => {
515                    row.insert(field.clone(), v);
516                }
517                None => {
518                    return CallOutcome::ToolErr(format!(
519                        "prop {field} is not a supported value type"
520                    ))
521                }
522            }
523        }
524        let opts = IngestOptions {
525            key_field: "id".to_string(),
526            auto_fk: AutoFk::Off,
527        };
528        let mut g = db.write();
529        match g.ingest(label, vec![row], &opts) {
530            Ok(_) => CallOutcome::ToolOk(json!({ "ok": true, "key": key, "created": true })),
531            Err(e) => CallOutcome::ToolErr(graph_err_msg(e)),
532        }
533    }
534}
535
536/// Return neighbors connected by a given edge type (default `"SIMILAR"`).
537///
538/// Results are read from edges already materialized by a derivation rule
539/// (e.g. a `VectorSimilar` rule). Without a matching rule the returned list
540/// is empty — no live cosine computation is performed here.
541/// Returns up to `limit` (default 10) neighbor entries.
542fn tool_find_similar(db: &SharedDb, args: &Js) -> CallOutcome {
543    // Parse the optional mask once — it applies to both vector and edge paths.
544    // An invalid mask value (non-array or non-string element) fails closed.
545    let mask_keys: Option<Vec<String>> = if let Some(mask_val) = args.get("mask") {
546        match mask_val.as_array() {
547            Some(arr) => {
548                let mut ks: Vec<String> = Vec::with_capacity(arr.len());
549                for v in arr {
550                    match v.as_str() {
551                        Some(s) => ks.push(s.to_string()),
552                        None => {
553                            return CallOutcome::ToolErr("mask must be an array of strings".into())
554                        }
555                    }
556                }
557                Some(ks)
558            }
559            None => return CallOutcome::ToolErr("mask must be an array of strings".into()),
560        }
561    } else {
562        None
563    };
564
565    // When a `vector` array is provided, use the HNSW / brute-force vector
566    // similarity path instead of looking up pre-derived edges.
567    if let Some(vec_js) = args.get("vector").and_then(Js::as_array) {
568        let q: Vec<f64> = vec_js.iter().filter_map(|v| v.as_f64()).collect();
569        if q.is_empty() {
570            return CallOutcome::ToolErr("vector must be a non-empty array of numbers".into());
571        }
572        let field = args
573            .get("field")
574            .and_then(Js::as_str)
575            .unwrap_or("embedding");
576        let label_str = args.get("label").and_then(Js::as_str).unwrap_or("");
577        let label = if label_str.is_empty() {
578            None
579        } else {
580            Some(label_str)
581        };
582        let k = args
583            .get("k")
584            .and_then(Js::as_u64)
585            .map(|n| n as usize)
586            .unwrap_or(10);
587        let min = args.get("min").and_then(Js::as_f64).unwrap_or(0.8);
588
589        let hits = {
590            let g = db.read();
591            if let Some(ref keys) = mask_keys {
592                let node_mask = NodeMask::from_keys(&*g, keys.iter().map(String::as_str));
593                g.find_similar_vector_masked(field, label, &q, k, min, &node_mask)
594            } else {
595                g.find_similar_vector(field, label, &q, k, min)
596            }
597        };
598        let results: Vec<Js> = hits
599            .into_iter()
600            .map(|(key, score)| json!({ "key": key, "score": score }))
601            .collect();
602        return CallOutcome::ToolOk(json!({
603            "mode": "vector",
604            "field": field,
605            "label": label,
606            "k": k,
607            "min": min,
608            "results": results
609        }));
610    }
611
612    // Edge-traversal path: return neighbors connected by the given edge type.
613    let Some(key) = args.get("key").and_then(Js::as_str) else {
614        return CallOutcome::ToolErr("missing key (or provide vector for vector search)".into());
615    };
616    let edge_type = args
617        .get("edge_type")
618        .and_then(Js::as_str)
619        .unwrap_or("SIMILAR");
620    let limit = args
621        .get("limit")
622        .and_then(Js::as_u64)
623        .map(|n| n as usize)
624        .unwrap_or(10);
625
626    // When a mask is present, a hidden query key behaves identically to a
627    // nonexistent key — we do not confirm its existence.
628    if let Some(ref mask) = mask_keys {
629        let mask_set: std::collections::HashSet<&str> = mask.iter().map(String::as_str).collect();
630        if !mask_set.contains(key) {
631            return CallOutcome::ToolErr(graph_err_msg(GraphError::KeyNotFound {
632                key: key.into(),
633            }));
634        }
635        let out = {
636            let g = db.read();
637            g.node_edges(key)
638        };
639        return match out {
640            Ok(edges) => {
641                let similar: Vec<Js> = edges
642                    .iter()
643                    .filter(|e| e.edge_type == edge_type)
644                    .filter(|e| {
645                        // Keep only edges where the neighbor is also visible.
646                        let neighbor_key = if e.src_key == key {
647                            &e.dst_key
648                        } else {
649                            &e.src_key
650                        };
651                        mask_set.contains(neighbor_key.as_str())
652                    })
653                    .take(limit)
654                    .map(|e| {
655                        let neighbor_key = if e.src_key == key {
656                            &e.dst_key
657                        } else {
658                            &e.src_key
659                        };
660                        let direction = if e.src_key == key { "out" } else { "in" };
661                        json!({
662                            "neighbor_key": neighbor_key,
663                            "direction": direction,
664                            "edge_type": e.edge_type,
665                            "derived": e.derived,
666                        })
667                    })
668                    .collect();
669                CallOutcome::ToolOk(json!({
670                    "key": key,
671                    "edge_type": edge_type,
672                    "similar": similar
673                }))
674            }
675            Err(e) => CallOutcome::ToolErr(graph_err_msg(e)),
676        };
677    }
678
679    let out = {
680        let g = db.read();
681        g.node_edges(key)
682    };
683    match out {
684        Ok(edges) => {
685            let similar: Vec<Js> = edges
686                .iter()
687                .filter(|e| e.edge_type == edge_type)
688                .take(limit)
689                .map(|e| {
690                    let neighbor_key = if e.src_key == key {
691                        &e.dst_key
692                    } else {
693                        &e.src_key
694                    };
695                    let direction = if e.src_key == key { "out" } else { "in" };
696                    json!({
697                        "neighbor_key": neighbor_key,
698                        "direction": direction,
699                        "edge_type": e.edge_type,
700                        "derived": e.derived,
701                    })
702                })
703                .collect();
704            CallOutcome::ToolOk(json!({
705                "key": key,
706                "edge_type": edge_type,
707                "similar": similar
708            }))
709        }
710        Err(e) => CallOutcome::ToolErr(graph_err_msg(e)),
711    }
712}
713
714fn tool_hybrid_search(db: &SharedDb, args: &Js) -> CallOutcome {
715    let Some(query_text) = args.get("query_text").and_then(Js::as_str) else {
716        return CallOutcome::ToolErr("missing required field: query_text".into());
717    };
718    let Some(text_field) = args.get("text_field").and_then(Js::as_str) else {
719        return CallOutcome::ToolErr("missing required field: text_field".into());
720    };
721
722    let vector_field = args
723        .get("vector_field")
724        .and_then(Js::as_str)
725        .unwrap_or("embedding");
726    let label = args.get("label").and_then(Js::as_str);
727    let k = args
728        .get("k")
729        .and_then(Js::as_u64)
730        .map(|n| n as usize)
731        .unwrap_or(10);
732
733    let query_vec: Vec<f64> = args
734        .get("vector")
735        .and_then(Js::as_array)
736        .map(|arr| arr.iter().filter_map(|v| v.as_f64()).collect())
737        .unwrap_or_default();
738
739    let hits = {
740        let g = db.read();
741        g.search_hybrid(text_field, query_text, vector_field, &query_vec, label, k)
742    };
743
744    let results: Vec<Js> = hits
745        .into_iter()
746        .map(|(key, score)| json!({ "key": key, "score": score }))
747        .collect();
748
749    CallOutcome::ToolOk(json!({
750        "query_text": query_text,
751        "text_field": text_field,
752        "vector_field": vector_field,
753        "label": label,
754        "k": k,
755        "results": results
756    }))
757}
758
759fn tool_node_history(db: &SharedDb, args: &Js) -> CallOutcome {
760    let Some(key) = args.get("key").and_then(Js::as_str) else {
761        return CallOutcome::ToolErr("missing key".into());
762    };
763    let g = db.read();
764    let entries = match g.node_history(key) {
765        Ok(e) => e,
766        Err(e) => return CallOutcome::ToolErr(graph_err_msg(e)),
767    };
768    let total_commits = match g.wal_total_commits() {
769        Ok(n) => n,
770        Err(e) => return CallOutcome::ToolErr(graph_err_msg(e)),
771    };
772    CallOutcome::ToolOk(node_history_json(key, &entries, total_commits))
773}
774
775fn tool_edge_history(db: &SharedDb, args: &Js) -> CallOutcome {
776    let Some(a) = args.get("a").and_then(Js::as_str).filter(|s| !s.is_empty()) else {
777        return CallOutcome::ToolErr("missing a".into());
778    };
779    let Some(b) = args.get("b").and_then(Js::as_str).filter(|s| !s.is_empty()) else {
780        return CallOutcome::ToolErr("missing b".into());
781    };
782    let result = {
783        let g = db.read();
784        g.edge_history(a, b)
785    };
786    match result {
787        Ok(hr) => CallOutcome::ToolOk(edge_history_result_json(a, b, &hr)),
788        Err(e) => CallOutcome::ToolErr(graph_err_msg(e)),
789    }
790}
791
792fn tool_was_linked(db: &SharedDb, args: &Js) -> CallOutcome {
793    let Some(a) = args.get("a").and_then(Js::as_str).filter(|s| !s.is_empty()) else {
794        return CallOutcome::ToolErr("missing a".into());
795    };
796    let Some(b) = args.get("b").and_then(Js::as_str).filter(|s| !s.is_empty()) else {
797        return CallOutcome::ToolErr("missing b".into());
798    };
799    let Some(edge_type) = args
800        .get("edge_type")
801        .and_then(Js::as_str)
802        .filter(|s| !s.is_empty())
803    else {
804        return CallOutcome::ToolErr("missing edge_type".into());
805    };
806    let at_commit = match args.get("at_commit").and_then(Js::as_u64) {
807        Some(n) => n,
808        None => return CallOutcome::ToolErr("missing or invalid at_commit".into()),
809    };
810    let result = {
811        let g = db.read();
812        g.was_linked(a, b, edge_type, at_commit)
813    };
814    match result {
815        Ok(linked) => CallOutcome::ToolOk(json!({
816            "a": a,
817            "b": b,
818            "edge_type": edge_type,
819            "at_commit": at_commit,
820            "linked": linked,
821        })),
822        Err(e) => CallOutcome::ToolErr(graph_err_msg(e)),
823    }
824}
825
826fn tool_rename_node(db: &SharedDb, args: &Js) -> CallOutcome {
827    let Some(old_key) = args.get("old_key").and_then(Js::as_str) else {
828        return CallOutcome::ToolErr("missing old_key".into());
829    };
830    let Some(new_key) = args.get("new_key").and_then(Js::as_str) else {
831        return CallOutcome::ToolErr("missing new_key".into());
832    };
833    let mut g = db.write();
834    match g.rename_node(old_key, new_key) {
835        Ok(()) => CallOutcome::ToolOk(json!({
836            "ok": true,
837            "old_key": old_key,
838            "new_key": new_key,
839        })),
840        Err(e) => CallOutcome::ToolErr(graph_err_msg(e)),
841    }
842}
843
844pub(crate) fn graph_err_msg(e: GraphError) -> String {
845    match e {
846        GraphError::QueryError { detail } | GraphError::IngestError { detail } => detail,
847        other => other.to_string(),
848    }
849}
850
851fn initialize_result() -> Js {
852    json!({
853        "protocolVersion": "2024-11-05",
854        "capabilities": { "tools": {} },
855        "serverInfo": { "name": "mushroomdb", "version": env!("CARGO_PKG_VERSION") }
856    })
857}
858
859/// The prefix every graph tool's description carries.
860///
861/// A host that ranks tools by their description now has one signal that the
862/// repository task tools are the ones to reach for first, and that everything
863/// under this prefix is the lower-level surface beneath them.
864const ADVANCED_PREFIX: &str = "Advanced: ";
865
866/// The three a code-graph store advertises: one tool to find, one to ask an
867/// arbitrary question, one to size the store.
868///
869/// `ingest_json` is not among them. A store built by `ingest-git` is written by
870/// `sync` and `touch`, not by an assistant bulk-loading rows into it, and the
871/// tool that is never the right one on this surface is the one worth not
872/// listing.
873pub const CODE_GRAPH_TOOLS: [&str; 3] = ["explore", "query", "stats"];
874
875/// The fifteen a memory store advertises, in the order it lists them.
876///
877/// A store with no repository in it used to be handed the code door's own task
878/// tools — `map`, `context`, `impact`, `owners`, `why`, `sync` — which answer
879/// from a code graph there is none of, plus `ingest_json`. Six of the eleven
880/// names an assistant found answered from a repository the store did not
881/// hold. These are
882/// the questions an entity graph *can* answer: what is there (`query` — now
883/// with a `role`), why two things are associated, what is around a node, what
884/// it is and what it is joined to, whether a link held at a commit and when it
885/// changed, what is like it, and what was written down about it.
886///
887/// Listing order is ranking: a host that defers schemas shows this list in
888/// order, so the two questions this door exists for come first.
889///
890/// `edges_at` and `what_if` are the two the first association benchmark run
891/// showed missing: a run asked what a node's relationships were at a past
892/// commit and spent twenty to sixty-seven turns replaying `edge_history` for
893/// it, and had no way at all to ask what a change would do. They sit after
894/// `was_linked`, which is the narrowest form of the same time question.
895///
896/// The code task tools stay served on a memory store, as these stay served on
897/// a code-graph one — [`tools_list`] decides what is *advertised*, never what
898/// is answered.
899pub const ASSOCIATION_TOOLS: [&str; 15] = [
900    "query",
901    "explain_association",
902    "neighborhood",
903    "node_info",
904    "node_edges",
905    "was_linked",
906    "edges_at",
907    "what_if",
908    "node_history",
909    "edge_history",
910    "find_similar",
911    "hybrid_search",
912    "remember",
913    "recall",
914    "stats",
915];
916
917/// Which door a store is: which default tool list it gets.
918///
919/// Decided from the store the server opened, once, at startup — not from an
920/// install flag — so one `.mcp.json` serves both and neither has to be
921/// configured for.
922#[derive(Clone, Copy, Debug, PartialEq, Eq)]
923pub(crate) enum Surface {
924    /// A repository was ingested into this store: the `GitSync` marker is
925    /// there, and `explore` has a code graph to explore.
926    CodeGraph,
927    /// Any other store, including an empty one: the fifteen-tool association
928    /// surface, where `explore` would have nothing to answer from.
929    Memory,
930}
931
932impl Surface {
933    /// The tools a default `tools/list` on this surface advertises, in the
934    /// order it advertises them.
935    fn listing(self) -> &'static [&'static str] {
936        match self {
937            Surface::CodeGraph => &CODE_GRAPH_TOOLS,
938            Surface::Memory => &ASSOCIATION_TOOLS,
939        }
940    }
941}
942
943/// The surface the store `db` holds: [`Surface::CodeGraph`] when it carries the
944/// `GitSync` marker `ingest-git` writes, [`Surface::Memory`] otherwise.
945fn surface_of(db: &SharedDb) -> Surface {
946    let ingested = {
947        let g = db.read();
948        g.has_node(crate::mcp_tasks::SYNC_KEY)
949    };
950    if ingested {
951        Surface::CodeGraph
952    } else {
953        Surface::Memory
954    }
955}
956
957/// The tools `tools/list` advertises: the fourteen task tools, then the
958/// graph tools with their descriptions prefixed.
959///
960/// `all` false — the default — lists what `surface` names, **in the order that
961/// surface names it**: three on a code graph, fifteen on a memory store. The
962/// order is the point. A host that defers tool schemas makes a model search
963/// for them, and the list it searches is read top-down, so each surface ranks
964/// its own tools rather than inheriting the task-tools-then-graph-tools order
965/// that only the code door has a reason for.
966///
967/// `all` true lists all twenty-seven in that established order whichever store
968/// this is, which is what `mushroomdb mcp --all-tools` runs and what the
969/// published server card documents: a caller that asked for everything asked
970/// for the whole surface, not for one door's ranking of it.
971///
972/// Either way every tool stays callable: the flag and the surface decide what
973/// is advertised, not what is served.
974fn tools_list(all: bool, surface: Surface) -> Js {
975    let mut served: Vec<Js> = crate::mcp_tasks::task_tools();
976    for mut tool in graph_tools() {
977        if let Some(d) = tool.get("description").and_then(Js::as_str) {
978            let prefixed = format!("{ADVANCED_PREFIX}{d}");
979            tool["description"] = Js::String(prefixed);
980        }
981        served.push(tool);
982    }
983    if all {
984        return json!({ "tools": served });
985    }
986    let listing = surface.listing();
987    let mut tools: Vec<Js> = Vec::with_capacity(listing.len());
988    for name in listing {
989        let Some(tool) = served
990            .iter()
991            .find(|t| t.get("name").and_then(Js::as_str) == Some(*name))
992        else {
993            debug_assert!(false, "{surface:?} lists {name}, which is not served");
994            continue;
995        };
996        tools.push(tool.clone());
997    }
998    json!({ "tools": tools })
999}
1000
1001/// The thirteen graph tools, in the order they have always been listed, with
1002/// their descriptions unprefixed. [`tools_list`] adds the prefix.
1003fn graph_tools() -> Vec<Js> {
1004    let Js::Array(tools) = json!([
1005            {
1006                "name": "query",
1007                "description": "Who may see this, and anything else one pattern can answer — run a Cypher query (read or write) against the graph. Pass 'role' to answer as one of the store's roles: only the nodes that role may see, writes refused. 'mask' is the same restriction written out as an explicit key allow-list. Cypher dialect: MATCH/WHERE/RETURN, CREATE, MERGE, SET, DELETE, with $named parameters in 'params'. A node's key and label read as properties (n.key, n.label) or as key(n)/labels(n). One MATCH takes comma-separated patterns that share variables — MATCH (t)-[:A]->(c), (t)-[:B]->(c) is the intersection of both, and count(DISTINCT t) after WITH counts each t once. WHERE takes STARTS WITH, ENDS WITH, CONTAINS, IN, and a list subscript (n.location[0]) — which is null when the index is out of range, the property is not a list, or the index is not an integer, so a subscript never errors and never matches.",
1008                "inputSchema": {
1009                    "type": "object",
1010                    "properties": {
1011                        "cypher": { "type": "string", "description": "Cypher query text." },
1012                        "params": {
1013                            "type": "object",
1014                            "description": "Named JSON-scalar query parameters."
1015                        },
1016                        "mask": {
1017                            "type": "array",
1018                            "items": { "type": "string" },
1019                            "description": "Optional node key allow-list. When present, only these nodes are visible; write statements are rejected."
1020                        },
1021                        "role": {
1022                            "type": "string",
1023                            "description": "Answer as this role from the store's roles: only the nodes it may see."
1024                        }
1025                    },
1026                    "required": ["cypher"]
1027                }
1028            },
1029            {
1030                "name": "ingest_json",
1031                "description": "Ingest a JSON array of objects as nodes of one label.",
1032                "inputSchema": {
1033                    "type": "object",
1034                    "properties": {
1035                        "label": { "type": "string" },
1036                        "rows_json": {
1037                            "type": "string",
1038                            "description": "JSON text of an array of objects."
1039                        },
1040                        "key_field": { "type": "string" },
1041                        "auto_fk_suffix": { "type": "string" },
1042                        "edges": {
1043                            "type": "array",
1044                            "description": "Optional user edges [{edge_type, src, dst}]."
1045                        }
1046                    },
1047                    "required": ["label", "rows_json"]
1048                }
1049            },
1050            {
1051                "name": "create_rule",
1052                "description": "How should this kind of relationship be derived from now on — declare a rule (RuleDef JSON) and the engine maintains its edges as the data changes. Propose it and show the edges it would derive before creating one.",
1053                "inputSchema": {
1054                    "type": "object",
1055                    "properties": {
1056                        "name": { "type": "string" },
1057                        "src_label": { "type": "string" },
1058                        "dst_label": { "type": "string" },
1059                        "predicate": { "type": "object" },
1060                        "edge_type": { "type": "string" },
1061                        "weight_prop": {
1062                            "type": ["string", "null"],
1063                            "description": "Edge property that stores the score (default: weight)."
1064                        },
1065                        "max_edges": { "type": ["integer", "null"] }
1066                    },
1067                    "required": ["name", "src_label", "dst_label", "predicate", "edge_type"]
1068                }
1069            },
1070            {
1071                "name": "explain",
1072                "description": "Why are A and B related, as a raw array — the same rule-derived edges explain_association renders, for a caller that wants the JSON without asking.",
1073                "inputSchema": {
1074                    "type": "object",
1075                    "properties": {
1076                        "a": { "type": "string", "minLength": 1 },
1077                        "b": { "type": "string", "minLength": 1 }
1078                    },
1079                    "required": ["a", "b"]
1080                }
1081            },
1082            {
1083                "name": "stats",
1084                "description": "How big is this store — live node, edge and rule counts.",
1085                "inputSchema": {
1086                    "type": "object",
1087                    "properties": {}
1088                }
1089            },
1090            {
1091                "name": "node_info",
1092                "description": "What is K — its label and every property it holds.",
1093                "inputSchema": {
1094                    "type": "object",
1095                    "properties": {
1096                        "key": { "type": "string" }
1097                    },
1098                    "required": ["key"]
1099                }
1100            },
1101            {
1102                "name": "upsert_entity",
1103                "description": "Record what is now true about K — insert or update a node by key. If the key exists, updates the supplied properties. If not, creates a new node with the given label and properties. Useful for agent memory: store or refresh an entity without checking existence first.",
1104                "inputSchema": {
1105                    "type": "object",
1106                    "properties": {
1107                        "key": { "type": "string", "description": "Unique node key." },
1108                        "label": { "type": "string", "description": "Node label (required when creating a new entity)." },
1109                        "props": {
1110                            "type": "object",
1111                            "description": "Properties to set. Values must be scalars (string, number, bool) or arrays of scalars."
1112                        }
1113                    },
1114                    "required": ["key", "props"]
1115                }
1116            },
1117            {
1118                "name": "find_similar",
1119                "description": "What is most like this — two modes: (1) Vector search — provide `vector` (and optionally `field`, `label`, `k`, `min`) to find the k most similar nodes by cosine similarity using the HNSW index when available, brute-force otherwise. (2) Edge traversal — provide `key` (and optionally `edge_type`, `limit`) to return neighbors previously connected by a derived rule edge. Results from mode 2 come only from edges already derived by a VectorSimilar rule. In both modes, the optional `mask` array limits visibility: hidden nodes never appear in results, and a hidden query key in edge mode behaves identically to a nonexistent key.",
1120                "inputSchema": {
1121                    "type": "object",
1122                    "properties": {
1123                        "vector": {
1124                            "type": "array",
1125                            "items": { "type": "number" },
1126                            "description": "Query embedding vector for vector-similarity search. When present, vector-search mode is used and `key` is ignored."
1127                        },
1128                        "field": { "type": "string", "description": "Property field holding the embedding vectors (default: embedding). Used in vector-search mode." },
1129                        "label": { "type": "string", "description": "Restrict search to nodes with this label. Empty string means all labels. Used in vector-search mode." },
1130                        "k": { "type": "integer", "description": "Maximum results to return in vector-search mode (default: 10)." },
1131                        "min": { "type": "number", "description": "Minimum cosine similarity threshold in vector-search mode (default: 0.8)." },
1132                        "mask": {
1133                            "type": "array",
1134                            "items": { "type": "string" },
1135                            "description": "Optional node key allow-list for vector-search mode. When present, only nodes whose key appears in this list are eligible for results. Hidden nodes are excluded before k-truncation so callers still receive up to k visible hits. Unknown keys are silently ignored."
1136                        },
1137                        "key": { "type": "string", "description": "Source node key for edge-traversal mode." },
1138                        "edge_type": { "type": "string", "description": "Edge type to filter by in edge-traversal mode (default: SIMILAR)." },
1139                        "limit": { "type": "integer", "description": "Maximum neighbors to return in edge-traversal mode (default: 10)." }
1140                    }
1141                }
1142            },
1143            {
1144                "name": "hybrid_search",
1145                "description": "What matches these words and this vector at once — Reciprocal Rank Fusion (RRF) over fulltext + vector results. Provide `query_text` and `text_field` for the fulltext leg. Optionally provide `vector` (embedding array) and `vector_field` (default: embedding) for the vector leg; omitting `vector` gives text-only ranking through the same RRF path. `label` restricts the vector search to nodes with that label (required for brute-force; omit to rely on HNSW rules). `k` controls result count (default: 10). RRF constant is fixed at 60; scores are 1/(60+rank) summed over lists a node appears in.",
1146                "inputSchema": {
1147                    "type": "object",
1148                    "properties": {
1149                        "query_text": { "type": "string", "description": "Fulltext query string." },
1150                        "text_field": { "type": "string", "description": "Property field to search with fulltext." },
1151                        "vector": {
1152                            "type": "array",
1153                            "items": { "type": "number" },
1154                            "description": "Query embedding vector. Omit for text-only ranking."
1155                        },
1156                        "vector_field": { "type": "string", "description": "Property field holding embedding vectors (default: embedding)." },
1157                        "label": { "type": "string", "description": "Restrict vector search to nodes with this label. Required when relying on brute-force (no HNSW rule covers the field). If omitted, the vector leg always returns empty results (no rule-created HNSW index covers the unlabeled path); ranking is text-only in that case." },
1158                        "k": { "type": "integer", "description": "Maximum results to return (default: 10)." }
1159                    },
1160                    "required": ["query_text", "text_field"]
1161                }
1162            },
1163            {
1164                "name": "node_history",
1165                "description": "What has happened to K — every recorded change to one node, newest last. Events include NodeInserted, PropSet, PropRemoved, EdgeAdded, EdgeRemoved, and NodeDeleted. The response includes `total_commits` (the horizon upper bound). History is WAL-scoped — pre-snapshot commits are not visible.",
1166                "inputSchema": {
1167                    "type": "object",
1168                    "properties": {
1169                        "key": { "type": "string", "description": "Node key to look up." }
1170                    },
1171                    "required": ["key"]
1172                }
1173            },
1174            {
1175                "name": "edge_history",
1176                "description": "When did A and B become linked, and when did it break — the full add/retract lifecycle for every edge between the two keys. Includes derived (rule-attributed) edges via DerivedEdgeAdded/DerivedEdgeRetracted WAL markers. The response includes `total_commits` (the horizon upper bound).",
1177                "inputSchema": {
1178                    "type": "object",
1179                    "properties": {
1180                        "a": { "type": "string", "minLength": 1, "description": "First node key." },
1181                        "b": { "type": "string", "minLength": 1, "description": "Second node key." }
1182                    },
1183                    "required": ["a", "b"]
1184                }
1185            },
1186            {
1187                "name": "was_linked",
1188                "description": "Were A and B linked at commit C — whether an edge of `edge_type` existed between the two keys (either direction) at that WAL commit. Returns an error when `at_commit` is outside the visible horizon (`0..total_commits`).",
1189                "inputSchema": {
1190                    "type": "object",
1191                    "properties": {
1192                        "a": { "type": "string", "minLength": 1, "description": "First node key." },
1193                        "b": { "type": "string", "minLength": 1, "description": "Second node key." },
1194                        "edge_type": { "type": "string", "minLength": 1, "description": "Edge type to check." },
1195                        "at_commit": { "type": "integer", "minimum": 0, "description": "0-based WAL commit index to query." }
1196                    },
1197                    "required": ["a", "b", "edge_type", "at_commit"]
1198                }
1199            },
1200            {
1201                "name": "rename_node",
1202                "description": "Rename K — the key changes and nothing else does. The dense id and all edges/properties remain stable. Returns 404 if `old_key` does not exist, 409 if `new_key` is already taken.",
1203                "inputSchema": {
1204                    "type": "object",
1205                    "properties": {
1206                        "old_key": { "type": "string", "minLength": 1, "description": "Current node key." },
1207                        "new_key": { "type": "string", "minLength": 1, "description": "Desired new node key." }
1208                    },
1209                    "required": ["old_key", "new_key"]
1210                }
1211            }
1212    ]) else {
1213        unreachable!("the literal above is an array")
1214    };
1215    tools
1216}
1217
1218fn tool_ok(payload: Js) -> Js {
1219    json!({
1220        "content": [{ "type": "text", "text": payload.to_string() }]
1221    })
1222}
1223
1224/// A task tool's result: one text block, and no `structuredContent`.
1225///
1226/// The report used to ride along beside the digest, repeating it verbatim
1227/// under a `text` key. Nothing bound it — no task tool declares an
1228/// `outputSchema` — and it tripled the size of every reply, so a caller that
1229/// wants the numbers now asks for them with `json: true` and gets the report
1230/// *as* the text.
1231fn task_ok(text: &str) -> Js {
1232    json!({
1233        "content": [{ "type": "text", "text": text }]
1234    })
1235}
1236
1237fn tool_err(message: &str) -> Js {
1238    json!({
1239        "content": [{ "type": "text", "text": message }],
1240        "isError": true
1241    })
1242}
1243
1244fn write_result(writer: &mut impl Write, id: Option<Js>, result: Js) -> io::Result<()> {
1245    write_json(
1246        writer,
1247        &json!({
1248            "jsonrpc": "2.0",
1249            "id": id.unwrap_or(Js::Null),
1250            "result": result
1251        }),
1252    )
1253}
1254
1255fn write_error(
1256    writer: &mut impl Write,
1257    id: Option<Js>,
1258    code: i64,
1259    message: &str,
1260) -> io::Result<()> {
1261    write_json(
1262        writer,
1263        &json!({
1264            "jsonrpc": "2.0",
1265            "id": id.unwrap_or(Js::Null),
1266            "error": { "code": code, "message": message }
1267        }),
1268    )
1269}
1270
1271fn write_json(writer: &mut impl Write, value: &Js) -> io::Result<()> {
1272    let s = serde_json::to_string(value).map_err(io::Error::other)?;
1273    writeln!(writer, "{s}")?;
1274    writer.flush()
1275}
1276
1277// ---------------------------------------------------------------------------
1278// Tests: MCP tool round-trips via stdio
1279// ---------------------------------------------------------------------------
1280
1281#[cfg(test)]
1282mod tests {
1283    use super::*;
1284    use core_api::{AutoFk, IngestOptions, Predicate, RuleDef, Value};
1285    use std::path::PathBuf;
1286    use std::sync::atomic::{AtomicU64, Ordering};
1287
1288    fn tmp_dir() -> PathBuf {
1289        static SEQ: AtomicU64 = AtomicU64::new(0);
1290        let n = SEQ.fetch_add(1, Ordering::Relaxed);
1291        let d = std::env::temp_dir().join(format!("mcp-test-{}-{}", std::process::id(), n));
1292        // These stores are never cleaned up, so a process id the OS hands out
1293        // again lands on a previous run's data and every assertion about counts
1294        // fails. `tests/mcp.rs::tmp` already clears its path for this reason.
1295        let _ = std::fs::remove_dir_all(&d);
1296        d
1297    }
1298
1299    /// Open a SharedDb with two Person nodes and one derived SIMILAR edge.
1300    fn demo_db() -> SharedDb {
1301        let db = SharedDb::open(&tmp_dir()).expect("open");
1302        {
1303            let mut g = db.write();
1304            let opts = IngestOptions {
1305                key_field: "id".into(),
1306                auto_fk: AutoFk::Off,
1307            };
1308            // Two people with identical embeddings → will fire SIMILAR rule.
1309            let people: Vec<BTreeMap<String, Value>> = vec![
1310                [
1311                    ("id", Value::Str("alice".into())),
1312                    ("name", Value::Str("Alice".into())),
1313                    (
1314                        "emb",
1315                        Value::List(vec![Value::Float(1.0), Value::Float(0.0)]),
1316                    ),
1317                ]
1318                .into_iter()
1319                .map(|(k, v)| (k.to_string(), v))
1320                .collect(),
1321                [
1322                    ("id", Value::Str("bob".into())),
1323                    ("name", Value::Str("Bob".into())),
1324                    (
1325                        "emb",
1326                        Value::List(vec![Value::Float(1.0), Value::Float(0.0)]),
1327                    ),
1328                ]
1329                .into_iter()
1330                .map(|(k, v)| (k.to_string(), v))
1331                .collect(),
1332            ];
1333            g.ingest("Person", people, &opts).expect("ingest");
1334
1335            // Rule: VectorSimilar on emb → SIMILAR edge (cosine(ident,ident)=1.0 ≥ 0.9).
1336            g.create_rule(RuleDef {
1337                name: "sim_emb".into(),
1338                src_label: "Person".into(),
1339                dst_label: "Person".into(),
1340                predicate: Predicate::VectorSimilar {
1341                    field: "emb".into(),
1342                    min: 0.9,
1343                },
1344                edge_type: "SIMILAR".into(),
1345                weight_prop: Some("score".into()),
1346                max_edges: None,
1347                approximate: false,
1348                via_label: None,
1349                via_edge: None,
1350                via_dir: None,
1351            })
1352            .expect("rule");
1353        }
1354        db
1355    }
1356
1357    fn roundtrip(db: &SharedDb, request: &str) -> Js {
1358        roundtrip_with(db, false, request)
1359    }
1360
1361    fn roundtrip_with(db: &SharedDb, all_tools: bool, request: &str) -> Js {
1362        let input = format!("{request}\n");
1363        let mut output = Vec::new();
1364        run_mcp_stdio_with(db.clone(), None, all_tools, input.as_bytes(), &mut output)
1365            .expect("mcp");
1366        let s = std::str::from_utf8(&output).expect("utf8");
1367        serde_json::from_str(s.trim()).expect("json response")
1368    }
1369
1370    fn tool_call(db: &SharedDb, id: u64, tool: &str, args: Js) -> Js {
1371        let req = json!({
1372            "jsonrpc": "2.0",
1373            "id": id,
1374            "method": "tools/call",
1375            "params": { "name": tool, "arguments": args }
1376        });
1377        roundtrip(db, &req.to_string())
1378    }
1379
1380    /// Unwrap the `text` field from a successful tool response.
1381    fn tool_text(resp: &Js) -> Js {
1382        let text = resp["result"]["content"][0]["text"]
1383            .as_str()
1384            .expect("content[0].text");
1385        serde_json::from_str(text).expect("tool text is json")
1386    }
1387
1388    fn is_error(resp: &Js) -> bool {
1389        resp["result"]["isError"].as_bool().unwrap_or(false)
1390    }
1391
1392    fn tool_err_text(resp: &Js) -> String {
1393        resp["result"]["content"][0]["text"]
1394            .as_str()
1395            .unwrap_or("")
1396            .to_string()
1397    }
1398
1399    // --- existing tools ---
1400
1401    #[test]
1402    fn test_tools_list_includes_all_expected() {
1403        let db = demo_db();
1404        let resp = roundtrip_with(
1405            &db,
1406            true,
1407            r#"{"jsonrpc":"2.0","id":1,"method":"tools/list"}"#,
1408        );
1409        let tools = resp["result"]["tools"].as_array().expect("tools array");
1410        let names: Vec<&str> = tools
1411            .iter()
1412            .map(|t| t["name"].as_str().expect("name"))
1413            .collect();
1414        for expected in &[
1415            // The fourteen task tools, first and in order.
1416            "explore",
1417            "map",
1418            "context",
1419            "impact",
1420            "owners",
1421            "why",
1422            "explain_association",
1423            "node_edges",
1424            "neighborhood",
1425            "edges_at",
1426            "what_if",
1427            "recall",
1428            "remember",
1429            "sync",
1430            // The thirteen graph tools.
1431            "query",
1432            "ingest_json",
1433            "create_rule",
1434            "explain",
1435            "stats",
1436            "node_info",
1437            "upsert_entity",
1438            "find_similar",
1439            "hybrid_search",
1440            "node_history",
1441            "edge_history",
1442            "was_linked",
1443            "rename_node",
1444        ] {
1445            assert!(names.contains(expected), "missing tool: {expected}");
1446        }
1447        assert_eq!(
1448            names.len(),
1449            27,
1450            "expected exactly 27 tools, got {}",
1451            names.len()
1452        );
1453        assert_eq!(
1454            &names[..14],
1455            [
1456                "explore",
1457                "map",
1458                "context",
1459                "impact",
1460                "owners",
1461                "why",
1462                "explain_association",
1463                "node_edges",
1464                "neighborhood",
1465                "edges_at",
1466                "what_if",
1467                "recall",
1468                "remember",
1469                "sync"
1470            ],
1471            "the task tools come first, in order"
1472        );
1473        assert_eq!(names[14], "query", "the graph tools follow them");
1474    }
1475
1476    /// Binding: on a store no repository was ingested into, the default
1477    /// listing is the fifteen association tools, in [`ASSOCIATION_TOOLS`]
1478    /// order, and nothing else.
1479    #[test]
1480    fn tools_list_defaults_to_fifteen_on_a_memory_store() {
1481        let db = demo_db();
1482        let resp = roundtrip(&db, r#"{"jsonrpc":"2.0","id":1,"method":"tools/list"}"#);
1483        let names: Vec<&str> = resp["result"]["tools"]
1484            .as_array()
1485            .expect("tools array")
1486            .iter()
1487            .map(|t| t["name"].as_str().expect("name"))
1488            .collect();
1489        assert_eq!(names, ASSOCIATION_TOOLS.to_vec());
1490    }
1491
1492    /// Binding: [`ASSOCIATION_TOOLS`] is a surface of its own, not the code
1493    /// door's list with a name changed.
1494    ///
1495    /// It keeps the two task tools an entity store can answer with — the notes
1496    /// it wrote and the notes it kept — and none of the seven that read a code
1497    /// graph there is none of. Every name in it is served.
1498    #[test]
1499    fn the_association_surface_is_entity_tools_only() {
1500        for kept in ["remember", "recall", "explain_association"] {
1501            assert!(
1502                ASSOCIATION_TOOLS.contains(&kept),
1503                "{kept} answers on an entity graph and must be listed"
1504            );
1505        }
1506        for code_only in [
1507            "explore", "map", "context", "impact", "owners", "why", "sync",
1508        ] {
1509            assert!(
1510                !ASSOCIATION_TOOLS.contains(&code_only),
1511                "{code_only} reads a code graph and must not be listed on a memory store"
1512            );
1513        }
1514        let served: Vec<String> = crate::mcp_tasks::task_tools()
1515            .iter()
1516            .chain(graph_tools().iter())
1517            .filter_map(|t| t.get("name").and_then(Js::as_str))
1518            .map(str::to_string)
1519            .collect();
1520        for name in ASSOCIATION_TOOLS {
1521            assert!(
1522                served.iter().any(|s| s == name),
1523                "{name} is listed but not served"
1524            );
1525        }
1526        assert!(
1527            CODE_GRAPH_TOOLS.contains(&"explore"),
1528            "and `explore` is the task tool the other surface lists"
1529        );
1530    }
1531
1532    /// Binding: the same server on a store carrying the `GitSync` marker lists
1533    /// three. One tool to find, one to query, one to size the store.
1534    #[test]
1535    fn tools_list_is_three_tools_on_a_code_graph_store() {
1536        let db = demo_db();
1537        db.write()
1538            .insert_node(
1539                "GitSync",
1540                crate::mcp_tasks::SYNC_KEY,
1541                vec![("id".into(), Value::Str(crate::mcp_tasks::SYNC_KEY.into()))],
1542            )
1543            .expect("marker");
1544        let resp = roundtrip(&db, r#"{"jsonrpc":"2.0","id":1,"method":"tools/list"}"#);
1545        let names: Vec<&str> = resp["result"]["tools"]
1546            .as_array()
1547            .expect("tools array")
1548            .iter()
1549            .map(|t| t["name"].as_str().expect("name"))
1550            .collect();
1551        assert_eq!(names, ["explore", "query", "stats"]);
1552    }
1553
1554    #[test]
1555    fn test_stats_returns_node_count() {
1556        let db = demo_db();
1557        let resp = tool_call(&db, 1, "stats", json!({}));
1558        assert!(!is_error(&resp));
1559        let result = tool_text(&resp);
1560        assert_eq!(result["nodes_live"], 2);
1561    }
1562
1563    #[test]
1564    fn test_query_runs_cypher() {
1565        let db = demo_db();
1566        let resp = tool_call(
1567            &db,
1568            1,
1569            "query",
1570            json!({ "cypher": "MATCH (n:Person) RETURN n.name ORDER BY n.name" }),
1571        );
1572        assert!(!is_error(&resp));
1573        let result = tool_text(&resp);
1574        // columns + 2 rows
1575        assert_eq!(result["columns"], json!(["n.name"]));
1576        assert_eq!(result["rows"].as_array().map(|r| r.len()), Some(2));
1577    }
1578
1579    #[test]
1580    fn test_query_create_is_a_write() {
1581        let db = SharedDb::open(&tmp_dir()).expect("open");
1582        let resp = tool_call(
1583            &db,
1584            1,
1585            "query",
1586            json!({ "cypher": "CREATE (n:L {id: 'k'}) RETURN n" }),
1587        );
1588        assert!(
1589            !is_error(&resp),
1590            "CREATE via MCP query must succeed: {resp}"
1591        );
1592        let stats = tool_text(&tool_call(&db, 2, "stats", json!({})));
1593        assert_eq!(stats["nodes_live"], 1);
1594    }
1595
1596    #[test]
1597    fn test_ingest_json_inserts_nodes() {
1598        let db = demo_db();
1599        let resp = tool_call(
1600            &db,
1601            1,
1602            "ingest_json",
1603            json!({
1604                "label": "Person",
1605                "rows_json": r#"[{"id":"carol","name":"Carol"}]"#,
1606                "key_field": "id"
1607            }),
1608        );
1609        assert!(!is_error(&resp));
1610        // Verify node visible via stats
1611        let stats = tool_text(&tool_call(&db, 2, "stats", json!({})));
1612        assert_eq!(stats["nodes_live"], 3);
1613    }
1614
1615    #[test]
1616    fn test_node_info_returns_props() {
1617        let db = demo_db();
1618        let resp = tool_call(&db, 1, "node_info", json!({ "key": "alice" }));
1619        assert!(!is_error(&resp));
1620        let result = tool_text(&resp);
1621        assert_eq!(result["key"], "alice");
1622        assert_eq!(result["label"], "Person");
1623        assert_eq!(result["props"]["name"], "Alice");
1624    }
1625
1626    /// Binding: `node_edges` groups by edge type and names the rule and score
1627    /// behind each derived edge, in the report as in the digest.
1628    #[test]
1629    fn test_node_edges_returns_edges() {
1630        let db = demo_db();
1631        let resp = tool_call(
1632            &db,
1633            1,
1634            "node_edges",
1635            json!({ "key": "alice", "json": true }),
1636        );
1637        assert!(!is_error(&resp));
1638        let result = tool_text(&resp);
1639        assert_eq!(result["key"], "alice");
1640        let types = result["types"].as_array().expect("types");
1641        assert!(
1642            !types.is_empty(),
1643            "alice should have at least one edge type"
1644        );
1645        let similar = types
1646            .iter()
1647            .find(|t| t["edge_type"] == "SIMILAR")
1648            .expect("the rule's edge type");
1649        // A symmetric rule derives the edge both ways, and both are listed
1650        // with the direction that tells them apart.
1651        assert_eq!(similar["count"], json!(2));
1652        let edges = similar["edges"].as_array().expect("edges");
1653        let dirs: Vec<&str> = edges
1654            .iter()
1655            .map(|e| e["direction"].as_str().expect("direction"))
1656            .collect();
1657        assert!(dirs.contains(&"out") && dirs.contains(&"in"), "{similar}");
1658        for edge in edges {
1659            assert_eq!(edge["other"], json!("bob"));
1660            assert_eq!(edge["derived"], json!(true));
1661            assert_eq!(edge["rule"], json!("sim_emb"));
1662            assert_eq!(edge["score"], json!(1.0));
1663            assert!(
1664                edge["predicate"]
1665                    .as_str()
1666                    .unwrap_or("")
1667                    .contains("vector_similar"),
1668                "the predicate travels with the edge: {edge}"
1669            );
1670        }
1671    }
1672
1673    /// Binding: a depth-1 `neighborhood` is the same relationship listing, and
1674    /// anything deeper is still the traversal table.
1675    #[test]
1676    fn test_neighborhood_traverses_one_hop() {
1677        let db = demo_db();
1678        let resp = tool_call(
1679            &db,
1680            1,
1681            "neighborhood",
1682            json!({ "key": "alice", "depth": 1, "json": true }),
1683        );
1684        assert!(!is_error(&resp));
1685        let result = tool_text(&resp);
1686        assert_eq!(result["key"], "alice");
1687        assert!(result["types"].as_array().is_some(), "{result}");
1688
1689        let deep = tool_call(
1690            &db,
1691            2,
1692            "neighborhood",
1693            json!({ "key": "alice", "depth": 2 }),
1694        );
1695        assert!(!is_error(&deep));
1696        let table = tool_text(&deep);
1697        assert_eq!(table["columns"], json!(["key", "label", "depth"]));
1698        assert!(table["rows"].as_array().is_some());
1699    }
1700
1701    #[test]
1702    fn test_explain_returns_rule_info() {
1703        let db = demo_db();
1704        let resp = tool_call(&db, 1, "explain", json!({ "a": "alice", "b": "bob" }));
1705        assert!(!is_error(&resp));
1706        let result = tool_text(&resp);
1707        let arr = result.as_array().expect("explain returns array");
1708        assert!(!arr.is_empty(), "expected at least one explanation");
1709        assert_eq!(arr[0]["rule"], "sim_emb");
1710    }
1711
1712    #[test]
1713    fn test_create_rule_backfills() {
1714        let db = SharedDb::open(&tmp_dir()).expect("open");
1715        {
1716            let mut g = db.write();
1717            let opts = IngestOptions {
1718                key_field: "id".into(),
1719                auto_fk: AutoFk::Off,
1720            };
1721            let rows: Vec<BTreeMap<String, Value>> = vec![
1722                [
1723                    ("id", Value::Str("x".into())),
1724                    ("tag", Value::Str("a".into())),
1725                ]
1726                .into_iter()
1727                .map(|(k, v)| (k.to_string(), v))
1728                .collect(),
1729                [
1730                    ("id", Value::Str("y".into())),
1731                    ("tag", Value::Str("a".into())),
1732                ]
1733                .into_iter()
1734                .map(|(k, v)| (k.to_string(), v))
1735                .collect(),
1736            ];
1737            g.ingest("Item", rows, &opts).expect("ingest");
1738        }
1739        let resp = tool_call(
1740            &db,
1741            1,
1742            "create_rule",
1743            json!({
1744                "name": "same_tag",
1745                "src_label": "Item",
1746                "dst_label": "Item",
1747                "predicate": { "FieldEqual": { "field": "tag" } },
1748                "edge_type": "SAME_TAG"
1749            }),
1750        );
1751        assert!(!is_error(&resp));
1752        let result = tool_text(&resp);
1753        assert_eq!(result["ok"], true);
1754        // Derived edges should now exist.
1755        let edges_resp = tool_call(&db, 2, "node_edges", json!({ "key": "x", "json": true }));
1756        let edges_result = tool_text(&edges_resp);
1757        let types = edges_result["types"].as_array().expect("types");
1758        assert!(
1759            types.iter().any(|t| t["edge_type"] == "SAME_TAG"),
1760            "SAME_TAG edge not found after create_rule"
1761        );
1762    }
1763
1764    // --- new tools ---
1765
1766    #[test]
1767    fn test_upsert_entity_creates_new_node() {
1768        let db = demo_db();
1769        let resp = tool_call(
1770            &db,
1771            1,
1772            "upsert_entity",
1773            json!({
1774                "key": "carol",
1775                "label": "Person",
1776                "props": { "name": "Carol", "age": 30 }
1777            }),
1778        );
1779        assert!(!is_error(&resp));
1780        let result = tool_text(&resp);
1781        assert_eq!(result["ok"], true);
1782        assert_eq!(result["created"], true);
1783        assert_eq!(result["key"], "carol");
1784        // Verify node exists
1785        let info = tool_text(&tool_call(&db, 2, "node_info", json!({ "key": "carol" })));
1786        assert_eq!(info["props"]["name"], "Carol");
1787    }
1788
1789    #[test]
1790    fn test_upsert_entity_updates_existing_node() {
1791        let db = demo_db();
1792        let resp = tool_call(
1793            &db,
1794            1,
1795            "upsert_entity",
1796            json!({
1797                "key": "alice",
1798                "props": { "name": "Alice Updated" }
1799            }),
1800        );
1801        assert!(!is_error(&resp));
1802        let result = tool_text(&resp);
1803        assert_eq!(result["ok"], true);
1804        assert_eq!(result["created"], false);
1805        assert_eq!(result["updated_fields"], 1);
1806        // Verify prop changed
1807        let info = tool_text(&tool_call(&db, 2, "node_info", json!({ "key": "alice" })));
1808        assert_eq!(info["props"]["name"], "Alice Updated");
1809    }
1810
1811    #[test]
1812    fn test_upsert_entity_missing_label_on_create_is_error() {
1813        let db = demo_db();
1814        let resp = tool_call(
1815            &db,
1816            1,
1817            "upsert_entity",
1818            json!({ "key": "new-node", "props": { "x": 1 } }),
1819        );
1820        assert!(is_error(&resp), "should error without label for new node");
1821    }
1822
1823    #[test]
1824    fn test_find_similar_returns_similar_edges() {
1825        let db = demo_db();
1826        let resp = tool_call(
1827            &db,
1828            1,
1829            "find_similar",
1830            json!({ "key": "alice", "edge_type": "SIMILAR" }),
1831        );
1832        assert!(!is_error(&resp));
1833        let result = tool_text(&resp);
1834        assert_eq!(result["key"], "alice");
1835        assert_eq!(result["edge_type"], "SIMILAR");
1836        let similar = result["similar"].as_array().expect("similar array");
1837        assert!(!similar.is_empty(), "expected SIMILAR neighbors for alice");
1838        assert_eq!(similar[0]["neighbor_key"], "bob");
1839    }
1840
1841    #[test]
1842    fn test_find_similar_limit_respected() {
1843        let db = demo_db();
1844        let resp = tool_call(
1845            &db,
1846            1,
1847            "find_similar",
1848            json!({ "key": "alice", "edge_type": "SIMILAR", "limit": 0 }),
1849        );
1850        assert!(!is_error(&resp));
1851        let result = tool_text(&resp);
1852        let similar = result["similar"].as_array().expect("similar array");
1853        assert_eq!(similar.len(), 0);
1854    }
1855
1856    /// When `min` is omitted from a vector-mode find_similar call, the server
1857    /// must apply the spec default of 0.8.  A node whose cosine similarity to
1858    /// the query is 0.0 (orthogonal) must not appear in the results.
1859    #[test]
1860    fn test_find_similar_vector_default_min_is_0_8() {
1861        let db = SharedDb::open(&tmp_dir()).expect("open");
1862        {
1863            let mut g = db.write();
1864            // close: [1,0] → cosine 1.0 with query [1,0] (above 0.8)
1865            g.insert_node(
1866                "Item",
1867                "close",
1868                vec![(
1869                    "emb".into(),
1870                    Value::List(vec![Value::Float(1.0), Value::Float(0.0)]),
1871                )],
1872            )
1873            .unwrap();
1874            // far: [0,1] → cosine 0.0 with query [1,0] (below 0.8, must be excluded)
1875            g.insert_node(
1876                "Item",
1877                "far",
1878                vec![(
1879                    "emb".into(),
1880                    Value::List(vec![Value::Float(0.0), Value::Float(1.0)]),
1881                )],
1882            )
1883            .unwrap();
1884        }
1885
1886        // No `min` in the request — must default to 0.8.
1887        let resp = tool_call(
1888            &db,
1889            1,
1890            "find_similar",
1891            json!({
1892                "vector": [1.0, 0.0],
1893                "field": "emb",
1894                "label": "Item",
1895                "k": 10
1896            }),
1897        );
1898        assert!(!is_error(&resp), "vector search must not error");
1899        let result = tool_text(&resp);
1900        let results = result["results"].as_array().expect("results array");
1901
1902        let keys: Vec<&str> = results.iter().filter_map(|r| r["key"].as_str()).collect();
1903        assert!(
1904            keys.contains(&"close"),
1905            "close node (sim=1.0) must be included"
1906        );
1907        assert!(
1908            !keys.contains(&"far"),
1909            "far node (sim=0.0) must be excluded by default min=0.8"
1910        );
1911    }
1912
1913    /// `find_similar` with `mask` must exclude hidden node keys from results.
1914    #[test]
1915    fn test_find_similar_vector_mask_excludes_hidden() {
1916        let db = SharedDb::open(&tmp_dir()).expect("open");
1917        {
1918            let mut g = db.write();
1919            // visible: [1,0] — should appear in results.
1920            g.insert_node(
1921                "Item",
1922                "visible",
1923                vec![(
1924                    "emb".into(),
1925                    Value::List(vec![Value::Float(1.0), Value::Float(0.0)]),
1926                )],
1927            )
1928            .unwrap();
1929            // hidden: [1,0] — same direction as query but must not appear.
1930            g.insert_node(
1931                "Item",
1932                "hidden",
1933                vec![(
1934                    "emb".into(),
1935                    Value::List(vec![Value::Float(1.0), Value::Float(0.0)]),
1936                )],
1937            )
1938            .unwrap();
1939        }
1940
1941        let resp = tool_call(
1942            &db,
1943            1,
1944            "find_similar",
1945            json!({
1946                "vector": [1.0, 0.0],
1947                "field": "emb",
1948                "label": "Item",
1949                "k": 10,
1950                "min": 0.0,
1951                "mask": ["visible"]
1952            }),
1953        );
1954        assert!(!is_error(&resp), "masked vector search must not error");
1955        let result = tool_text(&resp);
1956        let results = result["results"].as_array().expect("results array");
1957
1958        let keys: Vec<&str> = results.iter().filter_map(|r| r["key"].as_str()).collect();
1959        assert!(
1960            keys.contains(&"visible"),
1961            "visible node must appear in masked results"
1962        );
1963        assert!(
1964            !keys.contains(&"hidden"),
1965            "hidden node must be excluded by mask"
1966        );
1967    }
1968
1969    /// `find_similar` with `mask` — bad mask value returns a tool error.
1970    #[test]
1971    fn test_find_similar_vector_mask_bad_type_is_error() {
1972        let db = SharedDb::open(&tmp_dir()).expect("open");
1973        let resp = tool_call(
1974            &db,
1975            1,
1976            "find_similar",
1977            json!({
1978                "vector": [1.0, 0.0],
1979                "field": "emb",
1980                "k": 5,
1981                "mask": [42]
1982            }),
1983        );
1984        assert!(
1985            is_error(&resp),
1986            "non-string mask element must produce a tool error"
1987        );
1988    }
1989
1990    /// Edge-traversal mode with `mask` must exclude hidden neighbors.
1991    #[test]
1992    fn test_find_similar_edge_mask_excludes_hidden_neighbor() {
1993        let db = SharedDb::open(&tmp_dir()).expect("open");
1994        {
1995            let mut g = db.write();
1996            g.insert_node("P", "alice", vec![]).unwrap();
1997            g.insert_node("P", "bob", vec![]).unwrap(); // visible
1998            g.insert_node("P", "carol", vec![]).unwrap(); // hidden
1999            g.insert_edge("KNOWS", "alice", "bob").unwrap();
2000            g.insert_edge("KNOWS", "alice", "carol").unwrap();
2001        }
2002        // Mask: alice and bob visible; carol hidden.
2003        let resp = tool_call(
2004            &db,
2005            1,
2006            "find_similar",
2007            json!({
2008                "key": "alice",
2009                "edge_type": "KNOWS",
2010                "mask": ["alice", "bob"]
2011            }),
2012        );
2013        assert!(!is_error(&resp), "masked edge search must not error");
2014        let result = tool_text(&resp);
2015        let similar = result["similar"].as_array().expect("similar array");
2016        let neighbors: Vec<&str> = similar
2017            .iter()
2018            .filter_map(|e| e["neighbor_key"].as_str())
2019            .collect();
2020        assert!(neighbors.contains(&"bob"), "bob (visible) must appear");
2021        assert!(
2022            !neighbors.contains(&"carol"),
2023            "carol (hidden) must be excluded"
2024        );
2025    }
2026
2027    /// Edge-traversal mode with `mask`: a hidden query key must not reveal
2028    /// its existence — response must be a tool error identical to a nonexistent key.
2029    #[test]
2030    fn test_find_similar_edge_mask_hidden_key_is_not_found() {
2031        let db = SharedDb::open(&tmp_dir()).expect("open");
2032        {
2033            let mut g = db.write();
2034            g.insert_node("P", "alice", vec![]).unwrap();
2035            g.insert_node("P", "bob", vec![]).unwrap();
2036        }
2037        // alice exists but is not in the mask — must look like not-found.
2038        let resp_masked = tool_call(
2039            &db,
2040            1,
2041            "find_similar",
2042            json!({ "key": "alice", "edge_type": "KNOWS", "mask": ["bob"] }),
2043        );
2044        // ghost never exists — use as the reference for "not found".
2045        let resp_ghost = tool_call(
2046            &db,
2047            2,
2048            "find_similar",
2049            json!({ "key": "ghost", "edge_type": "KNOWS" }),
2050        );
2051        assert!(
2052            is_error(&resp_masked),
2053            "hidden query key must produce a tool error"
2054        );
2055        assert!(
2056            is_error(&resp_ghost),
2057            "nonexistent key must produce a tool error"
2058        );
2059        // Both errors must carry the same shape (both are key-not-found).
2060        assert_eq!(
2061            tool_err_text(&resp_masked).contains("alice"),
2062            tool_err_text(&resp_ghost).contains("ghost"),
2063            "error messages should follow same not-found template"
2064        );
2065    }
2066
2067    /// Binding: `explain_association` now answers in prose, and the report
2068    /// behind it — what `json: true` returns — is still `explain`'s array,
2069    /// with one `evidence` object added per relationship and every other
2070    /// field unchanged.
2071    #[test]
2072    fn test_explain_association_same_as_explain() {
2073        let db = demo_db();
2074        let explain = tool_text(&tool_call(
2075            &db,
2076            1,
2077            "explain",
2078            json!({ "a": "alice", "b": "bob" }),
2079        ));
2080        let assoc = tool_text(&tool_call(
2081            &db,
2082            2,
2083            "explain_association",
2084            json!({ "a": "alice", "b": "bob", "json": true }),
2085        ));
2086        let explain: Vec<Js> = serde_json::from_value(explain).expect("explain array");
2087        let mut assoc: Vec<Js> = serde_json::from_value(assoc).expect("assoc array");
2088        for row in &mut assoc {
2089            let ev = row
2090                .as_object_mut()
2091                .expect("object")
2092                .remove("evidence")
2093                .expect("every derived edge carries its evidence");
2094            assert!(
2095                ev["similarity"].is_number(),
2096                "a vector_similar edge reports the cosine it scored: {ev}"
2097            );
2098        }
2099        assert_eq!(explain, assoc, "evidence is the only addition");
2100
2101        let prose = tool_call(
2102            &db,
2103            3,
2104            "explain_association",
2105            json!({ "a": "alice", "b": "bob" }),
2106        );
2107        let text = prose["result"]["content"][0]["text"]
2108            .as_str()
2109            .expect("text content");
2110        assert!(
2111            text.contains("mushroomdb explain — alice ↔ bob:"),
2112            "the default reply is the digest: {text}"
2113        );
2114    }
2115
2116    // ── history tools ──────────────────────────────────────────────────────────
2117
2118    /// `edge_history` must return the derived-edge lifecycle (Added event with
2119    /// rule attribution) and include the `total_commits` horizon field.
2120    #[test]
2121    fn test_edge_history_returns_derived_lifecycle_with_rule() {
2122        let db = demo_db(); // alice+bob + sim_emb rule → SIMILAR derived edge
2123        let resp = tool_call(&db, 1, "edge_history", json!({ "a": "alice", "b": "bob" }));
2124        assert!(!is_error(&resp), "edge_history must not error: {resp}");
2125        let result = tool_text(&resp);
2126
2127        // Must carry horizon metadata.
2128        let total = result["total_commits"].as_u64().expect("total_commits");
2129        assert!(total > 0, "total_commits must be > 0 after ingest + rule");
2130
2131        // Must have at least one event (the SIMILAR derived-edge addition).
2132        let events = result["events"].as_array().expect("events array");
2133        assert!(!events.is_empty(), "expected at least one edge event");
2134
2135        // At least one event must be Added with a non-null rule (derived edge).
2136        let derived_added = events
2137            .iter()
2138            .any(|ev| ev["event"].as_str() == Some("Added") && !ev["rule"].is_null());
2139        assert!(
2140            derived_added,
2141            "expected a derived Added event with rule attribution: {events:?}"
2142        );
2143    }
2144
2145    /// `was_linked` must return `true` for an edge that was active at the given commit,
2146    /// and the response must include the echo fields.
2147    #[test]
2148    fn test_was_linked_at_valid_commit() {
2149        let db = SharedDb::open(&tmp_dir()).expect("open");
2150        {
2151            let mut g = db.write();
2152            let opts = IngestOptions {
2153                key_field: "id".into(),
2154                auto_fk: AutoFk::Off,
2155            };
2156            let rows: Vec<BTreeMap<String, Value>> = vec![
2157                [("id", Value::Str("x".into()))]
2158                    .into_iter()
2159                    .map(|(k, v)| (k.to_string(), v))
2160                    .collect(),
2161                [("id", Value::Str("y".into()))]
2162                    .into_iter()
2163                    .map(|(k, v)| (k.to_string(), v))
2164                    .collect(),
2165            ];
2166            g.ingest("N", rows, &opts).expect("ingest");
2167            g.insert_edge("LINK", "x", "y").expect("edge");
2168        }
2169        // There are now at least 2 commits (ingest + edge). Check at the last one.
2170        let g = db.read();
2171        let total = g.wal_total_commits().expect("wal_total_commits");
2172        drop(g);
2173
2174        let resp = tool_call(
2175            &db,
2176            1,
2177            "was_linked",
2178            json!({ "a": "x", "b": "y", "edge_type": "LINK", "at_commit": total - 1 }),
2179        );
2180        assert!(!is_error(&resp), "was_linked must not error: {resp}");
2181        let result = tool_text(&resp);
2182        assert_eq!(result["linked"], true);
2183        assert_eq!(result["a"], "x");
2184        assert_eq!(result["edge_type"], "LINK");
2185    }
2186
2187    /// `was_linked` with an out-of-horizon commit must return a tool error (not
2188    /// a protocol error), and the error message must mention the commit range.
2189    #[test]
2190    fn test_was_linked_out_of_horizon_returns_tool_error() {
2191        let db = SharedDb::open(&tmp_dir()).expect("open");
2192        {
2193            let mut g = db.write();
2194            g.insert_node("N", "a", vec![]).expect("node a");
2195            g.insert_node("N", "b", vec![]).expect("node b");
2196        }
2197        // Commit 999 is well beyond the WAL.
2198        let resp = tool_call(
2199            &db,
2200            1,
2201            "was_linked",
2202            json!({ "a": "a", "b": "b", "edge_type": "X", "at_commit": 999 }),
2203        );
2204        // isError true = tool-level error (not a JSON-RPC protocol error).
2205        assert!(
2206            is_error(&resp),
2207            "out-of-range commit must be a tool error: {resp}"
2208        );
2209        let text = resp["result"]["content"][0]["text"].as_str().expect("text");
2210        assert!(
2211            text.contains("out of range") || text.contains("range"),
2212            "error must mention range: {text}"
2213        );
2214    }
2215
2216    /// `node_history` tool must return the node's WAL history and the
2217    /// `total_commits` horizon field.
2218    #[test]
2219    fn test_node_history_via_mcp() {
2220        let db = demo_db(); // alice + bob, with a SIMILAR rule
2221        let resp = tool_call(&db, 1, "node_history", json!({ "key": "alice" }));
2222        assert!(!is_error(&resp), "node_history must not error: {resp}");
2223        let result = tool_text(&resp);
2224
2225        assert_eq!(result["key"], "alice");
2226        let total = result["total_commits"].as_u64().expect("total_commits");
2227        assert!(total > 0, "total_commits must be > 0");
2228
2229        let history = result["history"].as_array().expect("history array");
2230        assert!(
2231            !history.is_empty(),
2232            "alice should have at least one history entry"
2233        );
2234
2235        // First event should be a NodeInserted.
2236        let first_change = &history[0]["change"];
2237        assert_eq!(first_change["type"], "NodeInserted");
2238        assert_eq!(first_change["label"], "Person");
2239    }
2240}