1use 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 PropPredicate, 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
64pub 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
79pub 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 ToolOk(Js),
196 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 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 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 let namespace = match namespace_arg(args.get("namespace")) {
283 Ok(n) => n,
284 Err(e) => return CallOutcome::ToolErr(e),
285 };
286
287 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 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 Some(ns) => AsOfScope::Namespace(ns),
320 None => {
321 return match db.read().query_at(commit, cypher, ¶ms) {
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 (Some(ns), true) => g.query_at_scoped_in_namespace(commit, cypher, ¶ms, scope, ns),
332 _ => g.query_at_scoped(commit, cypher, ¶ms, 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 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 (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 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, ¶ms, &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, ¶ms)
392 } else {
393 let g = db.read();
394 g.query(cypher, ¶ms)
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
402fn 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 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 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
541fn 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 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
613fn tool_upsert_entity(db: &SharedDb, args: &Js) -> CallOutcome {
638 let Some(key) = args.get("key").and_then(Js::as_str) else {
639 return CallOutcome::ToolErr("missing key".into());
640 };
641 let label_opt = args.get("label").and_then(Js::as_str);
642 let Some(props_obj) = args.get("props").and_then(Js::as_object) else {
643 return CallOutcome::ToolErr("missing props".into());
644 };
645 let namespace = match namespace_arg(args.get("namespace")) {
646 Ok(n) => n,
647 Err(e) => return CallOutcome::ToolErr(e),
648 };
649
650 let mut row: BTreeMap<String, Value> = BTreeMap::new();
654 for (field, json_val) in props_obj {
655 if field == "id" {
656 continue;
657 }
658 match json_to_value(json_val.clone()) {
659 Some(v) => {
660 row.insert(field.clone(), v);
661 }
662 None => {
663 return CallOutcome::ToolErr(format!("prop {field} is not a supported value type"))
664 }
665 }
666 }
667 if let Some(ns) = namespace.as_deref() {
668 if let Err(e) = stamp_namespace_row(&mut row, ns) {
669 return CallOutcome::ToolErr(e);
670 }
671 }
672
673 let exists = {
674 let g = db.read();
675 g.has_node(key)
676 };
677
678 if exists {
679 let mut g = db.write();
680 let mut to_set: Vec<(String, Value)> = Vec::new();
681 for (field, v) in row {
682 if field == NS_PROP && Some(&v) == g.namespace_of(key).map(Value::Str).as_ref() {
688 continue;
689 }
690 to_set.push((field, v));
691 }
692 let count = to_set.len();
693 if let Err(e) = g.set_props(key, to_set) {
694 return CallOutcome::ToolErr(graph_err_msg(e));
695 }
696 CallOutcome::ToolOk(json!({
697 "ok": true,
698 "key": key,
699 "created": false,
700 "updated_fields": count
701 }))
702 } else {
703 let Some(label) = label_opt else {
704 return CallOutcome::ToolErr("label required when creating a new entity".into());
705 };
706 row.insert("id".to_string(), Value::Str(key.to_string()));
707 let opts = IngestOptions {
708 key_field: "id".to_string(),
709 auto_fk: AutoFk::Off,
710 };
711 let mut g = db.write();
712 match g.ingest(label, vec![row], &opts) {
713 Ok(_) => CallOutcome::ToolOk(json!({ "ok": true, "key": key, "created": true })),
714 Err(e) => CallOutcome::ToolErr(graph_err_msg(e)),
715 }
716 }
717}
718
719fn tool_find_similar(db: &SharedDb, args: &Js) -> CallOutcome {
726 let mask_keys: Option<Vec<String>> = if let Some(mask_val) = args.get("mask") {
729 match mask_val.as_array() {
730 Some(arr) => {
731 let mut ks: Vec<String> = Vec::with_capacity(arr.len());
732 for v in arr {
733 match v.as_str() {
734 Some(s) => ks.push(s.to_string()),
735 None => {
736 return CallOutcome::ToolErr("mask must be an array of strings".into())
737 }
738 }
739 }
740 Some(ks)
741 }
742 None => return CallOutcome::ToolErr("mask must be an array of strings".into()),
743 }
744 } else {
745 None
746 };
747
748 if let Some(vec_js) = args.get("vector").and_then(Js::as_array) {
751 let q: Vec<f64> = vec_js.iter().filter_map(|v| v.as_f64()).collect();
752 if q.is_empty() {
753 return CallOutcome::ToolErr("vector must be a non-empty array of numbers".into());
754 }
755 let field = args
756 .get("field")
757 .and_then(Js::as_str)
758 .unwrap_or("embedding");
759 let label_str = args.get("label").and_then(Js::as_str).unwrap_or("");
760 let label = if label_str.is_empty() {
761 None
762 } else {
763 Some(label_str)
764 };
765 let k = args
766 .get("k")
767 .and_then(Js::as_u64)
768 .map(|n| n as usize)
769 .unwrap_or(10);
770 let min = args.get("min").and_then(Js::as_f64).unwrap_or(0.8);
771 let where_pred = match parse_where_arg(args) {
772 Ok(p) => p,
773 Err(e) => return CallOutcome::ToolErr(e),
774 };
775 let exact = match args.get("exact") {
776 None => false,
777 Some(v) => match v.as_bool() {
778 Some(b) => b,
779 None => return CallOutcome::ToolErr("exact must be a boolean".into()),
780 },
781 };
782 let exact = exact || where_pred.is_some();
783
784 let hits = {
785 let g = db.read();
786 let node_mask = mask_keys
787 .as_ref()
788 .map(|keys| NodeMask::from_keys(&*g, keys.iter().map(String::as_str)));
789 match g.find_similar_vector_filtered(
790 field,
791 label,
792 &q,
793 k,
794 min,
795 node_mask.as_ref(),
796 where_pred.as_ref(),
797 exact,
798 ) {
799 Ok(h) => h,
800 Err(e) => return CallOutcome::ToolErr(graph_err_msg(e)),
801 }
802 };
803 let results: Vec<Js> = hits
804 .into_iter()
805 .map(|(key, score)| json!({ "key": key, "score": score }))
806 .collect();
807 return CallOutcome::ToolOk(json!({
808 "mode": "vector",
809 "field": field,
810 "label": label,
811 "k": k,
812 "min": min,
813 "results": results
814 }));
815 }
816
817 let Some(key) = args.get("key").and_then(Js::as_str) else {
819 return CallOutcome::ToolErr("missing key (or provide vector for vector search)".into());
820 };
821 let edge_type = args
822 .get("edge_type")
823 .and_then(Js::as_str)
824 .unwrap_or("SIMILAR");
825 let limit = args
826 .get("limit")
827 .and_then(Js::as_u64)
828 .map(|n| n as usize)
829 .unwrap_or(10);
830
831 if let Some(ref mask) = mask_keys {
834 let mask_set: std::collections::HashSet<&str> = mask.iter().map(String::as_str).collect();
835 if !mask_set.contains(key) {
836 return CallOutcome::ToolErr(graph_err_msg(GraphError::KeyNotFound {
837 key: key.into(),
838 }));
839 }
840 let out = {
841 let g = db.read();
842 g.node_edges(key)
843 };
844 return match out {
845 Ok(edges) => {
846 let similar: Vec<Js> = edges
847 .iter()
848 .filter(|e| e.edge_type == edge_type)
849 .filter(|e| {
850 let neighbor_key = if e.src_key == key {
852 &e.dst_key
853 } else {
854 &e.src_key
855 };
856 mask_set.contains(neighbor_key.as_str())
857 })
858 .take(limit)
859 .map(|e| {
860 let neighbor_key = if e.src_key == key {
861 &e.dst_key
862 } else {
863 &e.src_key
864 };
865 let direction = if e.src_key == key { "out" } else { "in" };
866 json!({
867 "neighbor_key": neighbor_key,
868 "direction": direction,
869 "edge_type": e.edge_type,
870 "derived": e.derived,
871 })
872 })
873 .collect();
874 CallOutcome::ToolOk(json!({
875 "key": key,
876 "edge_type": edge_type,
877 "similar": similar
878 }))
879 }
880 Err(e) => CallOutcome::ToolErr(graph_err_msg(e)),
881 };
882 }
883
884 let out = {
885 let g = db.read();
886 g.node_edges(key)
887 };
888 match out {
889 Ok(edges) => {
890 let similar: Vec<Js> = edges
891 .iter()
892 .filter(|e| e.edge_type == edge_type)
893 .take(limit)
894 .map(|e| {
895 let neighbor_key = if e.src_key == key {
896 &e.dst_key
897 } else {
898 &e.src_key
899 };
900 let direction = if e.src_key == key { "out" } else { "in" };
901 json!({
902 "neighbor_key": neighbor_key,
903 "direction": direction,
904 "edge_type": e.edge_type,
905 "derived": e.derived,
906 })
907 })
908 .collect();
909 CallOutcome::ToolOk(json!({
910 "key": key,
911 "edge_type": edge_type,
912 "similar": similar
913 }))
914 }
915 Err(e) => CallOutcome::ToolErr(graph_err_msg(e)),
916 }
917}
918
919fn tool_hybrid_search(db: &SharedDb, args: &Js) -> CallOutcome {
920 let Some(query_text) = args.get("query_text").and_then(Js::as_str) else {
921 return CallOutcome::ToolErr("missing required field: query_text".into());
922 };
923 let Some(text_field) = args.get("text_field").and_then(Js::as_str) else {
924 return CallOutcome::ToolErr("missing required field: text_field".into());
925 };
926
927 let vector_field = args
928 .get("vector_field")
929 .and_then(Js::as_str)
930 .unwrap_or("embedding");
931 let label = args.get("label").and_then(Js::as_str);
932 let k = args
933 .get("k")
934 .and_then(Js::as_u64)
935 .map(|n| n as usize)
936 .unwrap_or(10);
937
938 let query_vec: Vec<f64> = args
939 .get("vector")
940 .and_then(Js::as_array)
941 .map(|arr| arr.iter().filter_map(|v| v.as_f64()).collect())
942 .unwrap_or_default();
943
944 let hits = {
945 let g = db.read();
946 g.search_hybrid(text_field, query_text, vector_field, &query_vec, label, k)
947 };
948
949 let results: Vec<Js> = hits
950 .into_iter()
951 .map(|(key, score)| json!({ "key": key, "score": score }))
952 .collect();
953
954 CallOutcome::ToolOk(json!({
955 "query_text": query_text,
956 "text_field": text_field,
957 "vector_field": vector_field,
958 "label": label,
959 "k": k,
960 "results": results
961 }))
962}
963
964fn tool_node_history(db: &SharedDb, args: &Js) -> CallOutcome {
965 let Some(key) = args.get("key").and_then(Js::as_str) else {
966 return CallOutcome::ToolErr("missing key".into());
967 };
968 let g = db.read();
969 let result = match g.node_history(key) {
970 Ok(e) => e,
971 Err(e) => return CallOutcome::ToolErr(graph_err_msg(e)),
972 };
973 CallOutcome::ToolOk(node_history_json(key, &result))
974}
975
976fn tool_edge_history(db: &SharedDb, args: &Js) -> CallOutcome {
977 let Some(a) = args.get("a").and_then(Js::as_str).filter(|s| !s.is_empty()) else {
978 return CallOutcome::ToolErr("missing a".into());
979 };
980 let Some(b) = args.get("b").and_then(Js::as_str).filter(|s| !s.is_empty()) else {
981 return CallOutcome::ToolErr("missing b".into());
982 };
983 let result = {
984 let g = db.read();
985 g.edge_history(a, b)
986 };
987 match result {
988 Ok(hr) => CallOutcome::ToolOk(edge_history_result_json(a, b, &hr)),
989 Err(e) => CallOutcome::ToolErr(graph_err_msg(e)),
990 }
991}
992
993fn tool_was_linked(db: &SharedDb, args: &Js) -> CallOutcome {
994 let Some(a) = args.get("a").and_then(Js::as_str).filter(|s| !s.is_empty()) else {
995 return CallOutcome::ToolErr("missing a".into());
996 };
997 let Some(b) = args.get("b").and_then(Js::as_str).filter(|s| !s.is_empty()) else {
998 return CallOutcome::ToolErr("missing b".into());
999 };
1000 let Some(edge_type) = args
1001 .get("edge_type")
1002 .and_then(Js::as_str)
1003 .filter(|s| !s.is_empty())
1004 else {
1005 return CallOutcome::ToolErr("missing edge_type".into());
1006 };
1007 let at_commit = match args.get("at_commit").and_then(Js::as_u64) {
1008 Some(n) => n,
1009 None => return CallOutcome::ToolErr("missing or invalid at_commit".into()),
1010 };
1011 let result = {
1012 let g = db.read();
1013 g.was_linked(a, b, edge_type, at_commit)
1014 };
1015 match result {
1016 Ok(linked) => CallOutcome::ToolOk(json!({
1017 "a": a,
1018 "b": b,
1019 "edge_type": edge_type,
1020 "at_commit": at_commit,
1021 "linked": linked,
1022 })),
1023 Err(e) => CallOutcome::ToolErr(graph_err_msg(e)),
1024 }
1025}
1026
1027fn tool_rename_node(db: &SharedDb, args: &Js) -> CallOutcome {
1028 let Some(old_key) = args.get("old_key").and_then(Js::as_str) else {
1029 return CallOutcome::ToolErr("missing old_key".into());
1030 };
1031 let Some(new_key) = args.get("new_key").and_then(Js::as_str) else {
1032 return CallOutcome::ToolErr("missing new_key".into());
1033 };
1034 let mut g = db.write();
1035 match g.rename_node(old_key, new_key) {
1036 Ok(()) => CallOutcome::ToolOk(json!({
1037 "ok": true,
1038 "old_key": old_key,
1039 "new_key": new_key,
1040 })),
1041 Err(e) => CallOutcome::ToolErr(graph_err_msg(e)),
1042 }
1043}
1044
1045fn parse_where_arg(args: &Js) -> std::result::Result<Option<PropPredicate>, String> {
1046 let Some(w) = args.get("where") else {
1047 return Ok(None);
1048 };
1049 if w.is_null() {
1050 return Ok(None);
1051 }
1052 let pred: PropPredicate =
1053 serde_json::from_value(w.clone()).map_err(|e| format!("where: {e}"))?;
1054 pred.validate_named("where")?;
1055 Ok(Some(pred))
1056}
1057
1058pub(crate) fn graph_err_msg(e: GraphError) -> String {
1059 match e {
1060 GraphError::QueryError { detail } | GraphError::IngestError { detail } => detail,
1061 other => other.to_string(),
1062 }
1063}
1064
1065fn initialize_result() -> Js {
1066 json!({
1067 "protocolVersion": "2024-11-05",
1068 "capabilities": { "tools": {} },
1069 "serverInfo": { "name": "mushroomdb", "version": env!("CARGO_PKG_VERSION") }
1070 })
1071}
1072
1073const ADVANCED_PREFIX: &str = "Advanced: ";
1079
1080pub const CODE_GRAPH_TOOLS: [&str; 3] = ["explore", "query", "stats"];
1088
1089pub const ASSOCIATION_TOOLS: [&str; 15] = [
1114 "query",
1115 "explain_association",
1116 "neighborhood",
1117 "node_info",
1118 "node_edges",
1119 "was_linked",
1120 "edges_at",
1121 "what_if",
1122 "node_history",
1123 "edge_history",
1124 "find_similar",
1125 "hybrid_search",
1126 "remember",
1127 "recall",
1128 "stats",
1129];
1130
1131#[derive(Clone, Copy, Debug, PartialEq, Eq)]
1137pub(crate) enum Surface {
1138 CodeGraph,
1141 Memory,
1144}
1145
1146impl Surface {
1147 fn listing(self) -> &'static [&'static str] {
1150 match self {
1151 Surface::CodeGraph => &CODE_GRAPH_TOOLS,
1152 Surface::Memory => &ASSOCIATION_TOOLS,
1153 }
1154 }
1155}
1156
1157fn surface_of(db: &SharedDb) -> Surface {
1160 let ingested = {
1161 let g = db.read();
1162 g.has_node(crate::mcp_tasks::SYNC_KEY)
1163 };
1164 if ingested {
1165 Surface::CodeGraph
1166 } else {
1167 Surface::Memory
1168 }
1169}
1170
1171fn tools_list(all: bool, surface: Surface) -> Js {
1189 let mut served: Vec<Js> = crate::mcp_tasks::task_tools();
1190 for mut tool in graph_tools() {
1191 if let Some(d) = tool.get("description").and_then(Js::as_str) {
1192 let prefixed = format!("{ADVANCED_PREFIX}{d}");
1193 tool["description"] = Js::String(prefixed);
1194 }
1195 served.push(tool);
1196 }
1197 if all {
1198 return json!({ "tools": served });
1199 }
1200 let listing = surface.listing();
1201 let mut tools: Vec<Js> = Vec::with_capacity(listing.len());
1202 for name in listing {
1203 let Some(tool) = served
1204 .iter()
1205 .find(|t| t.get("name").and_then(Js::as_str) == Some(*name))
1206 else {
1207 debug_assert!(false, "{surface:?} lists {name}, which is not served");
1208 continue;
1209 };
1210 tools.push(tool.clone());
1211 }
1212 json!({ "tools": tools })
1213}
1214
1215fn graph_tools() -> Vec<Js> {
1218 let Js::Array(tools) = json!([
1219 {
1220 "name": "query",
1221 "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.",
1222 "inputSchema": {
1223 "type": "object",
1224 "properties": {
1225 "cypher": { "type": "string", "description": "Cypher query text." },
1226 "params": {
1227 "type": "object",
1228 "description": "Named JSON-scalar query parameters."
1229 },
1230 "mask": {
1231 "type": "array",
1232 "items": { "type": "string" },
1233 "description": "Optional node key allow-list. When present, only these nodes are visible; write statements are rejected."
1234 },
1235 "role": {
1236 "type": "string",
1237 "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."
1238 },
1239 "as_of": {
1240 "type": "integer",
1241 "minimum": 0,
1242 "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."
1243 },
1244 "namespace": {
1245 "type": "string",
1246 "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."
1247 }
1248 },
1249 "required": ["cypher"]
1250 }
1251 },
1252 {
1253 "name": "ingest_json",
1254 "description": "Ingest a JSON array of objects as nodes of one label.",
1255 "inputSchema": {
1256 "type": "object",
1257 "properties": {
1258 "label": { "type": "string" },
1259 "rows_json": {
1260 "type": "string",
1261 "description": "JSON text of an array of objects."
1262 },
1263 "key_field": { "type": "string" },
1264 "auto_fk_suffix": { "type": "string" },
1265 "edges": {
1266 "type": "array",
1267 "description": "Optional user edges [{edge_type, src, dst}]."
1268 },
1269 "namespace": {
1270 "type": "string",
1271 "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."
1272 }
1273 },
1274 "required": ["label", "rows_json"]
1275 }
1276 },
1277 {
1278 "name": "create_rule",
1279 "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.",
1280 "inputSchema": {
1281 "type": "object",
1282 "properties": {
1283 "name": { "type": "string" },
1284 "src_label": { "type": "string" },
1285 "dst_label": { "type": "string" },
1286 "predicate": { "type": "object" },
1287 "edge_type": { "type": "string" },
1288 "weight_prop": {
1289 "type": ["string", "null"],
1290 "description": "Edge property that stores the score (default: weight)."
1291 },
1292 "max_edges": { "type": ["integer", "null"] },
1293 "namespace": {
1294 "type": "string",
1295 "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."
1296 }
1297 },
1298 "required": ["name", "src_label", "dst_label", "predicate", "edge_type"]
1299 }
1300 },
1301 {
1302 "name": "explain",
1303 "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.",
1304 "inputSchema": {
1305 "type": "object",
1306 "properties": {
1307 "a": { "type": "string", "minLength": 1 },
1308 "b": { "type": "string", "minLength": 1 }
1309 },
1310 "required": ["a", "b"]
1311 }
1312 },
1313 {
1314 "name": "stats",
1315 "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.",
1316 "inputSchema": {
1317 "type": "object",
1318 "properties": {
1319 "role": {
1320 "type": "string",
1321 "description": "Report only the namespaces this role may see. The store-wide counts beside them are unchanged."
1322 },
1323 "namespace": {
1324 "type": "string",
1325 "description": "Report only this namespace. Intersects with 'role'."
1326 }
1327 }
1328 }
1329 },
1330 {
1331 "name": "node_info",
1332 "description": "What is K — its label and every property it holds.",
1333 "inputSchema": {
1334 "type": "object",
1335 "properties": {
1336 "key": { "type": "string" }
1337 },
1338 "required": ["key"]
1339 }
1340 },
1341 {
1342 "name": "upsert_entity",
1343 "description": "Record what is now true about K — insert or update a node by key. If the key exists, updates the supplied properties atomically: every property is checked before any is written, so a refusal leaves the node unchanged. 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.",
1344 "inputSchema": {
1345 "type": "object",
1346 "properties": {
1347 "key": { "type": "string", "description": "Unique node key." },
1348 "label": { "type": "string", "description": "Node label (required when creating a new entity)." },
1349 "props": {
1350 "type": "object",
1351 "description": "Properties to set. Values must be scalars (string, number, bool) or arrays of scalars."
1352 },
1353 "namespace": {
1354 "type": "string",
1355 "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."
1356 }
1357 },
1358 "required": ["key", "props"]
1359 }
1360 },
1361 {
1362 "name": "find_similar",
1363 "description": "What is most like this — two modes: (1) Vector search — provide `vector` (and optionally `field`, `label`, `k`, `min`, `where`, `exact`) to find the k most similar nodes by cosine similarity using the HNSW index when available, brute-force otherwise. `where` is a property predicate (`{field, eq}` or `{field, in}`) and implies exact search. `exact` true skips HNSW. (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. Edge-traversal mode ignores `where` and `exact`. 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.",
1364 "inputSchema": {
1365 "type": "object",
1366 "properties": {
1367 "vector": {
1368 "type": "array",
1369 "items": { "type": "number" },
1370 "description": "Query embedding vector for vector-similarity search. When present, vector-search mode is used and `key` is ignored."
1371 },
1372 "field": { "type": "string", "description": "Property field holding the embedding vectors (default: embedding). Used in vector-search mode." },
1373 "label": { "type": "string", "description": "Restrict search to nodes with this label. Empty string means all labels. Used in vector-search mode." },
1374 "k": { "type": "integer", "description": "Maximum results to return in vector-search mode (default: 10)." },
1375 "min": { "type": "number", "description": "Minimum cosine similarity threshold in vector-search mode (default: 0.8)." },
1376 "mask": {
1377 "type": "array",
1378 "items": { "type": "string" },
1379 "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. The beam widens until it has k visible hits, then falls back to an exhaustive masked scan at the same cap an exact VectorSimilar rule uses, so the result is not short while more visible hits exist. Unknown keys are silently ignored."
1380 },
1381 "where": {
1382 "type": "object",
1383 "description": "Optional property predicate for vector-search mode, same shape as visible_where: {\"field\": \"...\", \"eq\": value} or {\"field\": \"...\", \"in\": [values]}. Implies exact search (skips HNSW). Invalid predicates are a tool error. Edge-traversal mode ignores this."
1384 },
1385 "exact": {
1386 "type": "boolean",
1387 "description": "When true, vector-search mode uses exact GEMM brute force and does not consult HNSW. Default false. Edge-traversal mode ignores this."
1388 },
1389 "key": { "type": "string", "description": "Source node key for edge-traversal mode." },
1390 "edge_type": { "type": "string", "description": "Edge type to filter by in edge-traversal mode (default: SIMILAR)." },
1391 "limit": { "type": "integer", "description": "Maximum neighbors to return in edge-traversal mode (default: 10)." }
1392 }
1393 }
1394 },
1395 {
1396 "name": "hybrid_search",
1397 "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.",
1398 "inputSchema": {
1399 "type": "object",
1400 "properties": {
1401 "query_text": { "type": "string", "description": "Fulltext query string." },
1402 "text_field": { "type": "string", "description": "Property field to search with fulltext." },
1403 "vector": {
1404 "type": "array",
1405 "items": { "type": "number" },
1406 "description": "Query embedding vector. Omit for text-only ranking."
1407 },
1408 "vector_field": { "type": "string", "description": "Property field holding embedding vectors (default: embedding)." },
1409 "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." },
1410 "k": { "type": "integer", "description": "Maximum results to return (default: 10)." }
1411 },
1412 "required": ["query_text", "text_field"]
1413 }
1414 },
1415 {
1416 "name": "node_history",
1417 "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.",
1418 "inputSchema": {
1419 "type": "object",
1420 "properties": {
1421 "key": { "type": "string", "description": "Node key to look up." }
1422 },
1423 "required": ["key"]
1424 }
1425 },
1426 {
1427 "name": "edge_history",
1428 "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.",
1429 "inputSchema": {
1430 "type": "object",
1431 "properties": {
1432 "a": { "type": "string", "minLength": 1, "description": "First node key." },
1433 "b": { "type": "string", "minLength": 1, "description": "Second node key." }
1434 },
1435 "required": ["a", "b"]
1436 }
1437 },
1438 {
1439 "name": "was_linked",
1440 "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`).",
1441 "inputSchema": {
1442 "type": "object",
1443 "properties": {
1444 "a": { "type": "string", "minLength": 1, "description": "First node key." },
1445 "b": { "type": "string", "minLength": 1, "description": "Second node key." },
1446 "edge_type": { "type": "string", "minLength": 1, "description": "Edge type to check." },
1447 "at_commit": { "type": "integer", "minimum": 0, "description": "0-based WAL commit index to query." }
1448 },
1449 "required": ["a", "b", "edge_type", "at_commit"]
1450 }
1451 },
1452 {
1453 "name": "rename_node",
1454 "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.",
1455 "inputSchema": {
1456 "type": "object",
1457 "properties": {
1458 "old_key": { "type": "string", "minLength": 1, "description": "Current node key." },
1459 "new_key": { "type": "string", "minLength": 1, "description": "Desired new node key." }
1460 },
1461 "required": ["old_key", "new_key"]
1462 }
1463 }
1464 ]) else {
1465 unreachable!("the literal above is an array")
1466 };
1467 tools
1468}
1469
1470fn tool_ok(payload: Js) -> Js {
1471 json!({
1472 "content": [{ "type": "text", "text": payload.to_string() }]
1473 })
1474}
1475
1476fn task_ok(text: &str) -> Js {
1484 json!({
1485 "content": [{ "type": "text", "text": text }]
1486 })
1487}
1488
1489fn tool_err(message: &str) -> Js {
1490 json!({
1491 "content": [{ "type": "text", "text": message }],
1492 "isError": true
1493 })
1494}
1495
1496fn write_result(writer: &mut impl Write, id: Option<Js>, result: Js) -> io::Result<()> {
1497 write_json(
1498 writer,
1499 &json!({
1500 "jsonrpc": "2.0",
1501 "id": id.unwrap_or(Js::Null),
1502 "result": result
1503 }),
1504 )
1505}
1506
1507fn write_error(
1508 writer: &mut impl Write,
1509 id: Option<Js>,
1510 code: i64,
1511 message: &str,
1512) -> io::Result<()> {
1513 write_json(
1514 writer,
1515 &json!({
1516 "jsonrpc": "2.0",
1517 "id": id.unwrap_or(Js::Null),
1518 "error": { "code": code, "message": message }
1519 }),
1520 )
1521}
1522
1523fn write_json(writer: &mut impl Write, value: &Js) -> io::Result<()> {
1524 let s = serde_json::to_string(value).map_err(io::Error::other)?;
1525 writeln!(writer, "{s}")?;
1526 writer.flush()
1527}
1528
1529#[cfg(test)]
1534mod tests {
1535 use super::*;
1536 use core_api::{AutoFk, IngestOptions, Predicate, RuleDef, Value};
1537 use std::path::PathBuf;
1538 use std::sync::atomic::{AtomicU64, Ordering};
1539
1540 fn tmp_dir() -> PathBuf {
1541 static SEQ: AtomicU64 = AtomicU64::new(0);
1542 let n = SEQ.fetch_add(1, Ordering::Relaxed);
1543 let d = std::env::temp_dir().join(format!("mcp-test-{}-{}", std::process::id(), n));
1544 let _ = std::fs::remove_dir_all(&d);
1548 d
1549 }
1550
1551 fn demo_db() -> SharedDb {
1553 let db = SharedDb::open(&tmp_dir()).expect("open");
1554 {
1555 let mut g = db.write();
1556 let opts = IngestOptions {
1557 key_field: "id".into(),
1558 auto_fk: AutoFk::Off,
1559 };
1560 let people: Vec<BTreeMap<String, Value>> = vec![
1562 [
1563 ("id", Value::Str("alice".into())),
1564 ("name", Value::Str("Alice".into())),
1565 (
1566 "emb",
1567 Value::List(vec![Value::Float(1.0), Value::Float(0.0)]),
1568 ),
1569 ]
1570 .into_iter()
1571 .map(|(k, v)| (k.to_string(), v))
1572 .collect(),
1573 [
1574 ("id", Value::Str("bob".into())),
1575 ("name", Value::Str("Bob".into())),
1576 (
1577 "emb",
1578 Value::List(vec![Value::Float(1.0), Value::Float(0.0)]),
1579 ),
1580 ]
1581 .into_iter()
1582 .map(|(k, v)| (k.to_string(), v))
1583 .collect(),
1584 ];
1585 g.ingest("Person", people, &opts).expect("ingest");
1586
1587 g.create_rule(RuleDef {
1589 name: "sim_emb".into(),
1590 src_label: "Person".into(),
1591 dst_label: "Person".into(),
1592 predicate: Predicate::VectorSimilar {
1593 field: "emb".into(),
1594 min: 0.9,
1595 },
1596 edge_type: "SIMILAR".into(),
1597 weight_prop: Some("score".into()),
1598 max_edges: None,
1599 approximate: false,
1600 via_label: None,
1601 via_edge: None,
1602 via_dir: None,
1603 namespace: None,
1604 })
1605 .expect("rule");
1606 }
1607 db
1608 }
1609
1610 fn roundtrip(db: &SharedDb, request: &str) -> Js {
1611 roundtrip_with(db, false, request)
1612 }
1613
1614 fn roundtrip_with(db: &SharedDb, all_tools: bool, request: &str) -> Js {
1615 let input = format!("{request}\n");
1616 let mut output = Vec::new();
1617 run_mcp_stdio_with(db.clone(), None, all_tools, input.as_bytes(), &mut output)
1618 .expect("mcp");
1619 let s = std::str::from_utf8(&output).expect("utf8");
1620 serde_json::from_str(s.trim()).expect("json response")
1621 }
1622
1623 fn tool_call(db: &SharedDb, id: u64, tool: &str, args: Js) -> Js {
1624 let req = json!({
1625 "jsonrpc": "2.0",
1626 "id": id,
1627 "method": "tools/call",
1628 "params": { "name": tool, "arguments": args }
1629 });
1630 roundtrip(db, &req.to_string())
1631 }
1632
1633 fn tool_text(resp: &Js) -> Js {
1635 let text = resp["result"]["content"][0]["text"]
1636 .as_str()
1637 .expect("content[0].text");
1638 serde_json::from_str(text).expect("tool text is json")
1639 }
1640
1641 fn is_error(resp: &Js) -> bool {
1642 resp["result"]["isError"].as_bool().unwrap_or(false)
1643 }
1644
1645 fn tool_err_text(resp: &Js) -> String {
1646 resp["result"]["content"][0]["text"]
1647 .as_str()
1648 .unwrap_or("")
1649 .to_string()
1650 }
1651
1652 #[test]
1655 fn test_tools_list_includes_all_expected() {
1656 let db = demo_db();
1657 let resp = roundtrip_with(
1658 &db,
1659 true,
1660 r#"{"jsonrpc":"2.0","id":1,"method":"tools/list"}"#,
1661 );
1662 let tools = resp["result"]["tools"].as_array().expect("tools array");
1663 let names: Vec<&str> = tools
1664 .iter()
1665 .map(|t| t["name"].as_str().expect("name"))
1666 .collect();
1667 for expected in &[
1668 "explore",
1670 "map",
1671 "context",
1672 "impact",
1673 "owners",
1674 "why",
1675 "explain_association",
1676 "node_edges",
1677 "neighborhood",
1678 "edges_at",
1679 "what_if",
1680 "recall",
1681 "remember",
1682 "sync",
1683 "query",
1685 "ingest_json",
1686 "create_rule",
1687 "explain",
1688 "stats",
1689 "node_info",
1690 "upsert_entity",
1691 "find_similar",
1692 "hybrid_search",
1693 "node_history",
1694 "edge_history",
1695 "was_linked",
1696 "rename_node",
1697 ] {
1698 assert!(names.contains(expected), "missing tool: {expected}");
1699 }
1700 assert_eq!(
1701 names.len(),
1702 27,
1703 "expected exactly 27 tools, got {}",
1704 names.len()
1705 );
1706 assert_eq!(
1707 &names[..14],
1708 [
1709 "explore",
1710 "map",
1711 "context",
1712 "impact",
1713 "owners",
1714 "why",
1715 "explain_association",
1716 "node_edges",
1717 "neighborhood",
1718 "edges_at",
1719 "what_if",
1720 "recall",
1721 "remember",
1722 "sync"
1723 ],
1724 "the task tools come first, in order"
1725 );
1726 assert_eq!(names[14], "query", "the graph tools follow them");
1727 }
1728
1729 #[test]
1733 fn tools_list_defaults_to_fifteen_on_a_memory_store() {
1734 let db = demo_db();
1735 let resp = roundtrip(&db, r#"{"jsonrpc":"2.0","id":1,"method":"tools/list"}"#);
1736 let names: Vec<&str> = resp["result"]["tools"]
1737 .as_array()
1738 .expect("tools array")
1739 .iter()
1740 .map(|t| t["name"].as_str().expect("name"))
1741 .collect();
1742 assert_eq!(names, ASSOCIATION_TOOLS.to_vec());
1743 }
1744
1745 #[test]
1752 fn the_association_surface_is_entity_tools_only() {
1753 for kept in ["remember", "recall", "explain_association"] {
1754 assert!(
1755 ASSOCIATION_TOOLS.contains(&kept),
1756 "{kept} answers on an entity graph and must be listed"
1757 );
1758 }
1759 for code_only in [
1760 "explore", "map", "context", "impact", "owners", "why", "sync",
1761 ] {
1762 assert!(
1763 !ASSOCIATION_TOOLS.contains(&code_only),
1764 "{code_only} reads a code graph and must not be listed on a memory store"
1765 );
1766 }
1767 let served: Vec<String> = crate::mcp_tasks::task_tools()
1768 .iter()
1769 .chain(graph_tools().iter())
1770 .filter_map(|t| t.get("name").and_then(Js::as_str))
1771 .map(str::to_string)
1772 .collect();
1773 for name in ASSOCIATION_TOOLS {
1774 assert!(
1775 served.iter().any(|s| s == name),
1776 "{name} is listed but not served"
1777 );
1778 }
1779 assert!(
1780 CODE_GRAPH_TOOLS.contains(&"explore"),
1781 "and `explore` is the task tool the other surface lists"
1782 );
1783 }
1784
1785 #[test]
1788 fn tools_list_is_three_tools_on_a_code_graph_store() {
1789 let db = demo_db();
1790 db.write()
1791 .insert_node(
1792 "GitSync",
1793 crate::mcp_tasks::SYNC_KEY,
1794 vec![("id".into(), Value::Str(crate::mcp_tasks::SYNC_KEY.into()))],
1795 )
1796 .expect("marker");
1797 let resp = roundtrip(&db, r#"{"jsonrpc":"2.0","id":1,"method":"tools/list"}"#);
1798 let names: Vec<&str> = resp["result"]["tools"]
1799 .as_array()
1800 .expect("tools array")
1801 .iter()
1802 .map(|t| t["name"].as_str().expect("name"))
1803 .collect();
1804 assert_eq!(names, ["explore", "query", "stats"]);
1805 }
1806
1807 #[test]
1808 fn test_stats_returns_node_count() {
1809 let db = demo_db();
1810 let resp = tool_call(&db, 1, "stats", json!({}));
1811 assert!(!is_error(&resp));
1812 let result = tool_text(&resp);
1813 assert_eq!(result["nodes_live"], 2);
1814 }
1815
1816 #[test]
1820 fn create_rule_reports_a_build_it_could_not_finish() {
1821 let db = SharedDb::open(&tmp_dir()).expect("open");
1822 {
1823 let mut g = db.write();
1824 for i in 0..300usize {
1825 const D: usize = 32;
1826 let axis = (i / 10) % D;
1827 let mut xs = vec![0.0f64; D];
1828 xs[axis] = 1.0;
1829 xs[(axis + 1) % D] = (i % 10) as f64 * 0.001;
1830 g.insert_node(
1831 "V",
1832 &format!("v{i}"),
1833 vec![(
1834 "emb".into(),
1835 Value::List(xs.into_iter().map(Value::Float).collect()),
1836 )],
1837 )
1838 .expect("insert");
1839 }
1840 g.set_hnsw_build_batch(Some(64));
1841 }
1842 let args = json!({
1843 "name": "sim",
1844 "src_label": "V",
1845 "dst_label": "V",
1846 "predicate": {"VectorSimilar": {"field": "emb", "min": 0.9}},
1847 "edge_type": "SIM",
1848 "weight_prop": null,
1849 "max_edges": null,
1850 "approximate": true
1851 });
1852 let resp = tool_call(&db, 1, "create_rule", args);
1853 assert!(!is_error(&resp), "{resp}");
1854 let result = tool_text(&resp);
1855 assert_eq!(result["name"], json!("sim"));
1856 assert_eq!(result["building"], json!({"indexed": 64, "total": 300}));
1857 let note = result["note"].as_str().expect("a note explaining the wait");
1858 assert!(
1859 note.contains("derives no edges until it finishes") && note.contains("build-index"),
1860 "the note must say the edges are not there yet and how to finish: {note}"
1861 );
1862
1863 let stats = tool_text(&tool_call(&db, 2, "stats", json!({})));
1865 let rule = stats["rules"]
1866 .as_array()
1867 .expect("rules")
1868 .iter()
1869 .find(|r| r["name"] == "sim")
1870 .expect("the rule is installed while it builds");
1871 assert_eq!(rule["edges"], json!(0));
1872 assert_eq!(
1873 rule["building"],
1874 json!({"rule": "sim", "indexed": 64, "total": 300})
1875 );
1876
1877 while !db.write().pump_index_build().expect("pump").is_empty() {}
1879 let stats = tool_text(&tool_call(&db, 3, "stats", json!({})));
1880 let rule = stats["rules"]
1881 .as_array()
1882 .expect("rules")
1883 .iter()
1884 .find(|r| r["name"] == "sim")
1885 .expect("rule");
1886 assert!(rule.get("building").is_none(), "{rule}");
1887 assert!(rule["edges"].as_u64().expect("edges") > 0);
1888 }
1889
1890 #[test]
1891 fn test_query_runs_cypher() {
1892 let db = demo_db();
1893 let resp = tool_call(
1894 &db,
1895 1,
1896 "query",
1897 json!({ "cypher": "MATCH (n:Person) RETURN n.name ORDER BY n.name" }),
1898 );
1899 assert!(!is_error(&resp));
1900 let result = tool_text(&resp);
1901 assert_eq!(result["columns"], json!(["n.name"]));
1903 assert_eq!(result["rows"].as_array().map(|r| r.len()), Some(2));
1904 }
1905
1906 #[test]
1907 fn test_query_create_is_a_write() {
1908 let db = SharedDb::open(&tmp_dir()).expect("open");
1909 let resp = tool_call(
1910 &db,
1911 1,
1912 "query",
1913 json!({ "cypher": "CREATE (n:L {id: 'k'}) RETURN n" }),
1914 );
1915 assert!(
1916 !is_error(&resp),
1917 "CREATE via MCP query must succeed: {resp}"
1918 );
1919 let stats = tool_text(&tool_call(&db, 2, "stats", json!({})));
1920 assert_eq!(stats["nodes_live"], 1);
1921 }
1922
1923 #[test]
1924 fn test_ingest_json_inserts_nodes() {
1925 let db = demo_db();
1926 let resp = tool_call(
1927 &db,
1928 1,
1929 "ingest_json",
1930 json!({
1931 "label": "Person",
1932 "rows_json": r#"[{"id":"carol","name":"Carol"}]"#,
1933 "key_field": "id"
1934 }),
1935 );
1936 assert!(!is_error(&resp));
1937 let stats = tool_text(&tool_call(&db, 2, "stats", json!({})));
1939 assert_eq!(stats["nodes_live"], 3);
1940 }
1941
1942 #[test]
1943 fn test_node_info_returns_props() {
1944 let db = demo_db();
1945 let resp = tool_call(&db, 1, "node_info", json!({ "key": "alice" }));
1946 assert!(!is_error(&resp));
1947 let result = tool_text(&resp);
1948 assert_eq!(result["key"], "alice");
1949 assert_eq!(result["label"], "Person");
1950 assert_eq!(result["props"]["name"], "Alice");
1951 }
1952
1953 #[test]
1956 fn test_node_edges_returns_edges() {
1957 let db = demo_db();
1958 let resp = tool_call(
1959 &db,
1960 1,
1961 "node_edges",
1962 json!({ "key": "alice", "json": true }),
1963 );
1964 assert!(!is_error(&resp));
1965 let result = tool_text(&resp);
1966 assert_eq!(result["key"], "alice");
1967 let types = result["types"].as_array().expect("types");
1968 assert!(
1969 !types.is_empty(),
1970 "alice should have at least one edge type"
1971 );
1972 let similar = types
1973 .iter()
1974 .find(|t| t["edge_type"] == "SIMILAR")
1975 .expect("the rule's edge type");
1976 assert_eq!(similar["count"], json!(2));
1979 let edges = similar["edges"].as_array().expect("edges");
1980 let dirs: Vec<&str> = edges
1981 .iter()
1982 .map(|e| e["direction"].as_str().expect("direction"))
1983 .collect();
1984 assert!(dirs.contains(&"out") && dirs.contains(&"in"), "{similar}");
1985 for edge in edges {
1986 assert_eq!(edge["other"], json!("bob"));
1987 assert_eq!(edge["derived"], json!(true));
1988 assert_eq!(edge["rule"], json!("sim_emb"));
1989 assert_eq!(edge["score"], json!(1.0));
1990 assert!(
1991 edge["predicate"]
1992 .as_str()
1993 .unwrap_or("")
1994 .contains("vector_similar"),
1995 "the predicate travels with the edge: {edge}"
1996 );
1997 }
1998 }
1999
2000 #[test]
2003 fn test_neighborhood_traverses_one_hop() {
2004 let db = demo_db();
2005 let resp = tool_call(
2006 &db,
2007 1,
2008 "neighborhood",
2009 json!({ "key": "alice", "depth": 1, "json": true }),
2010 );
2011 assert!(!is_error(&resp));
2012 let result = tool_text(&resp);
2013 assert_eq!(result["key"], "alice");
2014 assert!(result["types"].as_array().is_some(), "{result}");
2015
2016 let deep = tool_call(
2017 &db,
2018 2,
2019 "neighborhood",
2020 json!({ "key": "alice", "depth": 2 }),
2021 );
2022 assert!(!is_error(&deep));
2023 let table = tool_text(&deep);
2024 assert_eq!(table["columns"], json!(["key", "label", "depth"]));
2025 assert!(table["rows"].as_array().is_some());
2026 }
2027
2028 #[test]
2029 fn test_explain_returns_rule_info() {
2030 let db = demo_db();
2031 let resp = tool_call(&db, 1, "explain", json!({ "a": "alice", "b": "bob" }));
2032 assert!(!is_error(&resp));
2033 let result = tool_text(&resp);
2034 let arr = result.as_array().expect("explain returns array");
2035 assert!(!arr.is_empty(), "expected at least one explanation");
2036 assert_eq!(arr[0]["rule"], "sim_emb");
2037 }
2038
2039 #[test]
2040 fn test_create_rule_backfills() {
2041 let db = SharedDb::open(&tmp_dir()).expect("open");
2042 {
2043 let mut g = db.write();
2044 let opts = IngestOptions {
2045 key_field: "id".into(),
2046 auto_fk: AutoFk::Off,
2047 };
2048 let rows: Vec<BTreeMap<String, Value>> = vec![
2049 [
2050 ("id", Value::Str("x".into())),
2051 ("tag", Value::Str("a".into())),
2052 ]
2053 .into_iter()
2054 .map(|(k, v)| (k.to_string(), v))
2055 .collect(),
2056 [
2057 ("id", Value::Str("y".into())),
2058 ("tag", Value::Str("a".into())),
2059 ]
2060 .into_iter()
2061 .map(|(k, v)| (k.to_string(), v))
2062 .collect(),
2063 ];
2064 g.ingest("Item", rows, &opts).expect("ingest");
2065 }
2066 let resp = tool_call(
2067 &db,
2068 1,
2069 "create_rule",
2070 json!({
2071 "name": "same_tag",
2072 "src_label": "Item",
2073 "dst_label": "Item",
2074 "predicate": { "FieldEqual": { "field": "tag" } },
2075 "edge_type": "SAME_TAG"
2076 }),
2077 );
2078 assert!(!is_error(&resp));
2079 let result = tool_text(&resp);
2080 assert_eq!(result["ok"], true);
2081 let edges_resp = tool_call(&db, 2, "node_edges", json!({ "key": "x", "json": true }));
2083 let edges_result = tool_text(&edges_resp);
2084 let types = edges_result["types"].as_array().expect("types");
2085 assert!(
2086 types.iter().any(|t| t["edge_type"] == "SAME_TAG"),
2087 "SAME_TAG edge not found after create_rule"
2088 );
2089 }
2090
2091 #[test]
2094 fn test_upsert_entity_creates_new_node() {
2095 let db = demo_db();
2096 let resp = tool_call(
2097 &db,
2098 1,
2099 "upsert_entity",
2100 json!({
2101 "key": "carol",
2102 "label": "Person",
2103 "props": { "name": "Carol", "age": 30 }
2104 }),
2105 );
2106 assert!(!is_error(&resp));
2107 let result = tool_text(&resp);
2108 assert_eq!(result["ok"], true);
2109 assert_eq!(result["created"], true);
2110 assert_eq!(result["key"], "carol");
2111 let info = tool_text(&tool_call(&db, 2, "node_info", json!({ "key": "carol" })));
2113 assert_eq!(info["props"]["name"], "Carol");
2114 }
2115
2116 #[test]
2117 fn test_upsert_entity_updates_existing_node() {
2118 let db = demo_db();
2119 let resp = tool_call(
2120 &db,
2121 1,
2122 "upsert_entity",
2123 json!({
2124 "key": "alice",
2125 "props": { "name": "Alice Updated" }
2126 }),
2127 );
2128 assert!(!is_error(&resp));
2129 let result = tool_text(&resp);
2130 assert_eq!(result["ok"], true);
2131 assert_eq!(result["created"], false);
2132 assert_eq!(result["updated_fields"], 1);
2133 let info = tool_text(&tool_call(&db, 2, "node_info", json!({ "key": "alice" })));
2135 assert_eq!(info["props"]["name"], "Alice Updated");
2136 }
2137
2138 #[test]
2139 fn test_upsert_entity_missing_label_on_create_is_error() {
2140 let db = demo_db();
2141 let resp = tool_call(
2142 &db,
2143 1,
2144 "upsert_entity",
2145 json!({ "key": "new-node", "props": { "x": 1 } }),
2146 );
2147 assert!(is_error(&resp), "should error without label for new node");
2148 }
2149
2150 #[test]
2151 fn test_find_similar_returns_similar_edges() {
2152 let db = demo_db();
2153 let resp = tool_call(
2154 &db,
2155 1,
2156 "find_similar",
2157 json!({ "key": "alice", "edge_type": "SIMILAR" }),
2158 );
2159 assert!(!is_error(&resp));
2160 let result = tool_text(&resp);
2161 assert_eq!(result["key"], "alice");
2162 assert_eq!(result["edge_type"], "SIMILAR");
2163 let similar = result["similar"].as_array().expect("similar array");
2164 assert!(!similar.is_empty(), "expected SIMILAR neighbors for alice");
2165 assert_eq!(similar[0]["neighbor_key"], "bob");
2166 }
2167
2168 #[test]
2169 fn test_find_similar_limit_respected() {
2170 let db = demo_db();
2171 let resp = tool_call(
2172 &db,
2173 1,
2174 "find_similar",
2175 json!({ "key": "alice", "edge_type": "SIMILAR", "limit": 0 }),
2176 );
2177 assert!(!is_error(&resp));
2178 let result = tool_text(&resp);
2179 let similar = result["similar"].as_array().expect("similar array");
2180 assert_eq!(similar.len(), 0);
2181 }
2182
2183 #[test]
2187 fn test_find_similar_vector_default_min_is_0_8() {
2188 let db = SharedDb::open(&tmp_dir()).expect("open");
2189 {
2190 let mut g = db.write();
2191 g.insert_node(
2193 "Item",
2194 "close",
2195 vec![(
2196 "emb".into(),
2197 Value::List(vec![Value::Float(1.0), Value::Float(0.0)]),
2198 )],
2199 )
2200 .unwrap();
2201 g.insert_node(
2203 "Item",
2204 "far",
2205 vec![(
2206 "emb".into(),
2207 Value::List(vec![Value::Float(0.0), Value::Float(1.0)]),
2208 )],
2209 )
2210 .unwrap();
2211 }
2212
2213 let resp = tool_call(
2215 &db,
2216 1,
2217 "find_similar",
2218 json!({
2219 "vector": [1.0, 0.0],
2220 "field": "emb",
2221 "label": "Item",
2222 "k": 10
2223 }),
2224 );
2225 assert!(!is_error(&resp), "vector search must not error");
2226 let result = tool_text(&resp);
2227 let results = result["results"].as_array().expect("results array");
2228
2229 let keys: Vec<&str> = results.iter().filter_map(|r| r["key"].as_str()).collect();
2230 assert!(
2231 keys.contains(&"close"),
2232 "close node (sim=1.0) must be included"
2233 );
2234 assert!(
2235 !keys.contains(&"far"),
2236 "far node (sim=0.0) must be excluded by default min=0.8"
2237 );
2238 }
2239
2240 #[test]
2242 fn test_find_similar_vector_mask_excludes_hidden() {
2243 let db = SharedDb::open(&tmp_dir()).expect("open");
2244 {
2245 let mut g = db.write();
2246 g.insert_node(
2248 "Item",
2249 "visible",
2250 vec![(
2251 "emb".into(),
2252 Value::List(vec![Value::Float(1.0), Value::Float(0.0)]),
2253 )],
2254 )
2255 .unwrap();
2256 g.insert_node(
2258 "Item",
2259 "hidden",
2260 vec![(
2261 "emb".into(),
2262 Value::List(vec![Value::Float(1.0), Value::Float(0.0)]),
2263 )],
2264 )
2265 .unwrap();
2266 }
2267
2268 let resp = tool_call(
2269 &db,
2270 1,
2271 "find_similar",
2272 json!({
2273 "vector": [1.0, 0.0],
2274 "field": "emb",
2275 "label": "Item",
2276 "k": 10,
2277 "min": 0.0,
2278 "mask": ["visible"]
2279 }),
2280 );
2281 assert!(!is_error(&resp), "masked vector search must not error");
2282 let result = tool_text(&resp);
2283 let results = result["results"].as_array().expect("results array");
2284
2285 let keys: Vec<&str> = results.iter().filter_map(|r| r["key"].as_str()).collect();
2286 assert!(
2287 keys.contains(&"visible"),
2288 "visible node must appear in masked results"
2289 );
2290 assert!(
2291 !keys.contains(&"hidden"),
2292 "hidden node must be excluded by mask"
2293 );
2294 }
2295
2296 #[test]
2298 fn test_find_similar_vector_mask_bad_type_is_error() {
2299 let db = SharedDb::open(&tmp_dir()).expect("open");
2300 let resp = tool_call(
2301 &db,
2302 1,
2303 "find_similar",
2304 json!({
2305 "vector": [1.0, 0.0],
2306 "field": "emb",
2307 "k": 5,
2308 "mask": [42]
2309 }),
2310 );
2311 assert!(
2312 is_error(&resp),
2313 "non-string mask element must produce a tool error"
2314 );
2315 }
2316
2317 #[test]
2319 fn test_find_similar_vector_where_eq() {
2320 let db = SharedDb::open(&tmp_dir()).expect("open");
2321 {
2322 let mut g = db.write();
2323 g.insert_node(
2324 "Document",
2325 "in-scope",
2326 vec![
2327 (
2328 "emb".into(),
2329 Value::List(vec![Value::Float(1.0), Value::Float(0.0)]),
2330 ),
2331 ("resource_scope_id".into(), Value::Str("a".into())),
2332 ],
2333 )
2334 .unwrap();
2335 g.insert_node(
2336 "Document",
2337 "out-scope",
2338 vec![
2339 (
2340 "emb".into(),
2341 Value::List(vec![Value::Float(1.0), Value::Float(0.0)]),
2342 ),
2343 ("resource_scope_id".into(), Value::Str("b".into())),
2344 ],
2345 )
2346 .unwrap();
2347 }
2348 let resp = tool_call(
2349 &db,
2350 1,
2351 "find_similar",
2352 json!({
2353 "vector": [1.0, 0.0],
2354 "field": "emb",
2355 "label": "Document",
2356 "k": 10,
2357 "min": 0.0,
2358 "where": { "field": "resource_scope_id", "eq": "a" }
2359 }),
2360 );
2361 assert!(!is_error(&resp), "where eq must not error: {resp:?}");
2362 let result = tool_text(&resp);
2363 let keys: Vec<&str> = result["results"]
2364 .as_array()
2365 .expect("results")
2366 .iter()
2367 .filter_map(|r| r["key"].as_str())
2368 .collect();
2369 assert_eq!(keys, vec!["in-scope"]);
2370 }
2371
2372 #[test]
2373 fn test_find_similar_vector_where_invalid_is_error() {
2374 let db = SharedDb::open(&tmp_dir()).expect("open");
2375 let resp = tool_call(
2376 &db,
2377 1,
2378 "find_similar",
2379 json!({
2380 "vector": [1.0, 0.0],
2381 "field": "emb",
2382 "where": { "field": "resource_scope_id", "eq": "a", "in": ["b"] }
2383 }),
2384 );
2385 assert!(is_error(&resp), "invalid where must be a tool error");
2386 let msg = format!("{resp:?}");
2387 assert!(
2388 msg.contains("where"),
2389 "tool error must name where, got {msg}"
2390 );
2391 }
2392
2393 #[test]
2395 fn test_find_similar_edge_ignores_where_and_exact() {
2396 let db = demo_db();
2397 let resp = tool_call(
2398 &db,
2399 1,
2400 "find_similar",
2401 json!({
2402 "key": "alice",
2403 "edge_type": "SIMILAR",
2404 "where": { "field": "x", "eq": "y", "in": ["z"] },
2405 "exact": true
2406 }),
2407 );
2408 assert!(
2409 !is_error(&resp),
2410 "edge mode must ignore invalid where: {resp:?}"
2411 );
2412 }
2413
2414 #[test]
2416 fn test_find_similar_edge_mask_excludes_hidden_neighbor() {
2417 let db = SharedDb::open(&tmp_dir()).expect("open");
2418 {
2419 let mut g = db.write();
2420 g.insert_node("P", "alice", vec![]).unwrap();
2421 g.insert_node("P", "bob", vec![]).unwrap(); g.insert_node("P", "carol", vec![]).unwrap(); g.insert_edge("KNOWS", "alice", "bob").unwrap();
2424 g.insert_edge("KNOWS", "alice", "carol").unwrap();
2425 }
2426 let resp = tool_call(
2428 &db,
2429 1,
2430 "find_similar",
2431 json!({
2432 "key": "alice",
2433 "edge_type": "KNOWS",
2434 "mask": ["alice", "bob"]
2435 }),
2436 );
2437 assert!(!is_error(&resp), "masked edge search must not error");
2438 let result = tool_text(&resp);
2439 let similar = result["similar"].as_array().expect("similar array");
2440 let neighbors: Vec<&str> = similar
2441 .iter()
2442 .filter_map(|e| e["neighbor_key"].as_str())
2443 .collect();
2444 assert!(neighbors.contains(&"bob"), "bob (visible) must appear");
2445 assert!(
2446 !neighbors.contains(&"carol"),
2447 "carol (hidden) must be excluded"
2448 );
2449 }
2450
2451 #[test]
2454 fn test_find_similar_edge_mask_hidden_key_is_not_found() {
2455 let db = SharedDb::open(&tmp_dir()).expect("open");
2456 {
2457 let mut g = db.write();
2458 g.insert_node("P", "alice", vec![]).unwrap();
2459 g.insert_node("P", "bob", vec![]).unwrap();
2460 }
2461 let resp_masked = tool_call(
2463 &db,
2464 1,
2465 "find_similar",
2466 json!({ "key": "alice", "edge_type": "KNOWS", "mask": ["bob"] }),
2467 );
2468 let resp_ghost = tool_call(
2470 &db,
2471 2,
2472 "find_similar",
2473 json!({ "key": "ghost", "edge_type": "KNOWS" }),
2474 );
2475 assert!(
2476 is_error(&resp_masked),
2477 "hidden query key must produce a tool error"
2478 );
2479 assert!(
2480 is_error(&resp_ghost),
2481 "nonexistent key must produce a tool error"
2482 );
2483 assert_eq!(
2485 tool_err_text(&resp_masked).contains("alice"),
2486 tool_err_text(&resp_ghost).contains("ghost"),
2487 "error messages should follow same not-found template"
2488 );
2489 }
2490
2491 #[test]
2496 fn test_explain_association_same_as_explain() {
2497 let db = demo_db();
2498 let explain = tool_text(&tool_call(
2499 &db,
2500 1,
2501 "explain",
2502 json!({ "a": "alice", "b": "bob" }),
2503 ));
2504 let assoc = tool_text(&tool_call(
2505 &db,
2506 2,
2507 "explain_association",
2508 json!({ "a": "alice", "b": "bob", "json": true }),
2509 ));
2510 let explain: Vec<Js> = serde_json::from_value(explain).expect("explain array");
2511 let mut assoc: Vec<Js> = serde_json::from_value(assoc).expect("assoc array");
2512 for row in &mut assoc {
2513 let ev = row
2514 .as_object_mut()
2515 .expect("object")
2516 .remove("evidence")
2517 .expect("every derived edge carries its evidence");
2518 assert!(
2519 ev["similarity"].is_number(),
2520 "a vector_similar edge reports the cosine it scored: {ev}"
2521 );
2522 }
2523 assert_eq!(explain, assoc, "evidence is the only addition");
2524
2525 let prose = tool_call(
2526 &db,
2527 3,
2528 "explain_association",
2529 json!({ "a": "alice", "b": "bob" }),
2530 );
2531 let text = prose["result"]["content"][0]["text"]
2532 .as_str()
2533 .expect("text content");
2534 assert!(
2535 text.contains("mushroomdb explain — alice ↔ bob:"),
2536 "the default reply is the digest: {text}"
2537 );
2538 }
2539
2540 #[test]
2545 fn test_edge_history_returns_derived_lifecycle_with_rule() {
2546 let db = demo_db(); let resp = tool_call(&db, 1, "edge_history", json!({ "a": "alice", "b": "bob" }));
2548 assert!(!is_error(&resp), "edge_history must not error: {resp}");
2549 let result = tool_text(&resp);
2550
2551 let total = result["total_commits"].as_u64().expect("total_commits");
2553 assert!(total > 0, "total_commits must be > 0 after ingest + rule");
2554
2555 let events = result["events"].as_array().expect("events array");
2557 assert!(!events.is_empty(), "expected at least one edge event");
2558
2559 let derived_added = events
2561 .iter()
2562 .any(|ev| ev["event"].as_str() == Some("Added") && !ev["rule"].is_null());
2563 assert!(
2564 derived_added,
2565 "expected a derived Added event with rule attribution: {events:?}"
2566 );
2567 }
2568
2569 #[test]
2572 fn test_was_linked_at_valid_commit() {
2573 let db = SharedDb::open(&tmp_dir()).expect("open");
2574 {
2575 let mut g = db.write();
2576 let opts = IngestOptions {
2577 key_field: "id".into(),
2578 auto_fk: AutoFk::Off,
2579 };
2580 let rows: Vec<BTreeMap<String, Value>> = vec![
2581 [("id", Value::Str("x".into()))]
2582 .into_iter()
2583 .map(|(k, v)| (k.to_string(), v))
2584 .collect(),
2585 [("id", Value::Str("y".into()))]
2586 .into_iter()
2587 .map(|(k, v)| (k.to_string(), v))
2588 .collect(),
2589 ];
2590 g.ingest("N", rows, &opts).expect("ingest");
2591 g.insert_edge("LINK", "x", "y").expect("edge");
2592 }
2593 let g = db.read();
2595 let total = g.wal_total_commits().expect("wal_total_commits");
2596 drop(g);
2597
2598 let resp = tool_call(
2599 &db,
2600 1,
2601 "was_linked",
2602 json!({ "a": "x", "b": "y", "edge_type": "LINK", "at_commit": total - 1 }),
2603 );
2604 assert!(!is_error(&resp), "was_linked must not error: {resp}");
2605 let result = tool_text(&resp);
2606 assert_eq!(result["linked"], true);
2607 assert_eq!(result["a"], "x");
2608 assert_eq!(result["edge_type"], "LINK");
2609 }
2610
2611 #[test]
2614 fn test_was_linked_out_of_horizon_returns_tool_error() {
2615 let db = SharedDb::open(&tmp_dir()).expect("open");
2616 {
2617 let mut g = db.write();
2618 g.insert_node("N", "a", vec![]).expect("node a");
2619 g.insert_node("N", "b", vec![]).expect("node b");
2620 }
2621 let resp = tool_call(
2623 &db,
2624 1,
2625 "was_linked",
2626 json!({ "a": "a", "b": "b", "edge_type": "X", "at_commit": 999 }),
2627 );
2628 assert!(
2630 is_error(&resp),
2631 "out-of-range commit must be a tool error: {resp}"
2632 );
2633 let text = resp["result"]["content"][0]["text"].as_str().expect("text");
2634 assert!(
2635 text.contains("out of range") || text.contains("range"),
2636 "error must mention range: {text}"
2637 );
2638 }
2639
2640 #[test]
2643 fn test_node_history_via_mcp() {
2644 let db = demo_db(); let resp = tool_call(&db, 1, "node_history", json!({ "key": "alice" }));
2646 assert!(!is_error(&resp), "node_history must not error: {resp}");
2647 let result = tool_text(&resp);
2648
2649 assert_eq!(result["key"], "alice");
2650 let total = result["total_commits"].as_u64().expect("total_commits");
2651 assert!(total > 0, "total_commits must be > 0");
2652
2653 let history = result["history"].as_array().expect("history array");
2654 assert!(
2655 !history.is_empty(),
2656 "alice should have at least one history entry"
2657 );
2658
2659 let first_change = &history[0]["change"];
2661 assert_eq!(first_change["type"], "NodeInserted");
2662 assert_eq!(first_change["label"], "Person");
2663 }
2664}