1use crate::json::{
52 edge_history_result_json, node_history_json, node_info_json, params_from_json,
53 parse_ingest_edges, result_set_json, rule_def_from_json,
54};
55use core_api::{
56 json_to_rows, json_to_value, AutoFk, GraphError, IngestOptions, MaskMode, NodeMask, SharedDb,
57 Value,
58};
59use serde_json::{json, Value as Js};
60use std::collections::BTreeMap;
61use std::io::{self, BufRead, Write};
62use std::path::{Path, PathBuf};
63
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),
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 if role.is_some() || mask_keys.is_some() {
277 let stub_hidden = args
278 .get("stub_hidden")
279 .and_then(|v| v.as_bool())
280 .unwrap_or(false);
281 let g = db.read();
282 let mask = match (role, &mask_keys) {
283 (Some(role), _) => match g.mask_for_role(role) {
284 Ok(m) => m,
285 Err(GraphError::KeyNotFound { .. }) => {
288 return CallOutcome::ToolErr(format!("unknown role '{role}'"))
289 }
290 Err(e) => return CallOutcome::ToolErr(graph_err_msg(e)),
291 },
292 (None, Some(keys)) => NodeMask::from_keys(&*g, keys.iter().map(String::as_str)),
293 (None, None) => unreachable!("one of the two is Some in this branch"),
294 };
295 let mask = if stub_hidden {
296 mask.with_mode(MaskMode::Stub)
297 } else {
298 mask
299 };
300 return match g.query_masked(cypher, ¶ms, &mask) {
301 Ok(rs) => CallOutcome::ToolOk(result_set_json(&rs)),
302 Err(e) => CallOutcome::ToolErr(graph_err_msg(e)),
303 };
304 }
305
306 let is_write = match core_api::is_write_query(cypher) {
307 Ok(b) => b,
308 Err(e) => return CallOutcome::ToolErr(e),
309 };
310 let rs = if is_write {
311 let mut g = db.write();
312 g.query_write(cypher, ¶ms)
313 } else {
314 let g = db.read();
315 g.query(cypher, ¶ms)
316 };
317 match rs {
318 Ok(rs) => CallOutcome::ToolOk(result_set_json(&rs)),
319 Err(e) => CallOutcome::ToolErr(graph_err_msg(e)),
320 }
321}
322
323fn mask_key_list(mask: &Js) -> Result<Vec<String>, String> {
326 let arr = mask
327 .as_array()
328 .ok_or_else(|| "mask must be an array of strings".to_string())?;
329 arr.iter()
330 .map(|v| {
331 v.as_str()
332 .map(str::to_string)
333 .ok_or_else(|| "mask must be an array of strings".to_string())
334 })
335 .collect()
336}
337
338fn tool_ingest(db: &SharedDb, args: &Js) -> CallOutcome {
339 let Some(label) = args.get("label").and_then(Js::as_str) else {
340 return CallOutcome::ToolErr("missing label".into());
341 };
342 let Some(rows_json) = args.get("rows_json").and_then(Js::as_str) else {
343 return CallOutcome::ToolErr("missing rows_json".into());
344 };
345 let mut opts = IngestOptions::default();
346 if let Some(kf) = args.get("key_field") {
347 match kf.as_str() {
348 Some(s) => opts.key_field = s.to_string(),
349 None => return CallOutcome::ToolErr("key_field must be a string".into()),
350 }
351 }
352 if let Some(suf) = args.get("auto_fk_suffix") {
353 match suf.as_str() {
354 Some(s) => {
355 opts.auto_fk = AutoFk::Auto {
356 suffix: s.to_string(),
357 }
358 }
359 None => return CallOutcome::ToolErr("auto_fk_suffix must be a string".into()),
360 }
361 }
362 let edges = match args.get("edges") {
363 None | Some(Js::Null) => Vec::new(),
364 Some(raw) => match parse_ingest_edges(raw) {
365 Ok(e) => e,
366 Err(e) => return CallOutcome::ToolErr(e),
367 },
368 };
369 let parsed: Js = match serde_json::from_str(rows_json) {
370 Ok(v) => v,
371 Err(e) => {
372 return CallOutcome::ToolErr(graph_err_msg(GraphError::IngestError {
373 detail: e.to_string(),
374 }))
375 }
376 };
377 let mut converted = match json_to_rows(&parsed) {
378 Ok(c) => c,
379 Err(e) => return CallOutcome::ToolErr(graph_err_msg(e)),
380 };
381 let taken = std::mem::take(&mut converted.rows);
382 let report = {
383 let mut g = db.write();
384 g.ingest_with_edges(label, taken, &opts, &edges)
385 };
386 match report.map(|r| converted.into_report(r)) {
387 Ok(r) => match serde_json::to_value(&r) {
388 Ok(v) => CallOutcome::ToolOk(v),
389 Err(e) => CallOutcome::ToolErr(e.to_string()),
390 },
391 Err(e) => CallOutcome::ToolErr(graph_err_msg(e)),
392 }
393}
394
395fn tool_create_rule(db: &SharedDb, args: &Js) -> CallOutcome {
396 let def = match rule_def_from_json(args.clone()) {
397 Ok(d) => d,
398 Err(e) => return CallOutcome::ToolErr(e),
399 };
400 let name = def.name.clone();
401 let res = {
402 let mut g = db.write();
403 g.create_rule(def)
404 };
405 match res {
406 Ok(()) => CallOutcome::ToolOk(json!({"ok": true, "name": name})),
407 Err(e) => CallOutcome::ToolErr(graph_err_msg(e)),
408 }
409}
410
411fn tool_explain(db: &SharedDb, args: &Js) -> CallOutcome {
412 let Some(a) = args.get("a").and_then(Js::as_str).filter(|s| !s.is_empty()) else {
413 return CallOutcome::ToolErr("missing a".into());
414 };
415 let Some(b) = args.get("b").and_then(Js::as_str).filter(|s| !s.is_empty()) else {
416 return CallOutcome::ToolErr("missing b".into());
417 };
418 let out = {
419 let g = db.read();
420 g.explain(a, b)
421 };
422 match out {
423 Ok(v) => match serde_json::to_value(&v) {
424 Ok(j) => CallOutcome::ToolOk(j),
425 Err(e) => CallOutcome::ToolErr(e.to_string()),
426 },
427 Err(e) => CallOutcome::ToolErr(graph_err_msg(e)),
428 }
429}
430
431fn tool_stats(db: &SharedDb) -> CallOutcome {
432 let snap = {
433 let g = db.read();
434 g.stats()
435 };
436 match serde_json::to_value(&snap) {
437 Ok(v) => CallOutcome::ToolOk(v),
438 Err(e) => CallOutcome::ToolErr(e.to_string()),
439 }
440}
441
442fn tool_node_info(db: &SharedDb, args: &Js) -> CallOutcome {
443 let Some(key) = args.get("key").and_then(Js::as_str) else {
444 return CallOutcome::ToolErr("missing key".into());
445 };
446 let info = {
447 let g = db.read();
448 g.node_info(key)
449 };
450 match info {
451 Some(info) => CallOutcome::ToolOk(node_info_json(&info)),
452 None => CallOutcome::ToolErr(graph_err_msg(GraphError::KeyNotFound {
453 key: key.to_string(),
454 })),
455 }
456}
457
458fn tool_upsert_entity(db: &SharedDb, args: &Js) -> CallOutcome {
466 let Some(key) = args.get("key").and_then(Js::as_str) else {
467 return CallOutcome::ToolErr("missing key".into());
468 };
469 let label_opt = args.get("label").and_then(Js::as_str);
470 let Some(props_obj) = args.get("props").and_then(Js::as_object) else {
471 return CallOutcome::ToolErr("missing props".into());
472 };
473
474 let exists = {
475 let g = db.read();
476 g.has_node(key)
477 };
478
479 if exists {
480 let mut g = db.write();
481 let mut count = 0usize;
482 for (field, json_val) in props_obj {
483 match json_to_value(json_val.clone()) {
484 Some(v) => {
485 if let Err(e) = g.set_prop(key, field, v) {
486 return CallOutcome::ToolErr(graph_err_msg(e));
487 }
488 count += 1;
489 }
490 None => {
491 return CallOutcome::ToolErr(format!(
492 "prop {field} is not a supported value type"
493 ))
494 }
495 }
496 }
497 CallOutcome::ToolOk(json!({
498 "ok": true,
499 "key": key,
500 "created": false,
501 "updated_fields": count
502 }))
503 } else {
504 let Some(label) = label_opt else {
505 return CallOutcome::ToolErr("label required when creating a new entity".into());
506 };
507 let mut row: BTreeMap<String, Value> = BTreeMap::new();
508 row.insert("id".to_string(), Value::Str(key.to_string()));
509 for (field, json_val) in props_obj {
510 if field == "id" {
511 continue;
512 }
513 match json_to_value(json_val.clone()) {
514 Some(v) => {
515 row.insert(field.clone(), v);
516 }
517 None => {
518 return CallOutcome::ToolErr(format!(
519 "prop {field} is not a supported value type"
520 ))
521 }
522 }
523 }
524 let opts = IngestOptions {
525 key_field: "id".to_string(),
526 auto_fk: AutoFk::Off,
527 };
528 let mut g = db.write();
529 match g.ingest(label, vec![row], &opts) {
530 Ok(_) => CallOutcome::ToolOk(json!({ "ok": true, "key": key, "created": true })),
531 Err(e) => CallOutcome::ToolErr(graph_err_msg(e)),
532 }
533 }
534}
535
536fn tool_find_similar(db: &SharedDb, args: &Js) -> CallOutcome {
543 let mask_keys: Option<Vec<String>> = if let Some(mask_val) = args.get("mask") {
546 match mask_val.as_array() {
547 Some(arr) => {
548 let mut ks: Vec<String> = Vec::with_capacity(arr.len());
549 for v in arr {
550 match v.as_str() {
551 Some(s) => ks.push(s.to_string()),
552 None => {
553 return CallOutcome::ToolErr("mask must be an array of strings".into())
554 }
555 }
556 }
557 Some(ks)
558 }
559 None => return CallOutcome::ToolErr("mask must be an array of strings".into()),
560 }
561 } else {
562 None
563 };
564
565 if let Some(vec_js) = args.get("vector").and_then(Js::as_array) {
568 let q: Vec<f64> = vec_js.iter().filter_map(|v| v.as_f64()).collect();
569 if q.is_empty() {
570 return CallOutcome::ToolErr("vector must be a non-empty array of numbers".into());
571 }
572 let field = args
573 .get("field")
574 .and_then(Js::as_str)
575 .unwrap_or("embedding");
576 let label_str = args.get("label").and_then(Js::as_str).unwrap_or("");
577 let label = if label_str.is_empty() {
578 None
579 } else {
580 Some(label_str)
581 };
582 let k = args
583 .get("k")
584 .and_then(Js::as_u64)
585 .map(|n| n as usize)
586 .unwrap_or(10);
587 let min = args.get("min").and_then(Js::as_f64).unwrap_or(0.8);
588
589 let hits = {
590 let g = db.read();
591 if let Some(ref keys) = mask_keys {
592 let node_mask = NodeMask::from_keys(&*g, keys.iter().map(String::as_str));
593 g.find_similar_vector_masked(field, label, &q, k, min, &node_mask)
594 } else {
595 g.find_similar_vector(field, label, &q, k, min)
596 }
597 };
598 let results: Vec<Js> = hits
599 .into_iter()
600 .map(|(key, score)| json!({ "key": key, "score": score }))
601 .collect();
602 return CallOutcome::ToolOk(json!({
603 "mode": "vector",
604 "field": field,
605 "label": label,
606 "k": k,
607 "min": min,
608 "results": results
609 }));
610 }
611
612 let Some(key) = args.get("key").and_then(Js::as_str) else {
614 return CallOutcome::ToolErr("missing key (or provide vector for vector search)".into());
615 };
616 let edge_type = args
617 .get("edge_type")
618 .and_then(Js::as_str)
619 .unwrap_or("SIMILAR");
620 let limit = args
621 .get("limit")
622 .and_then(Js::as_u64)
623 .map(|n| n as usize)
624 .unwrap_or(10);
625
626 if let Some(ref mask) = mask_keys {
629 let mask_set: std::collections::HashSet<&str> = mask.iter().map(String::as_str).collect();
630 if !mask_set.contains(key) {
631 return CallOutcome::ToolErr(graph_err_msg(GraphError::KeyNotFound {
632 key: key.into(),
633 }));
634 }
635 let out = {
636 let g = db.read();
637 g.node_edges(key)
638 };
639 return match out {
640 Ok(edges) => {
641 let similar: Vec<Js> = edges
642 .iter()
643 .filter(|e| e.edge_type == edge_type)
644 .filter(|e| {
645 let neighbor_key = if e.src_key == key {
647 &e.dst_key
648 } else {
649 &e.src_key
650 };
651 mask_set.contains(neighbor_key.as_str())
652 })
653 .take(limit)
654 .map(|e| {
655 let neighbor_key = if e.src_key == key {
656 &e.dst_key
657 } else {
658 &e.src_key
659 };
660 let direction = if e.src_key == key { "out" } else { "in" };
661 json!({
662 "neighbor_key": neighbor_key,
663 "direction": direction,
664 "edge_type": e.edge_type,
665 "derived": e.derived,
666 })
667 })
668 .collect();
669 CallOutcome::ToolOk(json!({
670 "key": key,
671 "edge_type": edge_type,
672 "similar": similar
673 }))
674 }
675 Err(e) => CallOutcome::ToolErr(graph_err_msg(e)),
676 };
677 }
678
679 let out = {
680 let g = db.read();
681 g.node_edges(key)
682 };
683 match out {
684 Ok(edges) => {
685 let similar: Vec<Js> = edges
686 .iter()
687 .filter(|e| e.edge_type == edge_type)
688 .take(limit)
689 .map(|e| {
690 let neighbor_key = if e.src_key == key {
691 &e.dst_key
692 } else {
693 &e.src_key
694 };
695 let direction = if e.src_key == key { "out" } else { "in" };
696 json!({
697 "neighbor_key": neighbor_key,
698 "direction": direction,
699 "edge_type": e.edge_type,
700 "derived": e.derived,
701 })
702 })
703 .collect();
704 CallOutcome::ToolOk(json!({
705 "key": key,
706 "edge_type": edge_type,
707 "similar": similar
708 }))
709 }
710 Err(e) => CallOutcome::ToolErr(graph_err_msg(e)),
711 }
712}
713
714fn tool_hybrid_search(db: &SharedDb, args: &Js) -> CallOutcome {
715 let Some(query_text) = args.get("query_text").and_then(Js::as_str) else {
716 return CallOutcome::ToolErr("missing required field: query_text".into());
717 };
718 let Some(text_field) = args.get("text_field").and_then(Js::as_str) else {
719 return CallOutcome::ToolErr("missing required field: text_field".into());
720 };
721
722 let vector_field = args
723 .get("vector_field")
724 .and_then(Js::as_str)
725 .unwrap_or("embedding");
726 let label = args.get("label").and_then(Js::as_str);
727 let k = args
728 .get("k")
729 .and_then(Js::as_u64)
730 .map(|n| n as usize)
731 .unwrap_or(10);
732
733 let query_vec: Vec<f64> = args
734 .get("vector")
735 .and_then(Js::as_array)
736 .map(|arr| arr.iter().filter_map(|v| v.as_f64()).collect())
737 .unwrap_or_default();
738
739 let hits = {
740 let g = db.read();
741 g.search_hybrid(text_field, query_text, vector_field, &query_vec, label, k)
742 };
743
744 let results: Vec<Js> = hits
745 .into_iter()
746 .map(|(key, score)| json!({ "key": key, "score": score }))
747 .collect();
748
749 CallOutcome::ToolOk(json!({
750 "query_text": query_text,
751 "text_field": text_field,
752 "vector_field": vector_field,
753 "label": label,
754 "k": k,
755 "results": results
756 }))
757}
758
759fn tool_node_history(db: &SharedDb, args: &Js) -> CallOutcome {
760 let Some(key) = args.get("key").and_then(Js::as_str) else {
761 return CallOutcome::ToolErr("missing key".into());
762 };
763 let g = db.read();
764 let entries = match g.node_history(key) {
765 Ok(e) => e,
766 Err(e) => return CallOutcome::ToolErr(graph_err_msg(e)),
767 };
768 let total_commits = match g.wal_total_commits() {
769 Ok(n) => n,
770 Err(e) => return CallOutcome::ToolErr(graph_err_msg(e)),
771 };
772 CallOutcome::ToolOk(node_history_json(key, &entries, total_commits))
773}
774
775fn tool_edge_history(db: &SharedDb, args: &Js) -> CallOutcome {
776 let Some(a) = args.get("a").and_then(Js::as_str).filter(|s| !s.is_empty()) else {
777 return CallOutcome::ToolErr("missing a".into());
778 };
779 let Some(b) = args.get("b").and_then(Js::as_str).filter(|s| !s.is_empty()) else {
780 return CallOutcome::ToolErr("missing b".into());
781 };
782 let result = {
783 let g = db.read();
784 g.edge_history(a, b)
785 };
786 match result {
787 Ok(hr) => CallOutcome::ToolOk(edge_history_result_json(a, b, &hr)),
788 Err(e) => CallOutcome::ToolErr(graph_err_msg(e)),
789 }
790}
791
792fn tool_was_linked(db: &SharedDb, args: &Js) -> CallOutcome {
793 let Some(a) = args.get("a").and_then(Js::as_str).filter(|s| !s.is_empty()) else {
794 return CallOutcome::ToolErr("missing a".into());
795 };
796 let Some(b) = args.get("b").and_then(Js::as_str).filter(|s| !s.is_empty()) else {
797 return CallOutcome::ToolErr("missing b".into());
798 };
799 let Some(edge_type) = args
800 .get("edge_type")
801 .and_then(Js::as_str)
802 .filter(|s| !s.is_empty())
803 else {
804 return CallOutcome::ToolErr("missing edge_type".into());
805 };
806 let at_commit = match args.get("at_commit").and_then(Js::as_u64) {
807 Some(n) => n,
808 None => return CallOutcome::ToolErr("missing or invalid at_commit".into()),
809 };
810 let result = {
811 let g = db.read();
812 g.was_linked(a, b, edge_type, at_commit)
813 };
814 match result {
815 Ok(linked) => CallOutcome::ToolOk(json!({
816 "a": a,
817 "b": b,
818 "edge_type": edge_type,
819 "at_commit": at_commit,
820 "linked": linked,
821 })),
822 Err(e) => CallOutcome::ToolErr(graph_err_msg(e)),
823 }
824}
825
826fn tool_rename_node(db: &SharedDb, args: &Js) -> CallOutcome {
827 let Some(old_key) = args.get("old_key").and_then(Js::as_str) else {
828 return CallOutcome::ToolErr("missing old_key".into());
829 };
830 let Some(new_key) = args.get("new_key").and_then(Js::as_str) else {
831 return CallOutcome::ToolErr("missing new_key".into());
832 };
833 let mut g = db.write();
834 match g.rename_node(old_key, new_key) {
835 Ok(()) => CallOutcome::ToolOk(json!({
836 "ok": true,
837 "old_key": old_key,
838 "new_key": new_key,
839 })),
840 Err(e) => CallOutcome::ToolErr(graph_err_msg(e)),
841 }
842}
843
844pub(crate) fn graph_err_msg(e: GraphError) -> String {
845 match e {
846 GraphError::QueryError { detail } | GraphError::IngestError { detail } => detail,
847 other => other.to_string(),
848 }
849}
850
851fn initialize_result() -> Js {
852 json!({
853 "protocolVersion": "2024-11-05",
854 "capabilities": { "tools": {} },
855 "serverInfo": { "name": "mushroomdb", "version": env!("CARGO_PKG_VERSION") }
856 })
857}
858
859const ADVANCED_PREFIX: &str = "Advanced: ";
865
866pub const CODE_GRAPH_TOOLS: [&str; 3] = ["explore", "query", "stats"];
874
875pub const ASSOCIATION_TOOLS: [&str; 15] = [
900 "query",
901 "explain_association",
902 "neighborhood",
903 "node_info",
904 "node_edges",
905 "was_linked",
906 "edges_at",
907 "what_if",
908 "node_history",
909 "edge_history",
910 "find_similar",
911 "hybrid_search",
912 "remember",
913 "recall",
914 "stats",
915];
916
917#[derive(Clone, Copy, Debug, PartialEq, Eq)]
923pub(crate) enum Surface {
924 CodeGraph,
927 Memory,
930}
931
932impl Surface {
933 fn listing(self) -> &'static [&'static str] {
936 match self {
937 Surface::CodeGraph => &CODE_GRAPH_TOOLS,
938 Surface::Memory => &ASSOCIATION_TOOLS,
939 }
940 }
941}
942
943fn surface_of(db: &SharedDb) -> Surface {
946 let ingested = {
947 let g = db.read();
948 g.has_node(crate::mcp_tasks::SYNC_KEY)
949 };
950 if ingested {
951 Surface::CodeGraph
952 } else {
953 Surface::Memory
954 }
955}
956
957fn tools_list(all: bool, surface: Surface) -> Js {
975 let mut served: Vec<Js> = crate::mcp_tasks::task_tools();
976 for mut tool in graph_tools() {
977 if let Some(d) = tool.get("description").and_then(Js::as_str) {
978 let prefixed = format!("{ADVANCED_PREFIX}{d}");
979 tool["description"] = Js::String(prefixed);
980 }
981 served.push(tool);
982 }
983 if all {
984 return json!({ "tools": served });
985 }
986 let listing = surface.listing();
987 let mut tools: Vec<Js> = Vec::with_capacity(listing.len());
988 for name in listing {
989 let Some(tool) = served
990 .iter()
991 .find(|t| t.get("name").and_then(Js::as_str) == Some(*name))
992 else {
993 debug_assert!(false, "{surface:?} lists {name}, which is not served");
994 continue;
995 };
996 tools.push(tool.clone());
997 }
998 json!({ "tools": tools })
999}
1000
1001fn graph_tools() -> Vec<Js> {
1004 let Js::Array(tools) = json!([
1005 {
1006 "name": "query",
1007 "description": "Who may see this, and anything else one pattern can answer — run a Cypher query (read or write) against the graph. Pass 'role' to answer as one of the store's roles: only the nodes that role may see, writes refused. 'mask' is the same restriction written out as an explicit key allow-list. Cypher dialect: MATCH/WHERE/RETURN, CREATE, MERGE, SET, DELETE, with $named parameters in 'params'. A node's key and label read as properties (n.key, n.label) or as key(n)/labels(n). One MATCH takes comma-separated patterns that share variables — MATCH (t)-[:A]->(c), (t)-[:B]->(c) is the intersection of both, and count(DISTINCT t) after WITH counts each t once. WHERE takes STARTS WITH, ENDS WITH, CONTAINS, IN, and a list subscript (n.location[0]) — which is null when the index is out of range, the property is not a list, or the index is not an integer, so a subscript never errors and never matches.",
1008 "inputSchema": {
1009 "type": "object",
1010 "properties": {
1011 "cypher": { "type": "string", "description": "Cypher query text." },
1012 "params": {
1013 "type": "object",
1014 "description": "Named JSON-scalar query parameters."
1015 },
1016 "mask": {
1017 "type": "array",
1018 "items": { "type": "string" },
1019 "description": "Optional node key allow-list. When present, only these nodes are visible; write statements are rejected."
1020 },
1021 "role": {
1022 "type": "string",
1023 "description": "Answer as this role from the store's roles: only the nodes it may see."
1024 }
1025 },
1026 "required": ["cypher"]
1027 }
1028 },
1029 {
1030 "name": "ingest_json",
1031 "description": "Ingest a JSON array of objects as nodes of one label.",
1032 "inputSchema": {
1033 "type": "object",
1034 "properties": {
1035 "label": { "type": "string" },
1036 "rows_json": {
1037 "type": "string",
1038 "description": "JSON text of an array of objects."
1039 },
1040 "key_field": { "type": "string" },
1041 "auto_fk_suffix": { "type": "string" },
1042 "edges": {
1043 "type": "array",
1044 "description": "Optional user edges [{edge_type, src, dst}]."
1045 }
1046 },
1047 "required": ["label", "rows_json"]
1048 }
1049 },
1050 {
1051 "name": "create_rule",
1052 "description": "How should this kind of relationship be derived from now on — declare a rule (RuleDef JSON) and the engine maintains its edges as the data changes. Propose it and show the edges it would derive before creating one.",
1053 "inputSchema": {
1054 "type": "object",
1055 "properties": {
1056 "name": { "type": "string" },
1057 "src_label": { "type": "string" },
1058 "dst_label": { "type": "string" },
1059 "predicate": { "type": "object" },
1060 "edge_type": { "type": "string" },
1061 "weight_prop": {
1062 "type": ["string", "null"],
1063 "description": "Edge property that stores the score (default: weight)."
1064 },
1065 "max_edges": { "type": ["integer", "null"] }
1066 },
1067 "required": ["name", "src_label", "dst_label", "predicate", "edge_type"]
1068 }
1069 },
1070 {
1071 "name": "explain",
1072 "description": "Why are A and B related, as a raw array — the same rule-derived edges explain_association renders, for a caller that wants the JSON without asking.",
1073 "inputSchema": {
1074 "type": "object",
1075 "properties": {
1076 "a": { "type": "string", "minLength": 1 },
1077 "b": { "type": "string", "minLength": 1 }
1078 },
1079 "required": ["a", "b"]
1080 }
1081 },
1082 {
1083 "name": "stats",
1084 "description": "How big is this store — live node, edge and rule counts.",
1085 "inputSchema": {
1086 "type": "object",
1087 "properties": {}
1088 }
1089 },
1090 {
1091 "name": "node_info",
1092 "description": "What is K — its label and every property it holds.",
1093 "inputSchema": {
1094 "type": "object",
1095 "properties": {
1096 "key": { "type": "string" }
1097 },
1098 "required": ["key"]
1099 }
1100 },
1101 {
1102 "name": "upsert_entity",
1103 "description": "Record what is now true about K — insert or update a node by key. If the key exists, updates the supplied properties. If not, creates a new node with the given label and properties. Useful for agent memory: store or refresh an entity without checking existence first.",
1104 "inputSchema": {
1105 "type": "object",
1106 "properties": {
1107 "key": { "type": "string", "description": "Unique node key." },
1108 "label": { "type": "string", "description": "Node label (required when creating a new entity)." },
1109 "props": {
1110 "type": "object",
1111 "description": "Properties to set. Values must be scalars (string, number, bool) or arrays of scalars."
1112 }
1113 },
1114 "required": ["key", "props"]
1115 }
1116 },
1117 {
1118 "name": "find_similar",
1119 "description": "What is most like this — two modes: (1) Vector search — provide `vector` (and optionally `field`, `label`, `k`, `min`) to find the k most similar nodes by cosine similarity using the HNSW index when available, brute-force otherwise. (2) Edge traversal — provide `key` (and optionally `edge_type`, `limit`) to return neighbors previously connected by a derived rule edge. Results from mode 2 come only from edges already derived by a VectorSimilar rule. In both modes, the optional `mask` array limits visibility: hidden nodes never appear in results, and a hidden query key in edge mode behaves identically to a nonexistent key.",
1120 "inputSchema": {
1121 "type": "object",
1122 "properties": {
1123 "vector": {
1124 "type": "array",
1125 "items": { "type": "number" },
1126 "description": "Query embedding vector for vector-similarity search. When present, vector-search mode is used and `key` is ignored."
1127 },
1128 "field": { "type": "string", "description": "Property field holding the embedding vectors (default: embedding). Used in vector-search mode." },
1129 "label": { "type": "string", "description": "Restrict search to nodes with this label. Empty string means all labels. Used in vector-search mode." },
1130 "k": { "type": "integer", "description": "Maximum results to return in vector-search mode (default: 10)." },
1131 "min": { "type": "number", "description": "Minimum cosine similarity threshold in vector-search mode (default: 0.8)." },
1132 "mask": {
1133 "type": "array",
1134 "items": { "type": "string" },
1135 "description": "Optional node key allow-list for vector-search mode. When present, only nodes whose key appears in this list are eligible for results. Hidden nodes are excluded before k-truncation so callers still receive up to k visible hits. Unknown keys are silently ignored."
1136 },
1137 "key": { "type": "string", "description": "Source node key for edge-traversal mode." },
1138 "edge_type": { "type": "string", "description": "Edge type to filter by in edge-traversal mode (default: SIMILAR)." },
1139 "limit": { "type": "integer", "description": "Maximum neighbors to return in edge-traversal mode (default: 10)." }
1140 }
1141 }
1142 },
1143 {
1144 "name": "hybrid_search",
1145 "description": "What matches these words and this vector at once — Reciprocal Rank Fusion (RRF) over fulltext + vector results. Provide `query_text` and `text_field` for the fulltext leg. Optionally provide `vector` (embedding array) and `vector_field` (default: embedding) for the vector leg; omitting `vector` gives text-only ranking through the same RRF path. `label` restricts the vector search to nodes with that label (required for brute-force; omit to rely on HNSW rules). `k` controls result count (default: 10). RRF constant is fixed at 60; scores are 1/(60+rank) summed over lists a node appears in.",
1146 "inputSchema": {
1147 "type": "object",
1148 "properties": {
1149 "query_text": { "type": "string", "description": "Fulltext query string." },
1150 "text_field": { "type": "string", "description": "Property field to search with fulltext." },
1151 "vector": {
1152 "type": "array",
1153 "items": { "type": "number" },
1154 "description": "Query embedding vector. Omit for text-only ranking."
1155 },
1156 "vector_field": { "type": "string", "description": "Property field holding embedding vectors (default: embedding)." },
1157 "label": { "type": "string", "description": "Restrict vector search to nodes with this label. Required when relying on brute-force (no HNSW rule covers the field). If omitted, the vector leg always returns empty results (no rule-created HNSW index covers the unlabeled path); ranking is text-only in that case." },
1158 "k": { "type": "integer", "description": "Maximum results to return (default: 10)." }
1159 },
1160 "required": ["query_text", "text_field"]
1161 }
1162 },
1163 {
1164 "name": "node_history",
1165 "description": "What has happened to K — every recorded change to one node, newest last. Events include NodeInserted, PropSet, PropRemoved, EdgeAdded, EdgeRemoved, and NodeDeleted. The response includes `total_commits` (the horizon upper bound). History is WAL-scoped — pre-snapshot commits are not visible.",
1166 "inputSchema": {
1167 "type": "object",
1168 "properties": {
1169 "key": { "type": "string", "description": "Node key to look up." }
1170 },
1171 "required": ["key"]
1172 }
1173 },
1174 {
1175 "name": "edge_history",
1176 "description": "When did A and B become linked, and when did it break — the full add/retract lifecycle for every edge between the two keys. Includes derived (rule-attributed) edges via DerivedEdgeAdded/DerivedEdgeRetracted WAL markers. The response includes `total_commits` (the horizon upper bound).",
1177 "inputSchema": {
1178 "type": "object",
1179 "properties": {
1180 "a": { "type": "string", "minLength": 1, "description": "First node key." },
1181 "b": { "type": "string", "minLength": 1, "description": "Second node key." }
1182 },
1183 "required": ["a", "b"]
1184 }
1185 },
1186 {
1187 "name": "was_linked",
1188 "description": "Were A and B linked at commit C — whether an edge of `edge_type` existed between the two keys (either direction) at that WAL commit. Returns an error when `at_commit` is outside the visible horizon (`0..total_commits`).",
1189 "inputSchema": {
1190 "type": "object",
1191 "properties": {
1192 "a": { "type": "string", "minLength": 1, "description": "First node key." },
1193 "b": { "type": "string", "minLength": 1, "description": "Second node key." },
1194 "edge_type": { "type": "string", "minLength": 1, "description": "Edge type to check." },
1195 "at_commit": { "type": "integer", "minimum": 0, "description": "0-based WAL commit index to query." }
1196 },
1197 "required": ["a", "b", "edge_type", "at_commit"]
1198 }
1199 },
1200 {
1201 "name": "rename_node",
1202 "description": "Rename K — the key changes and nothing else does. The dense id and all edges/properties remain stable. Returns 404 if `old_key` does not exist, 409 if `new_key` is already taken.",
1203 "inputSchema": {
1204 "type": "object",
1205 "properties": {
1206 "old_key": { "type": "string", "minLength": 1, "description": "Current node key." },
1207 "new_key": { "type": "string", "minLength": 1, "description": "Desired new node key." }
1208 },
1209 "required": ["old_key", "new_key"]
1210 }
1211 }
1212 ]) else {
1213 unreachable!("the literal above is an array")
1214 };
1215 tools
1216}
1217
1218fn tool_ok(payload: Js) -> Js {
1219 json!({
1220 "content": [{ "type": "text", "text": payload.to_string() }]
1221 })
1222}
1223
1224fn task_ok(text: &str) -> Js {
1232 json!({
1233 "content": [{ "type": "text", "text": text }]
1234 })
1235}
1236
1237fn tool_err(message: &str) -> Js {
1238 json!({
1239 "content": [{ "type": "text", "text": message }],
1240 "isError": true
1241 })
1242}
1243
1244fn write_result(writer: &mut impl Write, id: Option<Js>, result: Js) -> io::Result<()> {
1245 write_json(
1246 writer,
1247 &json!({
1248 "jsonrpc": "2.0",
1249 "id": id.unwrap_or(Js::Null),
1250 "result": result
1251 }),
1252 )
1253}
1254
1255fn write_error(
1256 writer: &mut impl Write,
1257 id: Option<Js>,
1258 code: i64,
1259 message: &str,
1260) -> io::Result<()> {
1261 write_json(
1262 writer,
1263 &json!({
1264 "jsonrpc": "2.0",
1265 "id": id.unwrap_or(Js::Null),
1266 "error": { "code": code, "message": message }
1267 }),
1268 )
1269}
1270
1271fn write_json(writer: &mut impl Write, value: &Js) -> io::Result<()> {
1272 let s = serde_json::to_string(value).map_err(io::Error::other)?;
1273 writeln!(writer, "{s}")?;
1274 writer.flush()
1275}
1276
1277#[cfg(test)]
1282mod tests {
1283 use super::*;
1284 use core_api::{AutoFk, IngestOptions, Predicate, RuleDef, Value};
1285 use std::path::PathBuf;
1286 use std::sync::atomic::{AtomicU64, Ordering};
1287
1288 fn tmp_dir() -> PathBuf {
1289 static SEQ: AtomicU64 = AtomicU64::new(0);
1290 let n = SEQ.fetch_add(1, Ordering::Relaxed);
1291 let d = std::env::temp_dir().join(format!("mcp-test-{}-{}", std::process::id(), n));
1292 let _ = std::fs::remove_dir_all(&d);
1296 d
1297 }
1298
1299 fn demo_db() -> SharedDb {
1301 let db = SharedDb::open(&tmp_dir()).expect("open");
1302 {
1303 let mut g = db.write();
1304 let opts = IngestOptions {
1305 key_field: "id".into(),
1306 auto_fk: AutoFk::Off,
1307 };
1308 let people: Vec<BTreeMap<String, Value>> = vec![
1310 [
1311 ("id", Value::Str("alice".into())),
1312 ("name", Value::Str("Alice".into())),
1313 (
1314 "emb",
1315 Value::List(vec![Value::Float(1.0), Value::Float(0.0)]),
1316 ),
1317 ]
1318 .into_iter()
1319 .map(|(k, v)| (k.to_string(), v))
1320 .collect(),
1321 [
1322 ("id", Value::Str("bob".into())),
1323 ("name", Value::Str("Bob".into())),
1324 (
1325 "emb",
1326 Value::List(vec![Value::Float(1.0), Value::Float(0.0)]),
1327 ),
1328 ]
1329 .into_iter()
1330 .map(|(k, v)| (k.to_string(), v))
1331 .collect(),
1332 ];
1333 g.ingest("Person", people, &opts).expect("ingest");
1334
1335 g.create_rule(RuleDef {
1337 name: "sim_emb".into(),
1338 src_label: "Person".into(),
1339 dst_label: "Person".into(),
1340 predicate: Predicate::VectorSimilar {
1341 field: "emb".into(),
1342 min: 0.9,
1343 },
1344 edge_type: "SIMILAR".into(),
1345 weight_prop: Some("score".into()),
1346 max_edges: None,
1347 approximate: false,
1348 via_label: None,
1349 via_edge: None,
1350 via_dir: None,
1351 })
1352 .expect("rule");
1353 }
1354 db
1355 }
1356
1357 fn roundtrip(db: &SharedDb, request: &str) -> Js {
1358 roundtrip_with(db, false, request)
1359 }
1360
1361 fn roundtrip_with(db: &SharedDb, all_tools: bool, request: &str) -> Js {
1362 let input = format!("{request}\n");
1363 let mut output = Vec::new();
1364 run_mcp_stdio_with(db.clone(), None, all_tools, input.as_bytes(), &mut output)
1365 .expect("mcp");
1366 let s = std::str::from_utf8(&output).expect("utf8");
1367 serde_json::from_str(s.trim()).expect("json response")
1368 }
1369
1370 fn tool_call(db: &SharedDb, id: u64, tool: &str, args: Js) -> Js {
1371 let req = json!({
1372 "jsonrpc": "2.0",
1373 "id": id,
1374 "method": "tools/call",
1375 "params": { "name": tool, "arguments": args }
1376 });
1377 roundtrip(db, &req.to_string())
1378 }
1379
1380 fn tool_text(resp: &Js) -> Js {
1382 let text = resp["result"]["content"][0]["text"]
1383 .as_str()
1384 .expect("content[0].text");
1385 serde_json::from_str(text).expect("tool text is json")
1386 }
1387
1388 fn is_error(resp: &Js) -> bool {
1389 resp["result"]["isError"].as_bool().unwrap_or(false)
1390 }
1391
1392 fn tool_err_text(resp: &Js) -> String {
1393 resp["result"]["content"][0]["text"]
1394 .as_str()
1395 .unwrap_or("")
1396 .to_string()
1397 }
1398
1399 #[test]
1402 fn test_tools_list_includes_all_expected() {
1403 let db = demo_db();
1404 let resp = roundtrip_with(
1405 &db,
1406 true,
1407 r#"{"jsonrpc":"2.0","id":1,"method":"tools/list"}"#,
1408 );
1409 let tools = resp["result"]["tools"].as_array().expect("tools array");
1410 let names: Vec<&str> = tools
1411 .iter()
1412 .map(|t| t["name"].as_str().expect("name"))
1413 .collect();
1414 for expected in &[
1415 "explore",
1417 "map",
1418 "context",
1419 "impact",
1420 "owners",
1421 "why",
1422 "explain_association",
1423 "node_edges",
1424 "neighborhood",
1425 "edges_at",
1426 "what_if",
1427 "recall",
1428 "remember",
1429 "sync",
1430 "query",
1432 "ingest_json",
1433 "create_rule",
1434 "explain",
1435 "stats",
1436 "node_info",
1437 "upsert_entity",
1438 "find_similar",
1439 "hybrid_search",
1440 "node_history",
1441 "edge_history",
1442 "was_linked",
1443 "rename_node",
1444 ] {
1445 assert!(names.contains(expected), "missing tool: {expected}");
1446 }
1447 assert_eq!(
1448 names.len(),
1449 27,
1450 "expected exactly 27 tools, got {}",
1451 names.len()
1452 );
1453 assert_eq!(
1454 &names[..14],
1455 [
1456 "explore",
1457 "map",
1458 "context",
1459 "impact",
1460 "owners",
1461 "why",
1462 "explain_association",
1463 "node_edges",
1464 "neighborhood",
1465 "edges_at",
1466 "what_if",
1467 "recall",
1468 "remember",
1469 "sync"
1470 ],
1471 "the task tools come first, in order"
1472 );
1473 assert_eq!(names[14], "query", "the graph tools follow them");
1474 }
1475
1476 #[test]
1480 fn tools_list_defaults_to_fifteen_on_a_memory_store() {
1481 let db = demo_db();
1482 let resp = roundtrip(&db, r#"{"jsonrpc":"2.0","id":1,"method":"tools/list"}"#);
1483 let names: Vec<&str> = resp["result"]["tools"]
1484 .as_array()
1485 .expect("tools array")
1486 .iter()
1487 .map(|t| t["name"].as_str().expect("name"))
1488 .collect();
1489 assert_eq!(names, ASSOCIATION_TOOLS.to_vec());
1490 }
1491
1492 #[test]
1499 fn the_association_surface_is_entity_tools_only() {
1500 for kept in ["remember", "recall", "explain_association"] {
1501 assert!(
1502 ASSOCIATION_TOOLS.contains(&kept),
1503 "{kept} answers on an entity graph and must be listed"
1504 );
1505 }
1506 for code_only in [
1507 "explore", "map", "context", "impact", "owners", "why", "sync",
1508 ] {
1509 assert!(
1510 !ASSOCIATION_TOOLS.contains(&code_only),
1511 "{code_only} reads a code graph and must not be listed on a memory store"
1512 );
1513 }
1514 let served: Vec<String> = crate::mcp_tasks::task_tools()
1515 .iter()
1516 .chain(graph_tools().iter())
1517 .filter_map(|t| t.get("name").and_then(Js::as_str))
1518 .map(str::to_string)
1519 .collect();
1520 for name in ASSOCIATION_TOOLS {
1521 assert!(
1522 served.iter().any(|s| s == name),
1523 "{name} is listed but not served"
1524 );
1525 }
1526 assert!(
1527 CODE_GRAPH_TOOLS.contains(&"explore"),
1528 "and `explore` is the task tool the other surface lists"
1529 );
1530 }
1531
1532 #[test]
1535 fn tools_list_is_three_tools_on_a_code_graph_store() {
1536 let db = demo_db();
1537 db.write()
1538 .insert_node(
1539 "GitSync",
1540 crate::mcp_tasks::SYNC_KEY,
1541 vec![("id".into(), Value::Str(crate::mcp_tasks::SYNC_KEY.into()))],
1542 )
1543 .expect("marker");
1544 let resp = roundtrip(&db, r#"{"jsonrpc":"2.0","id":1,"method":"tools/list"}"#);
1545 let names: Vec<&str> = resp["result"]["tools"]
1546 .as_array()
1547 .expect("tools array")
1548 .iter()
1549 .map(|t| t["name"].as_str().expect("name"))
1550 .collect();
1551 assert_eq!(names, ["explore", "query", "stats"]);
1552 }
1553
1554 #[test]
1555 fn test_stats_returns_node_count() {
1556 let db = demo_db();
1557 let resp = tool_call(&db, 1, "stats", json!({}));
1558 assert!(!is_error(&resp));
1559 let result = tool_text(&resp);
1560 assert_eq!(result["nodes_live"], 2);
1561 }
1562
1563 #[test]
1564 fn test_query_runs_cypher() {
1565 let db = demo_db();
1566 let resp = tool_call(
1567 &db,
1568 1,
1569 "query",
1570 json!({ "cypher": "MATCH (n:Person) RETURN n.name ORDER BY n.name" }),
1571 );
1572 assert!(!is_error(&resp));
1573 let result = tool_text(&resp);
1574 assert_eq!(result["columns"], json!(["n.name"]));
1576 assert_eq!(result["rows"].as_array().map(|r| r.len()), Some(2));
1577 }
1578
1579 #[test]
1580 fn test_query_create_is_a_write() {
1581 let db = SharedDb::open(&tmp_dir()).expect("open");
1582 let resp = tool_call(
1583 &db,
1584 1,
1585 "query",
1586 json!({ "cypher": "CREATE (n:L {id: 'k'}) RETURN n" }),
1587 );
1588 assert!(
1589 !is_error(&resp),
1590 "CREATE via MCP query must succeed: {resp}"
1591 );
1592 let stats = tool_text(&tool_call(&db, 2, "stats", json!({})));
1593 assert_eq!(stats["nodes_live"], 1);
1594 }
1595
1596 #[test]
1597 fn test_ingest_json_inserts_nodes() {
1598 let db = demo_db();
1599 let resp = tool_call(
1600 &db,
1601 1,
1602 "ingest_json",
1603 json!({
1604 "label": "Person",
1605 "rows_json": r#"[{"id":"carol","name":"Carol"}]"#,
1606 "key_field": "id"
1607 }),
1608 );
1609 assert!(!is_error(&resp));
1610 let stats = tool_text(&tool_call(&db, 2, "stats", json!({})));
1612 assert_eq!(stats["nodes_live"], 3);
1613 }
1614
1615 #[test]
1616 fn test_node_info_returns_props() {
1617 let db = demo_db();
1618 let resp = tool_call(&db, 1, "node_info", json!({ "key": "alice" }));
1619 assert!(!is_error(&resp));
1620 let result = tool_text(&resp);
1621 assert_eq!(result["key"], "alice");
1622 assert_eq!(result["label"], "Person");
1623 assert_eq!(result["props"]["name"], "Alice");
1624 }
1625
1626 #[test]
1629 fn test_node_edges_returns_edges() {
1630 let db = demo_db();
1631 let resp = tool_call(
1632 &db,
1633 1,
1634 "node_edges",
1635 json!({ "key": "alice", "json": true }),
1636 );
1637 assert!(!is_error(&resp));
1638 let result = tool_text(&resp);
1639 assert_eq!(result["key"], "alice");
1640 let types = result["types"].as_array().expect("types");
1641 assert!(
1642 !types.is_empty(),
1643 "alice should have at least one edge type"
1644 );
1645 let similar = types
1646 .iter()
1647 .find(|t| t["edge_type"] == "SIMILAR")
1648 .expect("the rule's edge type");
1649 assert_eq!(similar["count"], json!(2));
1652 let edges = similar["edges"].as_array().expect("edges");
1653 let dirs: Vec<&str> = edges
1654 .iter()
1655 .map(|e| e["direction"].as_str().expect("direction"))
1656 .collect();
1657 assert!(dirs.contains(&"out") && dirs.contains(&"in"), "{similar}");
1658 for edge in edges {
1659 assert_eq!(edge["other"], json!("bob"));
1660 assert_eq!(edge["derived"], json!(true));
1661 assert_eq!(edge["rule"], json!("sim_emb"));
1662 assert_eq!(edge["score"], json!(1.0));
1663 assert!(
1664 edge["predicate"]
1665 .as_str()
1666 .unwrap_or("")
1667 .contains("vector_similar"),
1668 "the predicate travels with the edge: {edge}"
1669 );
1670 }
1671 }
1672
1673 #[test]
1676 fn test_neighborhood_traverses_one_hop() {
1677 let db = demo_db();
1678 let resp = tool_call(
1679 &db,
1680 1,
1681 "neighborhood",
1682 json!({ "key": "alice", "depth": 1, "json": true }),
1683 );
1684 assert!(!is_error(&resp));
1685 let result = tool_text(&resp);
1686 assert_eq!(result["key"], "alice");
1687 assert!(result["types"].as_array().is_some(), "{result}");
1688
1689 let deep = tool_call(
1690 &db,
1691 2,
1692 "neighborhood",
1693 json!({ "key": "alice", "depth": 2 }),
1694 );
1695 assert!(!is_error(&deep));
1696 let table = tool_text(&deep);
1697 assert_eq!(table["columns"], json!(["key", "label", "depth"]));
1698 assert!(table["rows"].as_array().is_some());
1699 }
1700
1701 #[test]
1702 fn test_explain_returns_rule_info() {
1703 let db = demo_db();
1704 let resp = tool_call(&db, 1, "explain", json!({ "a": "alice", "b": "bob" }));
1705 assert!(!is_error(&resp));
1706 let result = tool_text(&resp);
1707 let arr = result.as_array().expect("explain returns array");
1708 assert!(!arr.is_empty(), "expected at least one explanation");
1709 assert_eq!(arr[0]["rule"], "sim_emb");
1710 }
1711
1712 #[test]
1713 fn test_create_rule_backfills() {
1714 let db = SharedDb::open(&tmp_dir()).expect("open");
1715 {
1716 let mut g = db.write();
1717 let opts = IngestOptions {
1718 key_field: "id".into(),
1719 auto_fk: AutoFk::Off,
1720 };
1721 let rows: Vec<BTreeMap<String, Value>> = vec![
1722 [
1723 ("id", Value::Str("x".into())),
1724 ("tag", Value::Str("a".into())),
1725 ]
1726 .into_iter()
1727 .map(|(k, v)| (k.to_string(), v))
1728 .collect(),
1729 [
1730 ("id", Value::Str("y".into())),
1731 ("tag", Value::Str("a".into())),
1732 ]
1733 .into_iter()
1734 .map(|(k, v)| (k.to_string(), v))
1735 .collect(),
1736 ];
1737 g.ingest("Item", rows, &opts).expect("ingest");
1738 }
1739 let resp = tool_call(
1740 &db,
1741 1,
1742 "create_rule",
1743 json!({
1744 "name": "same_tag",
1745 "src_label": "Item",
1746 "dst_label": "Item",
1747 "predicate": { "FieldEqual": { "field": "tag" } },
1748 "edge_type": "SAME_TAG"
1749 }),
1750 );
1751 assert!(!is_error(&resp));
1752 let result = tool_text(&resp);
1753 assert_eq!(result["ok"], true);
1754 let edges_resp = tool_call(&db, 2, "node_edges", json!({ "key": "x", "json": true }));
1756 let edges_result = tool_text(&edges_resp);
1757 let types = edges_result["types"].as_array().expect("types");
1758 assert!(
1759 types.iter().any(|t| t["edge_type"] == "SAME_TAG"),
1760 "SAME_TAG edge not found after create_rule"
1761 );
1762 }
1763
1764 #[test]
1767 fn test_upsert_entity_creates_new_node() {
1768 let db = demo_db();
1769 let resp = tool_call(
1770 &db,
1771 1,
1772 "upsert_entity",
1773 json!({
1774 "key": "carol",
1775 "label": "Person",
1776 "props": { "name": "Carol", "age": 30 }
1777 }),
1778 );
1779 assert!(!is_error(&resp));
1780 let result = tool_text(&resp);
1781 assert_eq!(result["ok"], true);
1782 assert_eq!(result["created"], true);
1783 assert_eq!(result["key"], "carol");
1784 let info = tool_text(&tool_call(&db, 2, "node_info", json!({ "key": "carol" })));
1786 assert_eq!(info["props"]["name"], "Carol");
1787 }
1788
1789 #[test]
1790 fn test_upsert_entity_updates_existing_node() {
1791 let db = demo_db();
1792 let resp = tool_call(
1793 &db,
1794 1,
1795 "upsert_entity",
1796 json!({
1797 "key": "alice",
1798 "props": { "name": "Alice Updated" }
1799 }),
1800 );
1801 assert!(!is_error(&resp));
1802 let result = tool_text(&resp);
1803 assert_eq!(result["ok"], true);
1804 assert_eq!(result["created"], false);
1805 assert_eq!(result["updated_fields"], 1);
1806 let info = tool_text(&tool_call(&db, 2, "node_info", json!({ "key": "alice" })));
1808 assert_eq!(info["props"]["name"], "Alice Updated");
1809 }
1810
1811 #[test]
1812 fn test_upsert_entity_missing_label_on_create_is_error() {
1813 let db = demo_db();
1814 let resp = tool_call(
1815 &db,
1816 1,
1817 "upsert_entity",
1818 json!({ "key": "new-node", "props": { "x": 1 } }),
1819 );
1820 assert!(is_error(&resp), "should error without label for new node");
1821 }
1822
1823 #[test]
1824 fn test_find_similar_returns_similar_edges() {
1825 let db = demo_db();
1826 let resp = tool_call(
1827 &db,
1828 1,
1829 "find_similar",
1830 json!({ "key": "alice", "edge_type": "SIMILAR" }),
1831 );
1832 assert!(!is_error(&resp));
1833 let result = tool_text(&resp);
1834 assert_eq!(result["key"], "alice");
1835 assert_eq!(result["edge_type"], "SIMILAR");
1836 let similar = result["similar"].as_array().expect("similar array");
1837 assert!(!similar.is_empty(), "expected SIMILAR neighbors for alice");
1838 assert_eq!(similar[0]["neighbor_key"], "bob");
1839 }
1840
1841 #[test]
1842 fn test_find_similar_limit_respected() {
1843 let db = demo_db();
1844 let resp = tool_call(
1845 &db,
1846 1,
1847 "find_similar",
1848 json!({ "key": "alice", "edge_type": "SIMILAR", "limit": 0 }),
1849 );
1850 assert!(!is_error(&resp));
1851 let result = tool_text(&resp);
1852 let similar = result["similar"].as_array().expect("similar array");
1853 assert_eq!(similar.len(), 0);
1854 }
1855
1856 #[test]
1860 fn test_find_similar_vector_default_min_is_0_8() {
1861 let db = SharedDb::open(&tmp_dir()).expect("open");
1862 {
1863 let mut g = db.write();
1864 g.insert_node(
1866 "Item",
1867 "close",
1868 vec![(
1869 "emb".into(),
1870 Value::List(vec![Value::Float(1.0), Value::Float(0.0)]),
1871 )],
1872 )
1873 .unwrap();
1874 g.insert_node(
1876 "Item",
1877 "far",
1878 vec![(
1879 "emb".into(),
1880 Value::List(vec![Value::Float(0.0), Value::Float(1.0)]),
1881 )],
1882 )
1883 .unwrap();
1884 }
1885
1886 let resp = tool_call(
1888 &db,
1889 1,
1890 "find_similar",
1891 json!({
1892 "vector": [1.0, 0.0],
1893 "field": "emb",
1894 "label": "Item",
1895 "k": 10
1896 }),
1897 );
1898 assert!(!is_error(&resp), "vector search must not error");
1899 let result = tool_text(&resp);
1900 let results = result["results"].as_array().expect("results array");
1901
1902 let keys: Vec<&str> = results.iter().filter_map(|r| r["key"].as_str()).collect();
1903 assert!(
1904 keys.contains(&"close"),
1905 "close node (sim=1.0) must be included"
1906 );
1907 assert!(
1908 !keys.contains(&"far"),
1909 "far node (sim=0.0) must be excluded by default min=0.8"
1910 );
1911 }
1912
1913 #[test]
1915 fn test_find_similar_vector_mask_excludes_hidden() {
1916 let db = SharedDb::open(&tmp_dir()).expect("open");
1917 {
1918 let mut g = db.write();
1919 g.insert_node(
1921 "Item",
1922 "visible",
1923 vec![(
1924 "emb".into(),
1925 Value::List(vec![Value::Float(1.0), Value::Float(0.0)]),
1926 )],
1927 )
1928 .unwrap();
1929 g.insert_node(
1931 "Item",
1932 "hidden",
1933 vec![(
1934 "emb".into(),
1935 Value::List(vec![Value::Float(1.0), Value::Float(0.0)]),
1936 )],
1937 )
1938 .unwrap();
1939 }
1940
1941 let resp = tool_call(
1942 &db,
1943 1,
1944 "find_similar",
1945 json!({
1946 "vector": [1.0, 0.0],
1947 "field": "emb",
1948 "label": "Item",
1949 "k": 10,
1950 "min": 0.0,
1951 "mask": ["visible"]
1952 }),
1953 );
1954 assert!(!is_error(&resp), "masked vector search must not error");
1955 let result = tool_text(&resp);
1956 let results = result["results"].as_array().expect("results array");
1957
1958 let keys: Vec<&str> = results.iter().filter_map(|r| r["key"].as_str()).collect();
1959 assert!(
1960 keys.contains(&"visible"),
1961 "visible node must appear in masked results"
1962 );
1963 assert!(
1964 !keys.contains(&"hidden"),
1965 "hidden node must be excluded by mask"
1966 );
1967 }
1968
1969 #[test]
1971 fn test_find_similar_vector_mask_bad_type_is_error() {
1972 let db = SharedDb::open(&tmp_dir()).expect("open");
1973 let resp = tool_call(
1974 &db,
1975 1,
1976 "find_similar",
1977 json!({
1978 "vector": [1.0, 0.0],
1979 "field": "emb",
1980 "k": 5,
1981 "mask": [42]
1982 }),
1983 );
1984 assert!(
1985 is_error(&resp),
1986 "non-string mask element must produce a tool error"
1987 );
1988 }
1989
1990 #[test]
1992 fn test_find_similar_edge_mask_excludes_hidden_neighbor() {
1993 let db = SharedDb::open(&tmp_dir()).expect("open");
1994 {
1995 let mut g = db.write();
1996 g.insert_node("P", "alice", vec![]).unwrap();
1997 g.insert_node("P", "bob", vec![]).unwrap(); g.insert_node("P", "carol", vec![]).unwrap(); g.insert_edge("KNOWS", "alice", "bob").unwrap();
2000 g.insert_edge("KNOWS", "alice", "carol").unwrap();
2001 }
2002 let resp = tool_call(
2004 &db,
2005 1,
2006 "find_similar",
2007 json!({
2008 "key": "alice",
2009 "edge_type": "KNOWS",
2010 "mask": ["alice", "bob"]
2011 }),
2012 );
2013 assert!(!is_error(&resp), "masked edge search must not error");
2014 let result = tool_text(&resp);
2015 let similar = result["similar"].as_array().expect("similar array");
2016 let neighbors: Vec<&str> = similar
2017 .iter()
2018 .filter_map(|e| e["neighbor_key"].as_str())
2019 .collect();
2020 assert!(neighbors.contains(&"bob"), "bob (visible) must appear");
2021 assert!(
2022 !neighbors.contains(&"carol"),
2023 "carol (hidden) must be excluded"
2024 );
2025 }
2026
2027 #[test]
2030 fn test_find_similar_edge_mask_hidden_key_is_not_found() {
2031 let db = SharedDb::open(&tmp_dir()).expect("open");
2032 {
2033 let mut g = db.write();
2034 g.insert_node("P", "alice", vec![]).unwrap();
2035 g.insert_node("P", "bob", vec![]).unwrap();
2036 }
2037 let resp_masked = tool_call(
2039 &db,
2040 1,
2041 "find_similar",
2042 json!({ "key": "alice", "edge_type": "KNOWS", "mask": ["bob"] }),
2043 );
2044 let resp_ghost = tool_call(
2046 &db,
2047 2,
2048 "find_similar",
2049 json!({ "key": "ghost", "edge_type": "KNOWS" }),
2050 );
2051 assert!(
2052 is_error(&resp_masked),
2053 "hidden query key must produce a tool error"
2054 );
2055 assert!(
2056 is_error(&resp_ghost),
2057 "nonexistent key must produce a tool error"
2058 );
2059 assert_eq!(
2061 tool_err_text(&resp_masked).contains("alice"),
2062 tool_err_text(&resp_ghost).contains("ghost"),
2063 "error messages should follow same not-found template"
2064 );
2065 }
2066
2067 #[test]
2072 fn test_explain_association_same_as_explain() {
2073 let db = demo_db();
2074 let explain = tool_text(&tool_call(
2075 &db,
2076 1,
2077 "explain",
2078 json!({ "a": "alice", "b": "bob" }),
2079 ));
2080 let assoc = tool_text(&tool_call(
2081 &db,
2082 2,
2083 "explain_association",
2084 json!({ "a": "alice", "b": "bob", "json": true }),
2085 ));
2086 let explain: Vec<Js> = serde_json::from_value(explain).expect("explain array");
2087 let mut assoc: Vec<Js> = serde_json::from_value(assoc).expect("assoc array");
2088 for row in &mut assoc {
2089 let ev = row
2090 .as_object_mut()
2091 .expect("object")
2092 .remove("evidence")
2093 .expect("every derived edge carries its evidence");
2094 assert!(
2095 ev["similarity"].is_number(),
2096 "a vector_similar edge reports the cosine it scored: {ev}"
2097 );
2098 }
2099 assert_eq!(explain, assoc, "evidence is the only addition");
2100
2101 let prose = tool_call(
2102 &db,
2103 3,
2104 "explain_association",
2105 json!({ "a": "alice", "b": "bob" }),
2106 );
2107 let text = prose["result"]["content"][0]["text"]
2108 .as_str()
2109 .expect("text content");
2110 assert!(
2111 text.contains("mushroomdb explain — alice ↔ bob:"),
2112 "the default reply is the digest: {text}"
2113 );
2114 }
2115
2116 #[test]
2121 fn test_edge_history_returns_derived_lifecycle_with_rule() {
2122 let db = demo_db(); let resp = tool_call(&db, 1, "edge_history", json!({ "a": "alice", "b": "bob" }));
2124 assert!(!is_error(&resp), "edge_history must not error: {resp}");
2125 let result = tool_text(&resp);
2126
2127 let total = result["total_commits"].as_u64().expect("total_commits");
2129 assert!(total > 0, "total_commits must be > 0 after ingest + rule");
2130
2131 let events = result["events"].as_array().expect("events array");
2133 assert!(!events.is_empty(), "expected at least one edge event");
2134
2135 let derived_added = events
2137 .iter()
2138 .any(|ev| ev["event"].as_str() == Some("Added") && !ev["rule"].is_null());
2139 assert!(
2140 derived_added,
2141 "expected a derived Added event with rule attribution: {events:?}"
2142 );
2143 }
2144
2145 #[test]
2148 fn test_was_linked_at_valid_commit() {
2149 let db = SharedDb::open(&tmp_dir()).expect("open");
2150 {
2151 let mut g = db.write();
2152 let opts = IngestOptions {
2153 key_field: "id".into(),
2154 auto_fk: AutoFk::Off,
2155 };
2156 let rows: Vec<BTreeMap<String, Value>> = vec![
2157 [("id", Value::Str("x".into()))]
2158 .into_iter()
2159 .map(|(k, v)| (k.to_string(), v))
2160 .collect(),
2161 [("id", Value::Str("y".into()))]
2162 .into_iter()
2163 .map(|(k, v)| (k.to_string(), v))
2164 .collect(),
2165 ];
2166 g.ingest("N", rows, &opts).expect("ingest");
2167 g.insert_edge("LINK", "x", "y").expect("edge");
2168 }
2169 let g = db.read();
2171 let total = g.wal_total_commits().expect("wal_total_commits");
2172 drop(g);
2173
2174 let resp = tool_call(
2175 &db,
2176 1,
2177 "was_linked",
2178 json!({ "a": "x", "b": "y", "edge_type": "LINK", "at_commit": total - 1 }),
2179 );
2180 assert!(!is_error(&resp), "was_linked must not error: {resp}");
2181 let result = tool_text(&resp);
2182 assert_eq!(result["linked"], true);
2183 assert_eq!(result["a"], "x");
2184 assert_eq!(result["edge_type"], "LINK");
2185 }
2186
2187 #[test]
2190 fn test_was_linked_out_of_horizon_returns_tool_error() {
2191 let db = SharedDb::open(&tmp_dir()).expect("open");
2192 {
2193 let mut g = db.write();
2194 g.insert_node("N", "a", vec![]).expect("node a");
2195 g.insert_node("N", "b", vec![]).expect("node b");
2196 }
2197 let resp = tool_call(
2199 &db,
2200 1,
2201 "was_linked",
2202 json!({ "a": "a", "b": "b", "edge_type": "X", "at_commit": 999 }),
2203 );
2204 assert!(
2206 is_error(&resp),
2207 "out-of-range commit must be a tool error: {resp}"
2208 );
2209 let text = resp["result"]["content"][0]["text"].as_str().expect("text");
2210 assert!(
2211 text.contains("out of range") || text.contains("range"),
2212 "error must mention range: {text}"
2213 );
2214 }
2215
2216 #[test]
2219 fn test_node_history_via_mcp() {
2220 let db = demo_db(); let resp = tool_call(&db, 1, "node_history", json!({ "key": "alice" }));
2222 assert!(!is_error(&resp), "node_history must not error: {resp}");
2223 let result = tool_text(&resp);
2224
2225 assert_eq!(result["key"], "alice");
2226 let total = result["total_commits"].as_u64().expect("total_commits");
2227 assert!(total > 0, "total_commits must be > 0");
2228
2229 let history = result["history"].as_array().expect("history array");
2230 assert!(
2231 !history.is_empty(),
2232 "alice should have at least one history entry"
2233 );
2234
2235 let first_change = &history[0]["change"];
2237 assert_eq!(first_change["type"], "NodeInserted");
2238 assert_eq!(first_change["label"], "Person");
2239 }
2240}