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, namespace_arg, node_history_json, node_info_json, params_from_json,
53    parse_ingest_edges, result_set_json, rule_def_from_json, stamp_namespace, stamp_namespace_row,
54};
55use core_api::{
56    json_to_rows, json_to_value, AsOfScope, AutoFk, GraphError, IngestOptions, MaskMode, NodeMask,
57    SharedDb, Value, NS_PROP,
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, args),
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    // The second visibility axis. `namespace` is not a third way to say what
277    // `role` and `mask` say — it is a leg that **intersects** whichever of them
278    // is present (and stands alone when neither is), so it can only narrow what
279    // they already allow. A role bound to namespaces honours them with no
280    // argument here; passing one outside the binding is the empty intersection,
281    // never the union.
282    let namespace = match namespace_arg(args.get("namespace")) {
283        Ok(n) => n,
284        Err(e) => return CallOutcome::ToolErr(e),
285    };
286
287    // Optional time travel: a 0-based WAL commit index. The graph is read as
288    // of that commit; a `role` is still the role the store defines now, since
289    // `roles.json` is a sidecar and is never a WAL record.
290    let as_of = match args.get("as_of") {
291        None | Some(Js::Null) => None,
292        Some(v) => match v.as_u64() {
293            Some(n) => Some(n),
294            None => {
295                return CallOutcome::ToolErr(
296                    "as_of must be a non-negative integer commit index".into(),
297                )
298            }
299        },
300    };
301
302    if let Some(commit) = as_of {
303        // Stub mode discloses node existence, which is exactly the question an
304        // as-of read is asking. The two do not compose.
305        if args
306            .get("stub_hidden")
307            .and_then(|v| v.as_bool())
308            .unwrap_or(false)
309        {
310            return CallOutcome::ToolErr(
311                "as_of (time-travel) does not compose with stub_hidden".into(),
312            );
313        }
314        let scope = match (role, &mask_keys) {
315            (Some(role), _) => AsOfScope::Role(role),
316            (None, Some(keys)) => AsOfScope::Keys(keys),
317            (None, None) => match namespace.as_deref() {
318                // A namespace alone is its own as-of scope.
319                Some(ns) => AsOfScope::Namespace(ns),
320                None => {
321                    return match db.read().query_at(commit, cypher, &params) {
322                        Ok(rs) => CallOutcome::ToolOk(result_set_json(&rs)),
323                        Err(e) => CallOutcome::ToolErr(graph_err_msg(e)),
324                    }
325                }
326            },
327        };
328        let g = db.read();
329        let out = match (namespace.as_deref(), role.is_some() || mask_keys.is_some()) {
330            // Both legs: the namespace intersects the scope at that commit.
331            (Some(ns), true) => g.query_at_scoped_in_namespace(commit, cypher, &params, scope, ns),
332            _ => g.query_at_scoped(commit, cypher, &params, scope),
333        };
334        return match out {
335            Ok(rs) => CallOutcome::ToolOk(result_set_json(&rs)),
336            Err(GraphError::KeyNotFound { key }) if key.starts_with("role:") => {
337                CallOutcome::ToolErr(format!("unknown role '{}'", &key["role:".len()..]))
338            }
339            Err(e) => CallOutcome::ToolErr(graph_err_msg(e)),
340        };
341    }
342
343    if role.is_some() || mask_keys.is_some() || namespace.is_some() {
344        let stub_hidden = args
345            .get("stub_hidden")
346            .and_then(|v| v.as_bool())
347            .unwrap_or(false);
348        let g = db.read();
349        let mask = match (role, &mask_keys) {
350            (Some(role), _) => match g.mask_for_role(role) {
351                Ok(m) => m,
352                // The one error a caller can fix by rereading `roles.json`,
353                // told apart from a store whose roles never loaded at all.
354                Err(GraphError::KeyNotFound { .. }) => {
355                    return CallOutcome::ToolErr(format!("unknown role '{role}'"))
356                }
357                Err(e) => return CallOutcome::ToolErr(graph_err_msg(e)),
358            },
359            (None, Some(keys)) => NodeMask::from_keys(&*g, keys.iter().map(String::as_str)),
360            // A namespace alone: the namespace leg is the whole mask.
361            (None, None) => g.mask_for_namespace(
362                namespace
363                    .as_deref()
364                    .expect("one of the three is Some in this branch"),
365            ),
366        };
367        // With a role or a client mask present, the namespace is a second leg
368        // intersected into it — the same `NodeMask::intersect` the
369        // role-plus-client-mask path uses, so never-widen holds by construction.
370        let mask = match (namespace.as_deref(), role.is_some() || mask_keys.is_some()) {
371            (Some(ns), true) => mask.intersect(&g.mask_for_namespace(ns)),
372            _ => mask,
373        };
374        let mask = if stub_hidden {
375            mask.with_mode(MaskMode::Stub)
376        } else {
377            mask
378        };
379        return match g.query_masked(cypher, &params, &mask) {
380            Ok(rs) => CallOutcome::ToolOk(result_set_json(&rs)),
381            Err(e) => CallOutcome::ToolErr(graph_err_msg(e)),
382        };
383    }
384
385    let is_write = match core_api::is_write_query(cypher) {
386        Ok(b) => b,
387        Err(e) => return CallOutcome::ToolErr(e),
388    };
389    let rs = if is_write {
390        let mut g = db.write();
391        g.query_write(cypher, &params)
392    } else {
393        let g = db.read();
394        g.query(cypher, &params)
395    };
396    match rs {
397        Ok(rs) => CallOutcome::ToolOk(result_set_json(&rs)),
398        Err(e) => CallOutcome::ToolErr(graph_err_msg(e)),
399    }
400}
401
402/// A `mask` argument as a key list. `Err` when it is anything but an array of
403/// strings — including `null`, which is a caller that meant to pass one.
404fn mask_key_list(mask: &Js) -> Result<Vec<String>, String> {
405    let arr = mask
406        .as_array()
407        .ok_or_else(|| "mask must be an array of strings".to_string())?;
408    arr.iter()
409        .map(|v| {
410            v.as_str()
411                .map(str::to_string)
412                .ok_or_else(|| "mask must be an array of strings".to_string())
413        })
414        .collect()
415}
416
417fn tool_ingest(db: &SharedDb, args: &Js) -> CallOutcome {
418    let Some(label) = args.get("label").and_then(Js::as_str) else {
419        return CallOutcome::ToolErr("missing label".into());
420    };
421    let Some(rows_json) = args.get("rows_json").and_then(Js::as_str) else {
422        return CallOutcome::ToolErr("missing rows_json".into());
423    };
424    let mut opts = IngestOptions::default();
425    if let Some(kf) = args.get("key_field") {
426        match kf.as_str() {
427            Some(s) => opts.key_field = s.to_string(),
428            None => return CallOutcome::ToolErr("key_field must be a string".into()),
429        }
430    }
431    if let Some(suf) = args.get("auto_fk_suffix") {
432        match suf.as_str() {
433            Some(s) => {
434                opts.auto_fk = AutoFk::Auto {
435                    suffix: s.to_string(),
436                }
437            }
438            None => return CallOutcome::ToolErr("auto_fk_suffix must be a string".into()),
439        }
440    }
441    let edges = match args.get("edges") {
442        None | Some(Js::Null) => Vec::new(),
443        Some(raw) => match parse_ingest_edges(raw) {
444            Ok(e) => e,
445            Err(e) => return CallOutcome::ToolErr(e),
446        },
447    };
448    let parsed: Js = match serde_json::from_str(rows_json) {
449        Ok(v) => v,
450        Err(e) => {
451            return CallOutcome::ToolErr(graph_err_msg(GraphError::IngestError {
452                detail: e.to_string(),
453            }))
454        }
455    };
456    let mut converted = match json_to_rows(&parsed) {
457        Ok(c) => c,
458        Err(e) => return CallOutcome::ToolErr(graph_err_msg(e)),
459    };
460    // `namespace` applies to every node this call creates.
461    let namespace = match namespace_arg(args.get("namespace")) {
462        Ok(n) => n,
463        Err(e) => return CallOutcome::ToolErr(e),
464    };
465    if let Err(e) = stamp_namespace(&mut converted.rows, namespace.as_deref()) {
466        return CallOutcome::ToolErr(e);
467    }
468    let taken = std::mem::take(&mut converted.rows);
469    let report = {
470        let mut g = db.write();
471        g.ingest_with_edges(label, taken, &opts, &edges)
472    };
473    match report.map(|r| converted.into_report(r)) {
474        Ok(r) => match serde_json::to_value(&r) {
475            Ok(v) => CallOutcome::ToolOk(v),
476            Err(e) => CallOutcome::ToolErr(e.to_string()),
477        },
478        Err(e) => CallOutcome::ToolErr(graph_err_msg(e)),
479    }
480}
481
482fn tool_create_rule(db: &SharedDb, args: &Js) -> CallOutcome {
483    let def = match rule_def_from_json(args.clone()) {
484        Ok(d) => d,
485        Err(e) => return CallOutcome::ToolErr(e),
486    };
487    let name = def.name.clone();
488    let res = {
489        let mut g = db.write();
490        g.create_rule(def)
491    };
492    if let Err(e) = res {
493        return CallOutcome::ToolErr(graph_err_msg(e));
494    }
495    // A rule over a corpus too large to index in one commit is installed but
496    // derives nothing yet. Saying "ok" there would tell the caller to go and
497    // query edges that do not exist, so report the build instead — `stats`
498    // carries the same progress under each rule's `building`.
499    let building = db
500        .read()
501        .builds_in_progress()
502        .into_iter()
503        .find(|b| b.rule == name);
504    match building {
505        Some(b) => CallOutcome::ToolOk(json!({
506            "ok": true,
507            "name": name,
508            "building": {"indexed": b.indexed, "total": b.total},
509            "note": format!(
510                "the vector index for {name:?} is still being built ({}/{} vectors); \
511                 this rule derives no edges until it finishes. Every write advances it, \
512                 and `mushroomdb build-index <db-dir>` finishes it now. Poll `stats` — \
513                 the rule's `building` field disappears when its edges are in.",
514                b.indexed, b.total
515            ),
516        })),
517        None => CallOutcome::ToolOk(json!({"ok": true, "name": name})),
518    }
519}
520
521fn tool_explain(db: &SharedDb, args: &Js) -> CallOutcome {
522    let Some(a) = args.get("a").and_then(Js::as_str).filter(|s| !s.is_empty()) else {
523        return CallOutcome::ToolErr("missing a".into());
524    };
525    let Some(b) = args.get("b").and_then(Js::as_str).filter(|s| !s.is_empty()) else {
526        return CallOutcome::ToolErr("missing b".into());
527    };
528    let out = {
529        let g = db.read();
530        g.explain(a, b)
531    };
532    match out {
533        Ok(v) => match serde_json::to_value(&v) {
534            Ok(j) => CallOutcome::ToolOk(j),
535            Err(e) => CallOutcome::ToolErr(e.to_string()),
536        },
537        Err(e) => CallOutcome::ToolErr(graph_err_msg(e)),
538    }
539}
540
541/// `stats`, with the namespace roster narrowed when the caller names a role or a
542/// namespace.
543///
544/// The roster is the one part of `stats` that is a list of *other tenants*:
545/// every namespace and its live count. A caller answering as a role should be
546/// told about its own namespaces and no others, so `role` narrows the roster to
547/// the role's binding and `namespace` to that one name. The store-wide counts
548/// beside it are unchanged — they were never per-namespace and narrowing them
549/// would make the two halves of one body disagree.
550fn tool_stats(db: &SharedDb, args: &Js) -> CallOutcome {
551    let role = match args.get("role") {
552        None | Some(Js::Null) => None,
553        Some(Js::String(s)) if !s.is_empty() => Some(s.clone()),
554        Some(_) => return CallOutcome::ToolErr("role must be a non-empty string".into()),
555    };
556    let namespace = match namespace_arg(args.get("namespace")) {
557        Ok(n) => n,
558        Err(e) => return CallOutcome::ToolErr(e),
559    };
560    let (snap, role_def) = {
561        let g = db.read();
562        let def = match &role {
563            Some(r) => {
564                // Resolve through the same resolver `query` uses, so a store
565                // whose `roles.json` was corrupt at open says so here too
566                // instead of reporting the role simply unknown — one answer per
567                // cause, the same one on both tools. The mask is memoised per
568                // (role, commit_seq), so asking costs nothing a `query` with the
569                // same role would not already have paid.
570                if let Err(e) = g.mask_for_role(r) {
571                    return match e {
572                        GraphError::KeyNotFound { .. } => {
573                            CallOutcome::ToolErr(format!("unknown role '{r}'"))
574                        }
575                        other => CallOutcome::ToolErr(graph_err_msg(other)),
576                    };
577                }
578                g.roles().into_iter().find(|d| &d.name == r)
579            }
580            None => None,
581        };
582        (g.stats(), def)
583    };
584    let mut snap = snap;
585    if role_def.is_some() || namespace.is_some() {
586        snap.namespaces.retain(|n| {
587            role_def.as_ref().is_none_or(|d| d.sees_namespace(&n.name))
588                && namespace.as_deref().is_none_or(|ns| ns == n.name)
589        });
590    }
591    match serde_json::to_value(&snap) {
592        Ok(v) => CallOutcome::ToolOk(v),
593        Err(e) => CallOutcome::ToolErr(e.to_string()),
594    }
595}
596
597fn tool_node_info(db: &SharedDb, args: &Js) -> CallOutcome {
598    let Some(key) = args.get("key").and_then(Js::as_str) else {
599        return CallOutcome::ToolErr("missing key".into());
600    };
601    let info = {
602        let g = db.read();
603        g.node_info(key)
604    };
605    match info {
606        Some(info) => CallOutcome::ToolOk(node_info_json(&info)),
607        None => CallOutcome::ToolErr(graph_err_msg(GraphError::KeyNotFound {
608            key: key.to_string(),
609        })),
610    }
611}
612
613/// Insert a new node or update an existing node's properties, keyed by `key`.
614///
615/// If the node exists: each prop in `props` is written via `set_prop`.
616/// If the node does not exist: `label` is required; the node is ingested with
617/// `key_field = "id"` and the supplied props.
618///
619/// `namespace` is the namespace a node this call **creates** is created in. On a
620/// node that already exists it is written like any other property, which is what
621/// makes naming the namespace the node is already in a no-op and naming another
622/// one the engine's `NamespaceImmutable` refusal — one rule, stated once, in the
623/// place that owns it. A no-op `ns` writes nothing and is not counted in
624/// `updated_fields`, because nothing was updated.
625///
626/// `id` in `props` is **dropped on both paths**: it is the node's key. The create
627/// path stores `id` from `key` (it ingests with `key_field: "id"`), and
628/// `rename_node` is the only way to change it. One row builder now serves the
629/// create and the update path, so the rule is the same on both — before v0.6.6 the
630/// update path wrote `props.id` straight through `set_prop`, which could leave a
631/// stored `id` disagreeing with the key the node is reached by, while the create
632/// path had always ignored it.
633///
634/// Returns `{ok, key, created, updated_fields?}`.
635fn tool_upsert_entity(db: &SharedDb, args: &Js) -> CallOutcome {
636    let Some(key) = args.get("key").and_then(Js::as_str) else {
637        return CallOutcome::ToolErr("missing key".into());
638    };
639    let label_opt = args.get("label").and_then(Js::as_str);
640    let Some(props_obj) = args.get("props").and_then(Js::as_object) else {
641        return CallOutcome::ToolErr("missing props".into());
642    };
643    let namespace = match namespace_arg(args.get("namespace")) {
644        Ok(n) => n,
645        Err(e) => return CallOutcome::ToolErr(e),
646    };
647
648    // One row, stamped with the namespace, whichever path takes it: the
649    // conflict rule between an explicit `props.ns` and `namespace` is then the
650    // same one `ingest_json` applies.
651    let mut row: BTreeMap<String, Value> = BTreeMap::new();
652    for (field, json_val) in props_obj {
653        if field == "id" {
654            continue;
655        }
656        match json_to_value(json_val.clone()) {
657            Some(v) => {
658                row.insert(field.clone(), v);
659            }
660            None => {
661                return CallOutcome::ToolErr(format!("prop {field} is not a supported value type"))
662            }
663        }
664    }
665    if let Some(ns) = namespace.as_deref() {
666        if let Err(e) = stamp_namespace_row(&mut row, ns) {
667            return CallOutcome::ToolErr(e);
668        }
669    }
670
671    let exists = {
672        let g = db.read();
673        g.has_node(key)
674    };
675
676    if exists {
677        let mut g = db.write();
678        let mut count = 0usize;
679        for (field, v) in &row {
680            // The namespace a node is already in is the engine's no-op: it
681            // writes no record and takes no commit, so counting it as an updated
682            // field would report an update that did not happen. Asking first
683            // also keeps the refusal for a *different* namespace coming from the
684            // engine rather than from a second rule stated here.
685            if field == NS_PROP && Some(v) == g.namespace_of(key).map(Value::Str).as_ref() {
686                continue;
687            }
688            if let Err(e) = g.set_prop(key, field, v.clone()) {
689                return CallOutcome::ToolErr(graph_err_msg(e));
690            }
691            count += 1;
692        }
693        CallOutcome::ToolOk(json!({
694            "ok": true,
695            "key": key,
696            "created": false,
697            "updated_fields": count
698        }))
699    } else {
700        let Some(label) = label_opt else {
701            return CallOutcome::ToolErr("label required when creating a new entity".into());
702        };
703        let mut row = row;
704        row.insert("id".to_string(), Value::Str(key.to_string()));
705        let opts = IngestOptions {
706            key_field: "id".to_string(),
707            auto_fk: AutoFk::Off,
708        };
709        let mut g = db.write();
710        match g.ingest(label, vec![row], &opts) {
711            Ok(_) => CallOutcome::ToolOk(json!({ "ok": true, "key": key, "created": true })),
712            Err(e) => CallOutcome::ToolErr(graph_err_msg(e)),
713        }
714    }
715}
716
717/// Return neighbors connected by a given edge type (default `"SIMILAR"`).
718///
719/// Results are read from edges already materialized by a derivation rule
720/// (e.g. a `VectorSimilar` rule). Without a matching rule the returned list
721/// is empty — no live cosine computation is performed here.
722/// Returns up to `limit` (default 10) neighbor entries.
723fn tool_find_similar(db: &SharedDb, args: &Js) -> CallOutcome {
724    // Parse the optional mask once — it applies to both vector and edge paths.
725    // An invalid mask value (non-array or non-string element) fails closed.
726    let mask_keys: Option<Vec<String>> = if let Some(mask_val) = args.get("mask") {
727        match mask_val.as_array() {
728            Some(arr) => {
729                let mut ks: Vec<String> = Vec::with_capacity(arr.len());
730                for v in arr {
731                    match v.as_str() {
732                        Some(s) => ks.push(s.to_string()),
733                        None => {
734                            return CallOutcome::ToolErr("mask must be an array of strings".into())
735                        }
736                    }
737                }
738                Some(ks)
739            }
740            None => return CallOutcome::ToolErr("mask must be an array of strings".into()),
741        }
742    } else {
743        None
744    };
745
746    // When a `vector` array is provided, use the HNSW / brute-force vector
747    // similarity path instead of looking up pre-derived edges.
748    if let Some(vec_js) = args.get("vector").and_then(Js::as_array) {
749        let q: Vec<f64> = vec_js.iter().filter_map(|v| v.as_f64()).collect();
750        if q.is_empty() {
751            return CallOutcome::ToolErr("vector must be a non-empty array of numbers".into());
752        }
753        let field = args
754            .get("field")
755            .and_then(Js::as_str)
756            .unwrap_or("embedding");
757        let label_str = args.get("label").and_then(Js::as_str).unwrap_or("");
758        let label = if label_str.is_empty() {
759            None
760        } else {
761            Some(label_str)
762        };
763        let k = args
764            .get("k")
765            .and_then(Js::as_u64)
766            .map(|n| n as usize)
767            .unwrap_or(10);
768        let min = args.get("min").and_then(Js::as_f64).unwrap_or(0.8);
769
770        let hits = {
771            let g = db.read();
772            if let Some(ref keys) = mask_keys {
773                let node_mask = NodeMask::from_keys(&*g, keys.iter().map(String::as_str));
774                g.find_similar_vector_masked(field, label, &q, k, min, &node_mask)
775            } else {
776                g.find_similar_vector(field, label, &q, k, min)
777            }
778        };
779        let results: Vec<Js> = hits
780            .into_iter()
781            .map(|(key, score)| json!({ "key": key, "score": score }))
782            .collect();
783        return CallOutcome::ToolOk(json!({
784            "mode": "vector",
785            "field": field,
786            "label": label,
787            "k": k,
788            "min": min,
789            "results": results
790        }));
791    }
792
793    // Edge-traversal path: return neighbors connected by the given edge type.
794    let Some(key) = args.get("key").and_then(Js::as_str) else {
795        return CallOutcome::ToolErr("missing key (or provide vector for vector search)".into());
796    };
797    let edge_type = args
798        .get("edge_type")
799        .and_then(Js::as_str)
800        .unwrap_or("SIMILAR");
801    let limit = args
802        .get("limit")
803        .and_then(Js::as_u64)
804        .map(|n| n as usize)
805        .unwrap_or(10);
806
807    // When a mask is present, a hidden query key behaves identically to a
808    // nonexistent key — we do not confirm its existence.
809    if let Some(ref mask) = mask_keys {
810        let mask_set: std::collections::HashSet<&str> = mask.iter().map(String::as_str).collect();
811        if !mask_set.contains(key) {
812            return CallOutcome::ToolErr(graph_err_msg(GraphError::KeyNotFound {
813                key: key.into(),
814            }));
815        }
816        let out = {
817            let g = db.read();
818            g.node_edges(key)
819        };
820        return match out {
821            Ok(edges) => {
822                let similar: Vec<Js> = edges
823                    .iter()
824                    .filter(|e| e.edge_type == edge_type)
825                    .filter(|e| {
826                        // Keep only edges where the neighbor is also visible.
827                        let neighbor_key = if e.src_key == key {
828                            &e.dst_key
829                        } else {
830                            &e.src_key
831                        };
832                        mask_set.contains(neighbor_key.as_str())
833                    })
834                    .take(limit)
835                    .map(|e| {
836                        let neighbor_key = if e.src_key == key {
837                            &e.dst_key
838                        } else {
839                            &e.src_key
840                        };
841                        let direction = if e.src_key == key { "out" } else { "in" };
842                        json!({
843                            "neighbor_key": neighbor_key,
844                            "direction": direction,
845                            "edge_type": e.edge_type,
846                            "derived": e.derived,
847                        })
848                    })
849                    .collect();
850                CallOutcome::ToolOk(json!({
851                    "key": key,
852                    "edge_type": edge_type,
853                    "similar": similar
854                }))
855            }
856            Err(e) => CallOutcome::ToolErr(graph_err_msg(e)),
857        };
858    }
859
860    let out = {
861        let g = db.read();
862        g.node_edges(key)
863    };
864    match out {
865        Ok(edges) => {
866            let similar: Vec<Js> = edges
867                .iter()
868                .filter(|e| e.edge_type == edge_type)
869                .take(limit)
870                .map(|e| {
871                    let neighbor_key = if e.src_key == key {
872                        &e.dst_key
873                    } else {
874                        &e.src_key
875                    };
876                    let direction = if e.src_key == key { "out" } else { "in" };
877                    json!({
878                        "neighbor_key": neighbor_key,
879                        "direction": direction,
880                        "edge_type": e.edge_type,
881                        "derived": e.derived,
882                    })
883                })
884                .collect();
885            CallOutcome::ToolOk(json!({
886                "key": key,
887                "edge_type": edge_type,
888                "similar": similar
889            }))
890        }
891        Err(e) => CallOutcome::ToolErr(graph_err_msg(e)),
892    }
893}
894
895fn tool_hybrid_search(db: &SharedDb, args: &Js) -> CallOutcome {
896    let Some(query_text) = args.get("query_text").and_then(Js::as_str) else {
897        return CallOutcome::ToolErr("missing required field: query_text".into());
898    };
899    let Some(text_field) = args.get("text_field").and_then(Js::as_str) else {
900        return CallOutcome::ToolErr("missing required field: text_field".into());
901    };
902
903    let vector_field = args
904        .get("vector_field")
905        .and_then(Js::as_str)
906        .unwrap_or("embedding");
907    let label = args.get("label").and_then(Js::as_str);
908    let k = args
909        .get("k")
910        .and_then(Js::as_u64)
911        .map(|n| n as usize)
912        .unwrap_or(10);
913
914    let query_vec: Vec<f64> = args
915        .get("vector")
916        .and_then(Js::as_array)
917        .map(|arr| arr.iter().filter_map(|v| v.as_f64()).collect())
918        .unwrap_or_default();
919
920    let hits = {
921        let g = db.read();
922        g.search_hybrid(text_field, query_text, vector_field, &query_vec, label, k)
923    };
924
925    let results: Vec<Js> = hits
926        .into_iter()
927        .map(|(key, score)| json!({ "key": key, "score": score }))
928        .collect();
929
930    CallOutcome::ToolOk(json!({
931        "query_text": query_text,
932        "text_field": text_field,
933        "vector_field": vector_field,
934        "label": label,
935        "k": k,
936        "results": results
937    }))
938}
939
940fn tool_node_history(db: &SharedDb, args: &Js) -> CallOutcome {
941    let Some(key) = args.get("key").and_then(Js::as_str) else {
942        return CallOutcome::ToolErr("missing key".into());
943    };
944    let g = db.read();
945    let result = match g.node_history(key) {
946        Ok(e) => e,
947        Err(e) => return CallOutcome::ToolErr(graph_err_msg(e)),
948    };
949    CallOutcome::ToolOk(node_history_json(key, &result))
950}
951
952fn tool_edge_history(db: &SharedDb, args: &Js) -> CallOutcome {
953    let Some(a) = args.get("a").and_then(Js::as_str).filter(|s| !s.is_empty()) else {
954        return CallOutcome::ToolErr("missing a".into());
955    };
956    let Some(b) = args.get("b").and_then(Js::as_str).filter(|s| !s.is_empty()) else {
957        return CallOutcome::ToolErr("missing b".into());
958    };
959    let result = {
960        let g = db.read();
961        g.edge_history(a, b)
962    };
963    match result {
964        Ok(hr) => CallOutcome::ToolOk(edge_history_result_json(a, b, &hr)),
965        Err(e) => CallOutcome::ToolErr(graph_err_msg(e)),
966    }
967}
968
969fn tool_was_linked(db: &SharedDb, args: &Js) -> CallOutcome {
970    let Some(a) = args.get("a").and_then(Js::as_str).filter(|s| !s.is_empty()) else {
971        return CallOutcome::ToolErr("missing a".into());
972    };
973    let Some(b) = args.get("b").and_then(Js::as_str).filter(|s| !s.is_empty()) else {
974        return CallOutcome::ToolErr("missing b".into());
975    };
976    let Some(edge_type) = args
977        .get("edge_type")
978        .and_then(Js::as_str)
979        .filter(|s| !s.is_empty())
980    else {
981        return CallOutcome::ToolErr("missing edge_type".into());
982    };
983    let at_commit = match args.get("at_commit").and_then(Js::as_u64) {
984        Some(n) => n,
985        None => return CallOutcome::ToolErr("missing or invalid at_commit".into()),
986    };
987    let result = {
988        let g = db.read();
989        g.was_linked(a, b, edge_type, at_commit)
990    };
991    match result {
992        Ok(linked) => CallOutcome::ToolOk(json!({
993            "a": a,
994            "b": b,
995            "edge_type": edge_type,
996            "at_commit": at_commit,
997            "linked": linked,
998        })),
999        Err(e) => CallOutcome::ToolErr(graph_err_msg(e)),
1000    }
1001}
1002
1003fn tool_rename_node(db: &SharedDb, args: &Js) -> CallOutcome {
1004    let Some(old_key) = args.get("old_key").and_then(Js::as_str) else {
1005        return CallOutcome::ToolErr("missing old_key".into());
1006    };
1007    let Some(new_key) = args.get("new_key").and_then(Js::as_str) else {
1008        return CallOutcome::ToolErr("missing new_key".into());
1009    };
1010    let mut g = db.write();
1011    match g.rename_node(old_key, new_key) {
1012        Ok(()) => CallOutcome::ToolOk(json!({
1013            "ok": true,
1014            "old_key": old_key,
1015            "new_key": new_key,
1016        })),
1017        Err(e) => CallOutcome::ToolErr(graph_err_msg(e)),
1018    }
1019}
1020
1021pub(crate) fn graph_err_msg(e: GraphError) -> String {
1022    match e {
1023        GraphError::QueryError { detail } | GraphError::IngestError { detail } => detail,
1024        other => other.to_string(),
1025    }
1026}
1027
1028fn initialize_result() -> Js {
1029    json!({
1030        "protocolVersion": "2024-11-05",
1031        "capabilities": { "tools": {} },
1032        "serverInfo": { "name": "mushroomdb", "version": env!("CARGO_PKG_VERSION") }
1033    })
1034}
1035
1036/// The prefix every graph tool's description carries.
1037///
1038/// A host that ranks tools by their description now has one signal that the
1039/// repository task tools are the ones to reach for first, and that everything
1040/// under this prefix is the lower-level surface beneath them.
1041const ADVANCED_PREFIX: &str = "Advanced: ";
1042
1043/// The three a code-graph store advertises: one tool to find, one to ask an
1044/// arbitrary question, one to size the store.
1045///
1046/// `ingest_json` is not among them. A store built by `ingest-git` is written by
1047/// `sync` and `touch`, not by an assistant bulk-loading rows into it, and the
1048/// tool that is never the right one on this surface is the one worth not
1049/// listing.
1050pub const CODE_GRAPH_TOOLS: [&str; 3] = ["explore", "query", "stats"];
1051
1052/// The fifteen a memory store advertises, in the order it lists them.
1053///
1054/// A store with no repository in it used to be handed the code door's own task
1055/// tools — `map`, `context`, `impact`, `owners`, `why`, `sync` — which answer
1056/// from a code graph there is none of, plus `ingest_json`. Six of the eleven
1057/// names an assistant found answered from a repository the store did not
1058/// hold. These are
1059/// the questions an entity graph *can* answer: what is there (`query` — now
1060/// with a `role`), why two things are associated, what is around a node, what
1061/// it is and what it is joined to, whether a link held at a commit and when it
1062/// changed, what is like it, and what was written down about it.
1063///
1064/// Listing order is ranking: a host that defers schemas shows this list in
1065/// order, so the two questions this door exists for come first.
1066///
1067/// `edges_at` and `what_if` are the two the first association benchmark run
1068/// showed missing: a run asked what a node's relationships were at a past
1069/// commit and spent twenty to sixty-seven turns replaying `edge_history` for
1070/// it, and had no way at all to ask what a change would do. They sit after
1071/// `was_linked`, which is the narrowest form of the same time question.
1072///
1073/// The code task tools stay served on a memory store, as these stay served on
1074/// a code-graph one — [`tools_list`] decides what is *advertised*, never what
1075/// is answered.
1076pub const ASSOCIATION_TOOLS: [&str; 15] = [
1077    "query",
1078    "explain_association",
1079    "neighborhood",
1080    "node_info",
1081    "node_edges",
1082    "was_linked",
1083    "edges_at",
1084    "what_if",
1085    "node_history",
1086    "edge_history",
1087    "find_similar",
1088    "hybrid_search",
1089    "remember",
1090    "recall",
1091    "stats",
1092];
1093
1094/// Which door a store is: which default tool list it gets.
1095///
1096/// Decided from the store the server opened, once, at startup — not from an
1097/// install flag — so one `.mcp.json` serves both and neither has to be
1098/// configured for.
1099#[derive(Clone, Copy, Debug, PartialEq, Eq)]
1100pub(crate) enum Surface {
1101    /// A repository was ingested into this store: the `GitSync` marker is
1102    /// there, and `explore` has a code graph to explore.
1103    CodeGraph,
1104    /// Any other store, including an empty one: the fifteen-tool association
1105    /// surface, where `explore` would have nothing to answer from.
1106    Memory,
1107}
1108
1109impl Surface {
1110    /// The tools a default `tools/list` on this surface advertises, in the
1111    /// order it advertises them.
1112    fn listing(self) -> &'static [&'static str] {
1113        match self {
1114            Surface::CodeGraph => &CODE_GRAPH_TOOLS,
1115            Surface::Memory => &ASSOCIATION_TOOLS,
1116        }
1117    }
1118}
1119
1120/// The surface the store `db` holds: [`Surface::CodeGraph`] when it carries the
1121/// `GitSync` marker `ingest-git` writes, [`Surface::Memory`] otherwise.
1122fn surface_of(db: &SharedDb) -> Surface {
1123    let ingested = {
1124        let g = db.read();
1125        g.has_node(crate::mcp_tasks::SYNC_KEY)
1126    };
1127    if ingested {
1128        Surface::CodeGraph
1129    } else {
1130        Surface::Memory
1131    }
1132}
1133
1134/// The tools `tools/list` advertises: the fourteen task tools, then the
1135/// graph tools with their descriptions prefixed.
1136///
1137/// `all` false — the default — lists what `surface` names, **in the order that
1138/// surface names it**: three on a code graph, fifteen on a memory store. The
1139/// order is the point. A host that defers tool schemas makes a model search
1140/// for them, and the list it searches is read top-down, so each surface ranks
1141/// its own tools rather than inheriting the task-tools-then-graph-tools order
1142/// that only the code door has a reason for.
1143///
1144/// `all` true lists all twenty-seven in that established order whichever store
1145/// this is, which is what `mushroomdb mcp --all-tools` runs and what the
1146/// published server card documents: a caller that asked for everything asked
1147/// for the whole surface, not for one door's ranking of it.
1148///
1149/// Either way every tool stays callable: the flag and the surface decide what
1150/// is advertised, not what is served.
1151fn tools_list(all: bool, surface: Surface) -> Js {
1152    let mut served: Vec<Js> = crate::mcp_tasks::task_tools();
1153    for mut tool in graph_tools() {
1154        if let Some(d) = tool.get("description").and_then(Js::as_str) {
1155            let prefixed = format!("{ADVANCED_PREFIX}{d}");
1156            tool["description"] = Js::String(prefixed);
1157        }
1158        served.push(tool);
1159    }
1160    if all {
1161        return json!({ "tools": served });
1162    }
1163    let listing = surface.listing();
1164    let mut tools: Vec<Js> = Vec::with_capacity(listing.len());
1165    for name in listing {
1166        let Some(tool) = served
1167            .iter()
1168            .find(|t| t.get("name").and_then(Js::as_str) == Some(*name))
1169        else {
1170            debug_assert!(false, "{surface:?} lists {name}, which is not served");
1171            continue;
1172        };
1173        tools.push(tool.clone());
1174    }
1175    json!({ "tools": tools })
1176}
1177
1178/// The thirteen graph tools, in the order they have always been listed, with
1179/// their descriptions unprefixed. [`tools_list`] adds the prefix.
1180fn graph_tools() -> Vec<Js> {
1181    let Js::Array(tools) = json!([
1182            {
1183                "name": "query",
1184                "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. Pass 'as_of' to answer from a past commit; it composes with 'role' or with 'mask'. Pass 'namespace' to answer from one namespace only. 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.",
1185                "inputSchema": {
1186                    "type": "object",
1187                    "properties": {
1188                        "cypher": { "type": "string", "description": "Cypher query text." },
1189                        "params": {
1190                            "type": "object",
1191                            "description": "Named JSON-scalar query parameters."
1192                        },
1193                        "mask": {
1194                            "type": "array",
1195                            "items": { "type": "string" },
1196                            "description": "Optional node key allow-list. When present, only these nodes are visible; write statements are rejected."
1197                        },
1198                        "role": {
1199                            "type": "string",
1200                            "description": "Answer as this role from the store's roles: only the nodes it may see. A role may also be narrowed by one property test (`status in [...]`), declared in the store's roles."
1201                        },
1202                        "as_of": {
1203                            "type": "integer",
1204                            "minimum": 0,
1205                            "description": "0-based WAL commit index: answer from the graph as it was at that commit. Composes with 'role' or with 'mask' — never both, which is refused as it is without 'as_of' — and whichever is passed is resolved against the graph as it was then. Deleting a node does not remove it from a role's past, and a role's 'keys' resolve to whichever node held the key at that commit. Writes are refused."
1206                        },
1207                        "namespace": {
1208                            "type": "string",
1209                            "description": "Answer only from this namespace. Intersects with 'role' and 'mask' — it can only narrow what they already allow. A role bound to namespaces honours them with no argument here. 'default' is the namespace of every node that names none; a name no node uses answers with nothing."
1210                        }
1211                    },
1212                    "required": ["cypher"]
1213                }
1214            },
1215            {
1216                "name": "ingest_json",
1217                "description": "Ingest a JSON array of objects as nodes of one label.",
1218                "inputSchema": {
1219                    "type": "object",
1220                    "properties": {
1221                        "label": { "type": "string" },
1222                        "rows_json": {
1223                            "type": "string",
1224                            "description": "JSON text of an array of objects."
1225                        },
1226                        "key_field": { "type": "string" },
1227                        "auto_fk_suffix": { "type": "string" },
1228                        "edges": {
1229                            "type": "array",
1230                            "description": "Optional user edges [{edge_type, src, dst}]."
1231                        },
1232                        "namespace": {
1233                            "type": "string",
1234                            "description": "Namespace for every node this call creates. Omitted means the 'default' namespace. A row that carries its own 'ns' must name the same namespace. A namespace is set at insert and cannot be changed afterwards."
1235                        }
1236                    },
1237                    "required": ["label", "rows_json"]
1238                }
1239            },
1240            {
1241                "name": "create_rule",
1242                "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.",
1243                "inputSchema": {
1244                    "type": "object",
1245                    "properties": {
1246                        "name": { "type": "string" },
1247                        "src_label": { "type": "string" },
1248                        "dst_label": { "type": "string" },
1249                        "predicate": { "type": "object" },
1250                        "edge_type": { "type": "string" },
1251                        "weight_prop": {
1252                            "type": ["string", "null"],
1253                            "description": "Edge property that stores the score (default: weight)."
1254                        },
1255                        "max_edges": { "type": ["integer", "null"] },
1256                        "namespace": {
1257                            "type": "string",
1258                            "description": "Scope the rule to one namespace: it sees only that namespace's nodes — source, via hop and destination — so every edge it derives stays inside. Omitted means a global rule, which is the only kind that may derive an edge across a boundary."
1259                        }
1260                    },
1261                    "required": ["name", "src_label", "dst_label", "predicate", "edge_type"]
1262                }
1263            },
1264            {
1265                "name": "explain",
1266                "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.",
1267                "inputSchema": {
1268                    "type": "object",
1269                    "properties": {
1270                        "a": { "type": "string", "minLength": 1 },
1271                        "b": { "type": "string", "minLength": 1 }
1272                    },
1273                    "required": ["a", "b"]
1274                }
1275            },
1276            {
1277                "name": "stats",
1278                "description": "How big is this store — live node, edge and rule counts, plus `history_floor`, the oldest commit history still reaches (0 when nothing has been pruned), and `namespaces`, every namespace with at least one live node and its count. Pass 'role' or 'namespace' to be told about those namespaces only.",
1279                "inputSchema": {
1280                    "type": "object",
1281                    "properties": {
1282                        "role": {
1283                            "type": "string",
1284                            "description": "Report only the namespaces this role may see. The store-wide counts beside them are unchanged."
1285                        },
1286                        "namespace": {
1287                            "type": "string",
1288                            "description": "Report only this namespace. Intersects with 'role'."
1289                        }
1290                    }
1291                }
1292            },
1293            {
1294                "name": "node_info",
1295                "description": "What is K — its label and every property it holds.",
1296                "inputSchema": {
1297                    "type": "object",
1298                    "properties": {
1299                        "key": { "type": "string" }
1300                    },
1301                    "required": ["key"]
1302                }
1303            },
1304            {
1305                "name": "upsert_entity",
1306                "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. 'id' in 'props' is ignored on both paths: a created node stores 'id' as its key, and 'rename_node' is the only way to change it. Useful for agent memory: store or refresh an entity without checking existence first.",
1307                "inputSchema": {
1308                    "type": "object",
1309                    "properties": {
1310                        "key": { "type": "string", "description": "Unique node key." },
1311                        "label": { "type": "string", "description": "Node label (required when creating a new entity)." },
1312                        "props": {
1313                            "type": "object",
1314                            "description": "Properties to set. Values must be scalars (string, number, bool) or arrays of scalars."
1315                        },
1316                        "namespace": {
1317                            "type": "string",
1318                            "description": "Namespace for a node this call creates. Omitted means the 'default' namespace. On a node that already exists, naming the namespace it is in is a no-op and naming another one is refused — a namespace is set at insert and cannot be changed."
1319                        }
1320                    },
1321                    "required": ["key", "props"]
1322                }
1323            },
1324            {
1325                "name": "find_similar",
1326                "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.",
1327                "inputSchema": {
1328                    "type": "object",
1329                    "properties": {
1330                        "vector": {
1331                            "type": "array",
1332                            "items": { "type": "number" },
1333                            "description": "Query embedding vector for vector-similarity search. When present, vector-search mode is used and `key` is ignored."
1334                        },
1335                        "field": { "type": "string", "description": "Property field holding the embedding vectors (default: embedding). Used in vector-search mode." },
1336                        "label": { "type": "string", "description": "Restrict search to nodes with this label. Empty string means all labels. Used in vector-search mode." },
1337                        "k": { "type": "integer", "description": "Maximum results to return in vector-search mode (default: 10)." },
1338                        "min": { "type": "number", "description": "Minimum cosine similarity threshold in vector-search mode (default: 0.8)." },
1339                        "mask": {
1340                            "type": "array",
1341                            "items": { "type": "string" },
1342                            "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."
1343                        },
1344                        "key": { "type": "string", "description": "Source node key for edge-traversal mode." },
1345                        "edge_type": { "type": "string", "description": "Edge type to filter by in edge-traversal mode (default: SIMILAR)." },
1346                        "limit": { "type": "integer", "description": "Maximum neighbors to return in edge-traversal mode (default: 10)." }
1347                    }
1348                }
1349            },
1350            {
1351                "name": "hybrid_search",
1352                "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.",
1353                "inputSchema": {
1354                    "type": "object",
1355                    "properties": {
1356                        "query_text": { "type": "string", "description": "Fulltext query string." },
1357                        "text_field": { "type": "string", "description": "Property field to search with fulltext." },
1358                        "vector": {
1359                            "type": "array",
1360                            "items": { "type": "number" },
1361                            "description": "Query embedding vector. Omit for text-only ranking."
1362                        },
1363                        "vector_field": { "type": "string", "description": "Property field holding embedding vectors (default: embedding)." },
1364                        "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." },
1365                        "k": { "type": "integer", "description": "Maximum results to return (default: 10)." }
1366                    },
1367                    "required": ["query_text", "text_field"]
1368                }
1369            },
1370            {
1371                "name": "node_history",
1372                "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) and `horizon`, the oldest commit still retained; events before it are gone. History is WAL-scoped — pre-snapshot commits are not visible.",
1373                "inputSchema": {
1374                    "type": "object",
1375                    "properties": {
1376                        "key": { "type": "string", "description": "Node key to look up." }
1377                    },
1378                    "required": ["key"]
1379                }
1380            },
1381            {
1382                "name": "edge_history",
1383                "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) and `horizon`, the oldest commit still retained; events before it are gone.",
1384                "inputSchema": {
1385                    "type": "object",
1386                    "properties": {
1387                        "a": { "type": "string", "minLength": 1, "description": "First node key." },
1388                        "b": { "type": "string", "minLength": 1, "description": "Second node key." }
1389                    },
1390                    "required": ["a", "b"]
1391                }
1392            },
1393            {
1394                "name": "was_linked",
1395                "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 retained horizon (`horizon..total_commits`).",
1396                "inputSchema": {
1397                    "type": "object",
1398                    "properties": {
1399                        "a": { "type": "string", "minLength": 1, "description": "First node key." },
1400                        "b": { "type": "string", "minLength": 1, "description": "Second node key." },
1401                        "edge_type": { "type": "string", "minLength": 1, "description": "Edge type to check." },
1402                        "at_commit": { "type": "integer", "minimum": 0, "description": "0-based WAL commit index to query." }
1403                    },
1404                    "required": ["a", "b", "edge_type", "at_commit"]
1405                }
1406            },
1407            {
1408                "name": "rename_node",
1409                "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.",
1410                "inputSchema": {
1411                    "type": "object",
1412                    "properties": {
1413                        "old_key": { "type": "string", "minLength": 1, "description": "Current node key." },
1414                        "new_key": { "type": "string", "minLength": 1, "description": "Desired new node key." }
1415                    },
1416                    "required": ["old_key", "new_key"]
1417                }
1418            }
1419    ]) else {
1420        unreachable!("the literal above is an array")
1421    };
1422    tools
1423}
1424
1425fn tool_ok(payload: Js) -> Js {
1426    json!({
1427        "content": [{ "type": "text", "text": payload.to_string() }]
1428    })
1429}
1430
1431/// A task tool's result: one text block, and no `structuredContent`.
1432///
1433/// The report used to ride along beside the digest, repeating it verbatim
1434/// under a `text` key. Nothing bound it — no task tool declares an
1435/// `outputSchema` — and it tripled the size of every reply, so a caller that
1436/// wants the numbers now asks for them with `json: true` and gets the report
1437/// *as* the text.
1438fn task_ok(text: &str) -> Js {
1439    json!({
1440        "content": [{ "type": "text", "text": text }]
1441    })
1442}
1443
1444fn tool_err(message: &str) -> Js {
1445    json!({
1446        "content": [{ "type": "text", "text": message }],
1447        "isError": true
1448    })
1449}
1450
1451fn write_result(writer: &mut impl Write, id: Option<Js>, result: Js) -> io::Result<()> {
1452    write_json(
1453        writer,
1454        &json!({
1455            "jsonrpc": "2.0",
1456            "id": id.unwrap_or(Js::Null),
1457            "result": result
1458        }),
1459    )
1460}
1461
1462fn write_error(
1463    writer: &mut impl Write,
1464    id: Option<Js>,
1465    code: i64,
1466    message: &str,
1467) -> io::Result<()> {
1468    write_json(
1469        writer,
1470        &json!({
1471            "jsonrpc": "2.0",
1472            "id": id.unwrap_or(Js::Null),
1473            "error": { "code": code, "message": message }
1474        }),
1475    )
1476}
1477
1478fn write_json(writer: &mut impl Write, value: &Js) -> io::Result<()> {
1479    let s = serde_json::to_string(value).map_err(io::Error::other)?;
1480    writeln!(writer, "{s}")?;
1481    writer.flush()
1482}
1483
1484// ---------------------------------------------------------------------------
1485// Tests: MCP tool round-trips via stdio
1486// ---------------------------------------------------------------------------
1487
1488#[cfg(test)]
1489mod tests {
1490    use super::*;
1491    use core_api::{AutoFk, IngestOptions, Predicate, RuleDef, Value};
1492    use std::path::PathBuf;
1493    use std::sync::atomic::{AtomicU64, Ordering};
1494
1495    fn tmp_dir() -> PathBuf {
1496        static SEQ: AtomicU64 = AtomicU64::new(0);
1497        let n = SEQ.fetch_add(1, Ordering::Relaxed);
1498        let d = std::env::temp_dir().join(format!("mcp-test-{}-{}", std::process::id(), n));
1499        // These stores are never cleaned up, so a process id the OS hands out
1500        // again lands on a previous run's data and every assertion about counts
1501        // fails. `tests/mcp.rs::tmp` already clears its path for this reason.
1502        let _ = std::fs::remove_dir_all(&d);
1503        d
1504    }
1505
1506    /// Open a SharedDb with two Person nodes and one derived SIMILAR edge.
1507    fn demo_db() -> SharedDb {
1508        let db = SharedDb::open(&tmp_dir()).expect("open");
1509        {
1510            let mut g = db.write();
1511            let opts = IngestOptions {
1512                key_field: "id".into(),
1513                auto_fk: AutoFk::Off,
1514            };
1515            // Two people with identical embeddings → will fire SIMILAR rule.
1516            let people: Vec<BTreeMap<String, Value>> = vec![
1517                [
1518                    ("id", Value::Str("alice".into())),
1519                    ("name", Value::Str("Alice".into())),
1520                    (
1521                        "emb",
1522                        Value::List(vec![Value::Float(1.0), Value::Float(0.0)]),
1523                    ),
1524                ]
1525                .into_iter()
1526                .map(|(k, v)| (k.to_string(), v))
1527                .collect(),
1528                [
1529                    ("id", Value::Str("bob".into())),
1530                    ("name", Value::Str("Bob".into())),
1531                    (
1532                        "emb",
1533                        Value::List(vec![Value::Float(1.0), Value::Float(0.0)]),
1534                    ),
1535                ]
1536                .into_iter()
1537                .map(|(k, v)| (k.to_string(), v))
1538                .collect(),
1539            ];
1540            g.ingest("Person", people, &opts).expect("ingest");
1541
1542            // Rule: VectorSimilar on emb → SIMILAR edge (cosine(ident,ident)=1.0 ≥ 0.9).
1543            g.create_rule(RuleDef {
1544                name: "sim_emb".into(),
1545                src_label: "Person".into(),
1546                dst_label: "Person".into(),
1547                predicate: Predicate::VectorSimilar {
1548                    field: "emb".into(),
1549                    min: 0.9,
1550                },
1551                edge_type: "SIMILAR".into(),
1552                weight_prop: Some("score".into()),
1553                max_edges: None,
1554                approximate: false,
1555                via_label: None,
1556                via_edge: None,
1557                via_dir: None,
1558                namespace: None,
1559            })
1560            .expect("rule");
1561        }
1562        db
1563    }
1564
1565    fn roundtrip(db: &SharedDb, request: &str) -> Js {
1566        roundtrip_with(db, false, request)
1567    }
1568
1569    fn roundtrip_with(db: &SharedDb, all_tools: bool, request: &str) -> Js {
1570        let input = format!("{request}\n");
1571        let mut output = Vec::new();
1572        run_mcp_stdio_with(db.clone(), None, all_tools, input.as_bytes(), &mut output)
1573            .expect("mcp");
1574        let s = std::str::from_utf8(&output).expect("utf8");
1575        serde_json::from_str(s.trim()).expect("json response")
1576    }
1577
1578    fn tool_call(db: &SharedDb, id: u64, tool: &str, args: Js) -> Js {
1579        let req = json!({
1580            "jsonrpc": "2.0",
1581            "id": id,
1582            "method": "tools/call",
1583            "params": { "name": tool, "arguments": args }
1584        });
1585        roundtrip(db, &req.to_string())
1586    }
1587
1588    /// Unwrap the `text` field from a successful tool response.
1589    fn tool_text(resp: &Js) -> Js {
1590        let text = resp["result"]["content"][0]["text"]
1591            .as_str()
1592            .expect("content[0].text");
1593        serde_json::from_str(text).expect("tool text is json")
1594    }
1595
1596    fn is_error(resp: &Js) -> bool {
1597        resp["result"]["isError"].as_bool().unwrap_or(false)
1598    }
1599
1600    fn tool_err_text(resp: &Js) -> String {
1601        resp["result"]["content"][0]["text"]
1602            .as_str()
1603            .unwrap_or("")
1604            .to_string()
1605    }
1606
1607    // --- existing tools ---
1608
1609    #[test]
1610    fn test_tools_list_includes_all_expected() {
1611        let db = demo_db();
1612        let resp = roundtrip_with(
1613            &db,
1614            true,
1615            r#"{"jsonrpc":"2.0","id":1,"method":"tools/list"}"#,
1616        );
1617        let tools = resp["result"]["tools"].as_array().expect("tools array");
1618        let names: Vec<&str> = tools
1619            .iter()
1620            .map(|t| t["name"].as_str().expect("name"))
1621            .collect();
1622        for expected in &[
1623            // The fourteen task tools, first and in order.
1624            "explore",
1625            "map",
1626            "context",
1627            "impact",
1628            "owners",
1629            "why",
1630            "explain_association",
1631            "node_edges",
1632            "neighborhood",
1633            "edges_at",
1634            "what_if",
1635            "recall",
1636            "remember",
1637            "sync",
1638            // The thirteen graph tools.
1639            "query",
1640            "ingest_json",
1641            "create_rule",
1642            "explain",
1643            "stats",
1644            "node_info",
1645            "upsert_entity",
1646            "find_similar",
1647            "hybrid_search",
1648            "node_history",
1649            "edge_history",
1650            "was_linked",
1651            "rename_node",
1652        ] {
1653            assert!(names.contains(expected), "missing tool: {expected}");
1654        }
1655        assert_eq!(
1656            names.len(),
1657            27,
1658            "expected exactly 27 tools, got {}",
1659            names.len()
1660        );
1661        assert_eq!(
1662            &names[..14],
1663            [
1664                "explore",
1665                "map",
1666                "context",
1667                "impact",
1668                "owners",
1669                "why",
1670                "explain_association",
1671                "node_edges",
1672                "neighborhood",
1673                "edges_at",
1674                "what_if",
1675                "recall",
1676                "remember",
1677                "sync"
1678            ],
1679            "the task tools come first, in order"
1680        );
1681        assert_eq!(names[14], "query", "the graph tools follow them");
1682    }
1683
1684    /// Binding: on a store no repository was ingested into, the default
1685    /// listing is the fifteen association tools, in [`ASSOCIATION_TOOLS`]
1686    /// order, and nothing else.
1687    #[test]
1688    fn tools_list_defaults_to_fifteen_on_a_memory_store() {
1689        let db = demo_db();
1690        let resp = roundtrip(&db, r#"{"jsonrpc":"2.0","id":1,"method":"tools/list"}"#);
1691        let names: Vec<&str> = resp["result"]["tools"]
1692            .as_array()
1693            .expect("tools array")
1694            .iter()
1695            .map(|t| t["name"].as_str().expect("name"))
1696            .collect();
1697        assert_eq!(names, ASSOCIATION_TOOLS.to_vec());
1698    }
1699
1700    /// Binding: [`ASSOCIATION_TOOLS`] is a surface of its own, not the code
1701    /// door's list with a name changed.
1702    ///
1703    /// It keeps the two task tools an entity store can answer with — the notes
1704    /// it wrote and the notes it kept — and none of the seven that read a code
1705    /// graph there is none of. Every name in it is served.
1706    #[test]
1707    fn the_association_surface_is_entity_tools_only() {
1708        for kept in ["remember", "recall", "explain_association"] {
1709            assert!(
1710                ASSOCIATION_TOOLS.contains(&kept),
1711                "{kept} answers on an entity graph and must be listed"
1712            );
1713        }
1714        for code_only in [
1715            "explore", "map", "context", "impact", "owners", "why", "sync",
1716        ] {
1717            assert!(
1718                !ASSOCIATION_TOOLS.contains(&code_only),
1719                "{code_only} reads a code graph and must not be listed on a memory store"
1720            );
1721        }
1722        let served: Vec<String> = crate::mcp_tasks::task_tools()
1723            .iter()
1724            .chain(graph_tools().iter())
1725            .filter_map(|t| t.get("name").and_then(Js::as_str))
1726            .map(str::to_string)
1727            .collect();
1728        for name in ASSOCIATION_TOOLS {
1729            assert!(
1730                served.iter().any(|s| s == name),
1731                "{name} is listed but not served"
1732            );
1733        }
1734        assert!(
1735            CODE_GRAPH_TOOLS.contains(&"explore"),
1736            "and `explore` is the task tool the other surface lists"
1737        );
1738    }
1739
1740    /// Binding: the same server on a store carrying the `GitSync` marker lists
1741    /// three. One tool to find, one to query, one to size the store.
1742    #[test]
1743    fn tools_list_is_three_tools_on_a_code_graph_store() {
1744        let db = demo_db();
1745        db.write()
1746            .insert_node(
1747                "GitSync",
1748                crate::mcp_tasks::SYNC_KEY,
1749                vec![("id".into(), Value::Str(crate::mcp_tasks::SYNC_KEY.into()))],
1750            )
1751            .expect("marker");
1752        let resp = roundtrip(&db, r#"{"jsonrpc":"2.0","id":1,"method":"tools/list"}"#);
1753        let names: Vec<&str> = resp["result"]["tools"]
1754            .as_array()
1755            .expect("tools array")
1756            .iter()
1757            .map(|t| t["name"].as_str().expect("name"))
1758            .collect();
1759        assert_eq!(names, ["explore", "query", "stats"]);
1760    }
1761
1762    #[test]
1763    fn test_stats_returns_node_count() {
1764        let db = demo_db();
1765        let resp = tool_call(&db, 1, "stats", json!({}));
1766        assert!(!is_error(&resp));
1767        let result = tool_text(&resp);
1768        assert_eq!(result["nodes_live"], 2);
1769    }
1770
1771    /// A rule whose corpus is too large to index in one commit must not come
1772    /// back as a bare "ok": the caller would go straight to querying edges that
1773    /// do not exist yet.
1774    #[test]
1775    fn create_rule_reports_a_build_it_could_not_finish() {
1776        let db = SharedDb::open(&tmp_dir()).expect("open");
1777        {
1778            let mut g = db.write();
1779            for i in 0..300usize {
1780                const D: usize = 32;
1781                let axis = (i / 10) % D;
1782                let mut xs = vec![0.0f64; D];
1783                xs[axis] = 1.0;
1784                xs[(axis + 1) % D] = (i % 10) as f64 * 0.001;
1785                g.insert_node(
1786                    "V",
1787                    &format!("v{i}"),
1788                    vec![(
1789                        "emb".into(),
1790                        Value::List(xs.into_iter().map(Value::Float).collect()),
1791                    )],
1792                )
1793                .expect("insert");
1794            }
1795            g.set_hnsw_build_batch(Some(64));
1796        }
1797        let args = json!({
1798            "name": "sim",
1799            "src_label": "V",
1800            "dst_label": "V",
1801            "predicate": {"VectorSimilar": {"field": "emb", "min": 0.9}},
1802            "edge_type": "SIM",
1803            "weight_prop": null,
1804            "max_edges": null,
1805            "approximate": true
1806        });
1807        let resp = tool_call(&db, 1, "create_rule", args);
1808        assert!(!is_error(&resp), "{resp}");
1809        let result = tool_text(&resp);
1810        assert_eq!(result["name"], json!("sim"));
1811        assert_eq!(result["building"], json!({"indexed": 64, "total": 300}));
1812        let note = result["note"].as_str().expect("a note explaining the wait");
1813        assert!(
1814            note.contains("derives no edges until it finishes") && note.contains("build-index"),
1815            "the note must say the edges are not there yet and how to finish: {note}"
1816        );
1817
1818        // `stats` carries the same progress.
1819        let stats = tool_text(&tool_call(&db, 2, "stats", json!({})));
1820        let rule = stats["rules"]
1821            .as_array()
1822            .expect("rules")
1823            .iter()
1824            .find(|r| r["name"] == "sim")
1825            .expect("the rule is installed while it builds");
1826        assert_eq!(rule["edges"], json!(0));
1827        assert_eq!(
1828            rule["building"],
1829            json!({"rule": "sim", "indexed": 64, "total": 300})
1830        );
1831
1832        // Finished, the report is the plain one again.
1833        while !db.write().pump_index_build().expect("pump").is_empty() {}
1834        let stats = tool_text(&tool_call(&db, 3, "stats", json!({})));
1835        let rule = stats["rules"]
1836            .as_array()
1837            .expect("rules")
1838            .iter()
1839            .find(|r| r["name"] == "sim")
1840            .expect("rule");
1841        assert!(rule.get("building").is_none(), "{rule}");
1842        assert!(rule["edges"].as_u64().expect("edges") > 0);
1843    }
1844
1845    #[test]
1846    fn test_query_runs_cypher() {
1847        let db = demo_db();
1848        let resp = tool_call(
1849            &db,
1850            1,
1851            "query",
1852            json!({ "cypher": "MATCH (n:Person) RETURN n.name ORDER BY n.name" }),
1853        );
1854        assert!(!is_error(&resp));
1855        let result = tool_text(&resp);
1856        // columns + 2 rows
1857        assert_eq!(result["columns"], json!(["n.name"]));
1858        assert_eq!(result["rows"].as_array().map(|r| r.len()), Some(2));
1859    }
1860
1861    #[test]
1862    fn test_query_create_is_a_write() {
1863        let db = SharedDb::open(&tmp_dir()).expect("open");
1864        let resp = tool_call(
1865            &db,
1866            1,
1867            "query",
1868            json!({ "cypher": "CREATE (n:L {id: 'k'}) RETURN n" }),
1869        );
1870        assert!(
1871            !is_error(&resp),
1872            "CREATE via MCP query must succeed: {resp}"
1873        );
1874        let stats = tool_text(&tool_call(&db, 2, "stats", json!({})));
1875        assert_eq!(stats["nodes_live"], 1);
1876    }
1877
1878    #[test]
1879    fn test_ingest_json_inserts_nodes() {
1880        let db = demo_db();
1881        let resp = tool_call(
1882            &db,
1883            1,
1884            "ingest_json",
1885            json!({
1886                "label": "Person",
1887                "rows_json": r#"[{"id":"carol","name":"Carol"}]"#,
1888                "key_field": "id"
1889            }),
1890        );
1891        assert!(!is_error(&resp));
1892        // Verify node visible via stats
1893        let stats = tool_text(&tool_call(&db, 2, "stats", json!({})));
1894        assert_eq!(stats["nodes_live"], 3);
1895    }
1896
1897    #[test]
1898    fn test_node_info_returns_props() {
1899        let db = demo_db();
1900        let resp = tool_call(&db, 1, "node_info", json!({ "key": "alice" }));
1901        assert!(!is_error(&resp));
1902        let result = tool_text(&resp);
1903        assert_eq!(result["key"], "alice");
1904        assert_eq!(result["label"], "Person");
1905        assert_eq!(result["props"]["name"], "Alice");
1906    }
1907
1908    /// Binding: `node_edges` groups by edge type and names the rule and score
1909    /// behind each derived edge, in the report as in the digest.
1910    #[test]
1911    fn test_node_edges_returns_edges() {
1912        let db = demo_db();
1913        let resp = tool_call(
1914            &db,
1915            1,
1916            "node_edges",
1917            json!({ "key": "alice", "json": true }),
1918        );
1919        assert!(!is_error(&resp));
1920        let result = tool_text(&resp);
1921        assert_eq!(result["key"], "alice");
1922        let types = result["types"].as_array().expect("types");
1923        assert!(
1924            !types.is_empty(),
1925            "alice should have at least one edge type"
1926        );
1927        let similar = types
1928            .iter()
1929            .find(|t| t["edge_type"] == "SIMILAR")
1930            .expect("the rule's edge type");
1931        // A symmetric rule derives the edge both ways, and both are listed
1932        // with the direction that tells them apart.
1933        assert_eq!(similar["count"], json!(2));
1934        let edges = similar["edges"].as_array().expect("edges");
1935        let dirs: Vec<&str> = edges
1936            .iter()
1937            .map(|e| e["direction"].as_str().expect("direction"))
1938            .collect();
1939        assert!(dirs.contains(&"out") && dirs.contains(&"in"), "{similar}");
1940        for edge in edges {
1941            assert_eq!(edge["other"], json!("bob"));
1942            assert_eq!(edge["derived"], json!(true));
1943            assert_eq!(edge["rule"], json!("sim_emb"));
1944            assert_eq!(edge["score"], json!(1.0));
1945            assert!(
1946                edge["predicate"]
1947                    .as_str()
1948                    .unwrap_or("")
1949                    .contains("vector_similar"),
1950                "the predicate travels with the edge: {edge}"
1951            );
1952        }
1953    }
1954
1955    /// Binding: a depth-1 `neighborhood` is the same relationship listing, and
1956    /// anything deeper is still the traversal table.
1957    #[test]
1958    fn test_neighborhood_traverses_one_hop() {
1959        let db = demo_db();
1960        let resp = tool_call(
1961            &db,
1962            1,
1963            "neighborhood",
1964            json!({ "key": "alice", "depth": 1, "json": true }),
1965        );
1966        assert!(!is_error(&resp));
1967        let result = tool_text(&resp);
1968        assert_eq!(result["key"], "alice");
1969        assert!(result["types"].as_array().is_some(), "{result}");
1970
1971        let deep = tool_call(
1972            &db,
1973            2,
1974            "neighborhood",
1975            json!({ "key": "alice", "depth": 2 }),
1976        );
1977        assert!(!is_error(&deep));
1978        let table = tool_text(&deep);
1979        assert_eq!(table["columns"], json!(["key", "label", "depth"]));
1980        assert!(table["rows"].as_array().is_some());
1981    }
1982
1983    #[test]
1984    fn test_explain_returns_rule_info() {
1985        let db = demo_db();
1986        let resp = tool_call(&db, 1, "explain", json!({ "a": "alice", "b": "bob" }));
1987        assert!(!is_error(&resp));
1988        let result = tool_text(&resp);
1989        let arr = result.as_array().expect("explain returns array");
1990        assert!(!arr.is_empty(), "expected at least one explanation");
1991        assert_eq!(arr[0]["rule"], "sim_emb");
1992    }
1993
1994    #[test]
1995    fn test_create_rule_backfills() {
1996        let db = SharedDb::open(&tmp_dir()).expect("open");
1997        {
1998            let mut g = db.write();
1999            let opts = IngestOptions {
2000                key_field: "id".into(),
2001                auto_fk: AutoFk::Off,
2002            };
2003            let rows: Vec<BTreeMap<String, Value>> = vec![
2004                [
2005                    ("id", Value::Str("x".into())),
2006                    ("tag", Value::Str("a".into())),
2007                ]
2008                .into_iter()
2009                .map(|(k, v)| (k.to_string(), v))
2010                .collect(),
2011                [
2012                    ("id", Value::Str("y".into())),
2013                    ("tag", Value::Str("a".into())),
2014                ]
2015                .into_iter()
2016                .map(|(k, v)| (k.to_string(), v))
2017                .collect(),
2018            ];
2019            g.ingest("Item", rows, &opts).expect("ingest");
2020        }
2021        let resp = tool_call(
2022            &db,
2023            1,
2024            "create_rule",
2025            json!({
2026                "name": "same_tag",
2027                "src_label": "Item",
2028                "dst_label": "Item",
2029                "predicate": { "FieldEqual": { "field": "tag" } },
2030                "edge_type": "SAME_TAG"
2031            }),
2032        );
2033        assert!(!is_error(&resp));
2034        let result = tool_text(&resp);
2035        assert_eq!(result["ok"], true);
2036        // Derived edges should now exist.
2037        let edges_resp = tool_call(&db, 2, "node_edges", json!({ "key": "x", "json": true }));
2038        let edges_result = tool_text(&edges_resp);
2039        let types = edges_result["types"].as_array().expect("types");
2040        assert!(
2041            types.iter().any(|t| t["edge_type"] == "SAME_TAG"),
2042            "SAME_TAG edge not found after create_rule"
2043        );
2044    }
2045
2046    // --- new tools ---
2047
2048    #[test]
2049    fn test_upsert_entity_creates_new_node() {
2050        let db = demo_db();
2051        let resp = tool_call(
2052            &db,
2053            1,
2054            "upsert_entity",
2055            json!({
2056                "key": "carol",
2057                "label": "Person",
2058                "props": { "name": "Carol", "age": 30 }
2059            }),
2060        );
2061        assert!(!is_error(&resp));
2062        let result = tool_text(&resp);
2063        assert_eq!(result["ok"], true);
2064        assert_eq!(result["created"], true);
2065        assert_eq!(result["key"], "carol");
2066        // Verify node exists
2067        let info = tool_text(&tool_call(&db, 2, "node_info", json!({ "key": "carol" })));
2068        assert_eq!(info["props"]["name"], "Carol");
2069    }
2070
2071    #[test]
2072    fn test_upsert_entity_updates_existing_node() {
2073        let db = demo_db();
2074        let resp = tool_call(
2075            &db,
2076            1,
2077            "upsert_entity",
2078            json!({
2079                "key": "alice",
2080                "props": { "name": "Alice Updated" }
2081            }),
2082        );
2083        assert!(!is_error(&resp));
2084        let result = tool_text(&resp);
2085        assert_eq!(result["ok"], true);
2086        assert_eq!(result["created"], false);
2087        assert_eq!(result["updated_fields"], 1);
2088        // Verify prop changed
2089        let info = tool_text(&tool_call(&db, 2, "node_info", json!({ "key": "alice" })));
2090        assert_eq!(info["props"]["name"], "Alice Updated");
2091    }
2092
2093    #[test]
2094    fn test_upsert_entity_missing_label_on_create_is_error() {
2095        let db = demo_db();
2096        let resp = tool_call(
2097            &db,
2098            1,
2099            "upsert_entity",
2100            json!({ "key": "new-node", "props": { "x": 1 } }),
2101        );
2102        assert!(is_error(&resp), "should error without label for new node");
2103    }
2104
2105    #[test]
2106    fn test_find_similar_returns_similar_edges() {
2107        let db = demo_db();
2108        let resp = tool_call(
2109            &db,
2110            1,
2111            "find_similar",
2112            json!({ "key": "alice", "edge_type": "SIMILAR" }),
2113        );
2114        assert!(!is_error(&resp));
2115        let result = tool_text(&resp);
2116        assert_eq!(result["key"], "alice");
2117        assert_eq!(result["edge_type"], "SIMILAR");
2118        let similar = result["similar"].as_array().expect("similar array");
2119        assert!(!similar.is_empty(), "expected SIMILAR neighbors for alice");
2120        assert_eq!(similar[0]["neighbor_key"], "bob");
2121    }
2122
2123    #[test]
2124    fn test_find_similar_limit_respected() {
2125        let db = demo_db();
2126        let resp = tool_call(
2127            &db,
2128            1,
2129            "find_similar",
2130            json!({ "key": "alice", "edge_type": "SIMILAR", "limit": 0 }),
2131        );
2132        assert!(!is_error(&resp));
2133        let result = tool_text(&resp);
2134        let similar = result["similar"].as_array().expect("similar array");
2135        assert_eq!(similar.len(), 0);
2136    }
2137
2138    /// When `min` is omitted from a vector-mode find_similar call, the server
2139    /// must apply the spec default of 0.8.  A node whose cosine similarity to
2140    /// the query is 0.0 (orthogonal) must not appear in the results.
2141    #[test]
2142    fn test_find_similar_vector_default_min_is_0_8() {
2143        let db = SharedDb::open(&tmp_dir()).expect("open");
2144        {
2145            let mut g = db.write();
2146            // close: [1,0] → cosine 1.0 with query [1,0] (above 0.8)
2147            g.insert_node(
2148                "Item",
2149                "close",
2150                vec![(
2151                    "emb".into(),
2152                    Value::List(vec![Value::Float(1.0), Value::Float(0.0)]),
2153                )],
2154            )
2155            .unwrap();
2156            // far: [0,1] → cosine 0.0 with query [1,0] (below 0.8, must be excluded)
2157            g.insert_node(
2158                "Item",
2159                "far",
2160                vec![(
2161                    "emb".into(),
2162                    Value::List(vec![Value::Float(0.0), Value::Float(1.0)]),
2163                )],
2164            )
2165            .unwrap();
2166        }
2167
2168        // No `min` in the request — must default to 0.8.
2169        let resp = tool_call(
2170            &db,
2171            1,
2172            "find_similar",
2173            json!({
2174                "vector": [1.0, 0.0],
2175                "field": "emb",
2176                "label": "Item",
2177                "k": 10
2178            }),
2179        );
2180        assert!(!is_error(&resp), "vector search must not error");
2181        let result = tool_text(&resp);
2182        let results = result["results"].as_array().expect("results array");
2183
2184        let keys: Vec<&str> = results.iter().filter_map(|r| r["key"].as_str()).collect();
2185        assert!(
2186            keys.contains(&"close"),
2187            "close node (sim=1.0) must be included"
2188        );
2189        assert!(
2190            !keys.contains(&"far"),
2191            "far node (sim=0.0) must be excluded by default min=0.8"
2192        );
2193    }
2194
2195    /// `find_similar` with `mask` must exclude hidden node keys from results.
2196    #[test]
2197    fn test_find_similar_vector_mask_excludes_hidden() {
2198        let db = SharedDb::open(&tmp_dir()).expect("open");
2199        {
2200            let mut g = db.write();
2201            // visible: [1,0] — should appear in results.
2202            g.insert_node(
2203                "Item",
2204                "visible",
2205                vec![(
2206                    "emb".into(),
2207                    Value::List(vec![Value::Float(1.0), Value::Float(0.0)]),
2208                )],
2209            )
2210            .unwrap();
2211            // hidden: [1,0] — same direction as query but must not appear.
2212            g.insert_node(
2213                "Item",
2214                "hidden",
2215                vec![(
2216                    "emb".into(),
2217                    Value::List(vec![Value::Float(1.0), Value::Float(0.0)]),
2218                )],
2219            )
2220            .unwrap();
2221        }
2222
2223        let resp = tool_call(
2224            &db,
2225            1,
2226            "find_similar",
2227            json!({
2228                "vector": [1.0, 0.0],
2229                "field": "emb",
2230                "label": "Item",
2231                "k": 10,
2232                "min": 0.0,
2233                "mask": ["visible"]
2234            }),
2235        );
2236        assert!(!is_error(&resp), "masked vector search must not error");
2237        let result = tool_text(&resp);
2238        let results = result["results"].as_array().expect("results array");
2239
2240        let keys: Vec<&str> = results.iter().filter_map(|r| r["key"].as_str()).collect();
2241        assert!(
2242            keys.contains(&"visible"),
2243            "visible node must appear in masked results"
2244        );
2245        assert!(
2246            !keys.contains(&"hidden"),
2247            "hidden node must be excluded by mask"
2248        );
2249    }
2250
2251    /// `find_similar` with `mask` — bad mask value returns a tool error.
2252    #[test]
2253    fn test_find_similar_vector_mask_bad_type_is_error() {
2254        let db = SharedDb::open(&tmp_dir()).expect("open");
2255        let resp = tool_call(
2256            &db,
2257            1,
2258            "find_similar",
2259            json!({
2260                "vector": [1.0, 0.0],
2261                "field": "emb",
2262                "k": 5,
2263                "mask": [42]
2264            }),
2265        );
2266        assert!(
2267            is_error(&resp),
2268            "non-string mask element must produce a tool error"
2269        );
2270    }
2271
2272    /// Edge-traversal mode with `mask` must exclude hidden neighbors.
2273    #[test]
2274    fn test_find_similar_edge_mask_excludes_hidden_neighbor() {
2275        let db = SharedDb::open(&tmp_dir()).expect("open");
2276        {
2277            let mut g = db.write();
2278            g.insert_node("P", "alice", vec![]).unwrap();
2279            g.insert_node("P", "bob", vec![]).unwrap(); // visible
2280            g.insert_node("P", "carol", vec![]).unwrap(); // hidden
2281            g.insert_edge("KNOWS", "alice", "bob").unwrap();
2282            g.insert_edge("KNOWS", "alice", "carol").unwrap();
2283        }
2284        // Mask: alice and bob visible; carol hidden.
2285        let resp = tool_call(
2286            &db,
2287            1,
2288            "find_similar",
2289            json!({
2290                "key": "alice",
2291                "edge_type": "KNOWS",
2292                "mask": ["alice", "bob"]
2293            }),
2294        );
2295        assert!(!is_error(&resp), "masked edge search must not error");
2296        let result = tool_text(&resp);
2297        let similar = result["similar"].as_array().expect("similar array");
2298        let neighbors: Vec<&str> = similar
2299            .iter()
2300            .filter_map(|e| e["neighbor_key"].as_str())
2301            .collect();
2302        assert!(neighbors.contains(&"bob"), "bob (visible) must appear");
2303        assert!(
2304            !neighbors.contains(&"carol"),
2305            "carol (hidden) must be excluded"
2306        );
2307    }
2308
2309    /// Edge-traversal mode with `mask`: a hidden query key must not reveal
2310    /// its existence — response must be a tool error identical to a nonexistent key.
2311    #[test]
2312    fn test_find_similar_edge_mask_hidden_key_is_not_found() {
2313        let db = SharedDb::open(&tmp_dir()).expect("open");
2314        {
2315            let mut g = db.write();
2316            g.insert_node("P", "alice", vec![]).unwrap();
2317            g.insert_node("P", "bob", vec![]).unwrap();
2318        }
2319        // alice exists but is not in the mask — must look like not-found.
2320        let resp_masked = tool_call(
2321            &db,
2322            1,
2323            "find_similar",
2324            json!({ "key": "alice", "edge_type": "KNOWS", "mask": ["bob"] }),
2325        );
2326        // ghost never exists — use as the reference for "not found".
2327        let resp_ghost = tool_call(
2328            &db,
2329            2,
2330            "find_similar",
2331            json!({ "key": "ghost", "edge_type": "KNOWS" }),
2332        );
2333        assert!(
2334            is_error(&resp_masked),
2335            "hidden query key must produce a tool error"
2336        );
2337        assert!(
2338            is_error(&resp_ghost),
2339            "nonexistent key must produce a tool error"
2340        );
2341        // Both errors must carry the same shape (both are key-not-found).
2342        assert_eq!(
2343            tool_err_text(&resp_masked).contains("alice"),
2344            tool_err_text(&resp_ghost).contains("ghost"),
2345            "error messages should follow same not-found template"
2346        );
2347    }
2348
2349    /// Binding: `explain_association` now answers in prose, and the report
2350    /// behind it — what `json: true` returns — is still `explain`'s array,
2351    /// with one `evidence` object added per relationship and every other
2352    /// field unchanged.
2353    #[test]
2354    fn test_explain_association_same_as_explain() {
2355        let db = demo_db();
2356        let explain = tool_text(&tool_call(
2357            &db,
2358            1,
2359            "explain",
2360            json!({ "a": "alice", "b": "bob" }),
2361        ));
2362        let assoc = tool_text(&tool_call(
2363            &db,
2364            2,
2365            "explain_association",
2366            json!({ "a": "alice", "b": "bob", "json": true }),
2367        ));
2368        let explain: Vec<Js> = serde_json::from_value(explain).expect("explain array");
2369        let mut assoc: Vec<Js> = serde_json::from_value(assoc).expect("assoc array");
2370        for row in &mut assoc {
2371            let ev = row
2372                .as_object_mut()
2373                .expect("object")
2374                .remove("evidence")
2375                .expect("every derived edge carries its evidence");
2376            assert!(
2377                ev["similarity"].is_number(),
2378                "a vector_similar edge reports the cosine it scored: {ev}"
2379            );
2380        }
2381        assert_eq!(explain, assoc, "evidence is the only addition");
2382
2383        let prose = tool_call(
2384            &db,
2385            3,
2386            "explain_association",
2387            json!({ "a": "alice", "b": "bob" }),
2388        );
2389        let text = prose["result"]["content"][0]["text"]
2390            .as_str()
2391            .expect("text content");
2392        assert!(
2393            text.contains("mushroomdb explain — alice ↔ bob:"),
2394            "the default reply is the digest: {text}"
2395        );
2396    }
2397
2398    // ── history tools ──────────────────────────────────────────────────────────
2399
2400    /// `edge_history` must return the derived-edge lifecycle (Added event with
2401    /// rule attribution) and include the `total_commits` horizon field.
2402    #[test]
2403    fn test_edge_history_returns_derived_lifecycle_with_rule() {
2404        let db = demo_db(); // alice+bob + sim_emb rule → SIMILAR derived edge
2405        let resp = tool_call(&db, 1, "edge_history", json!({ "a": "alice", "b": "bob" }));
2406        assert!(!is_error(&resp), "edge_history must not error: {resp}");
2407        let result = tool_text(&resp);
2408
2409        // Must carry horizon metadata.
2410        let total = result["total_commits"].as_u64().expect("total_commits");
2411        assert!(total > 0, "total_commits must be > 0 after ingest + rule");
2412
2413        // Must have at least one event (the SIMILAR derived-edge addition).
2414        let events = result["events"].as_array().expect("events array");
2415        assert!(!events.is_empty(), "expected at least one edge event");
2416
2417        // At least one event must be Added with a non-null rule (derived edge).
2418        let derived_added = events
2419            .iter()
2420            .any(|ev| ev["event"].as_str() == Some("Added") && !ev["rule"].is_null());
2421        assert!(
2422            derived_added,
2423            "expected a derived Added event with rule attribution: {events:?}"
2424        );
2425    }
2426
2427    /// `was_linked` must return `true` for an edge that was active at the given commit,
2428    /// and the response must include the echo fields.
2429    #[test]
2430    fn test_was_linked_at_valid_commit() {
2431        let db = SharedDb::open(&tmp_dir()).expect("open");
2432        {
2433            let mut g = db.write();
2434            let opts = IngestOptions {
2435                key_field: "id".into(),
2436                auto_fk: AutoFk::Off,
2437            };
2438            let rows: Vec<BTreeMap<String, Value>> = vec![
2439                [("id", Value::Str("x".into()))]
2440                    .into_iter()
2441                    .map(|(k, v)| (k.to_string(), v))
2442                    .collect(),
2443                [("id", Value::Str("y".into()))]
2444                    .into_iter()
2445                    .map(|(k, v)| (k.to_string(), v))
2446                    .collect(),
2447            ];
2448            g.ingest("N", rows, &opts).expect("ingest");
2449            g.insert_edge("LINK", "x", "y").expect("edge");
2450        }
2451        // There are now at least 2 commits (ingest + edge). Check at the last one.
2452        let g = db.read();
2453        let total = g.wal_total_commits().expect("wal_total_commits");
2454        drop(g);
2455
2456        let resp = tool_call(
2457            &db,
2458            1,
2459            "was_linked",
2460            json!({ "a": "x", "b": "y", "edge_type": "LINK", "at_commit": total - 1 }),
2461        );
2462        assert!(!is_error(&resp), "was_linked must not error: {resp}");
2463        let result = tool_text(&resp);
2464        assert_eq!(result["linked"], true);
2465        assert_eq!(result["a"], "x");
2466        assert_eq!(result["edge_type"], "LINK");
2467    }
2468
2469    /// `was_linked` with an out-of-horizon commit must return a tool error (not
2470    /// a protocol error), and the error message must mention the commit range.
2471    #[test]
2472    fn test_was_linked_out_of_horizon_returns_tool_error() {
2473        let db = SharedDb::open(&tmp_dir()).expect("open");
2474        {
2475            let mut g = db.write();
2476            g.insert_node("N", "a", vec![]).expect("node a");
2477            g.insert_node("N", "b", vec![]).expect("node b");
2478        }
2479        // Commit 999 is well beyond the WAL.
2480        let resp = tool_call(
2481            &db,
2482            1,
2483            "was_linked",
2484            json!({ "a": "a", "b": "b", "edge_type": "X", "at_commit": 999 }),
2485        );
2486        // isError true = tool-level error (not a JSON-RPC protocol error).
2487        assert!(
2488            is_error(&resp),
2489            "out-of-range commit must be a tool error: {resp}"
2490        );
2491        let text = resp["result"]["content"][0]["text"].as_str().expect("text");
2492        assert!(
2493            text.contains("out of range") || text.contains("range"),
2494            "error must mention range: {text}"
2495        );
2496    }
2497
2498    /// `node_history` tool must return the node's WAL history and the
2499    /// `total_commits` horizon field.
2500    #[test]
2501    fn test_node_history_via_mcp() {
2502        let db = demo_db(); // alice + bob, with a SIMILAR rule
2503        let resp = tool_call(&db, 1, "node_history", json!({ "key": "alice" }));
2504        assert!(!is_error(&resp), "node_history must not error: {resp}");
2505        let result = tool_text(&resp);
2506
2507        assert_eq!(result["key"], "alice");
2508        let total = result["total_commits"].as_u64().expect("total_commits");
2509        assert!(total > 0, "total_commits must be > 0");
2510
2511        let history = result["history"].as_array().expect("history array");
2512        assert!(
2513            !history.is_empty(),
2514            "alice should have at least one history entry"
2515        );
2516
2517        // First event should be a NodeInserted.
2518        let first_change = &history[0]["change"];
2519        assert_eq!(first_change["type"], "NodeInserted");
2520        assert_eq!(first_change["label"], "Person");
2521    }
2522}