1use crate::AppState;
23use anyhow::{anyhow, Context, Result};
24use serde_json::{json, Value};
25use trusty_memory_core::palace::{Palace, PalaceId, RoomType};
26use trusty_memory_core::retrieval::{recall, recall_across_palaces, recall_deep};
27use trusty_memory_core::store::kg::Triple;
28use uuid::Uuid;
29
30pub struct MemoryMcpServer;
37
38impl MemoryMcpServer {
39 pub fn new() -> Self {
40 Self
41 }
42}
43
44impl Default for MemoryMcpServer {
45 fn default() -> Self {
46 Self::new()
47 }
48}
49
50pub fn tool_definitions() -> Value {
61 tool_definitions_with(false)
62}
63
64pub fn tool_definitions_with(has_default: bool) -> Value {
75 let memory_remember_required: Vec<&str> = if has_default {
76 vec!["text"]
77 } else {
78 vec!["palace", "text"]
79 };
80 let memory_recall_required: Vec<&str> = if has_default {
81 vec!["query"]
82 } else {
83 vec!["palace", "query"]
84 };
85 let kg_assert_required: Vec<&str> = if has_default {
86 vec!["subject", "predicate", "object"]
87 } else {
88 vec!["palace", "subject", "predicate", "object"]
89 };
90 let kg_query_required: Vec<&str> = if has_default {
91 vec!["subject"]
92 } else {
93 vec!["palace", "subject"]
94 };
95 let memory_list_required: Vec<&str> = if has_default { vec![] } else { vec!["palace"] };
96 let memory_forget_required: Vec<&str> = if has_default {
97 vec!["drawer_id"]
98 } else {
99 vec!["palace", "drawer_id"]
100 };
101 let palace_info_required: Vec<&str> = if has_default { vec![] } else { vec!["palace"] };
102 let palace_compact_required: Vec<&str> = if has_default { vec![] } else { vec!["palace"] };
103
104 json!({
105 "tools": [
106 {
107 "name": "memory_remember",
108 "description": "Store a memory (drawer) in a palace room.",
109 "inputSchema": {
110 "type": "object",
111 "properties": {
112 "palace": {"type": "string", "description": "Palace ID (optional if server started with --palace)"},
113 "text": {"type": "string", "description": "Memory content"},
114 "room": {"type": "string", "description": "Room type (optional)"},
115 "tags": {"type": "array", "items": {"type": "string"}}
116 },
117 "required": memory_remember_required,
118 }
119 },
120 {
121 "name": "memory_recall",
122 "description": "Recall memories using L0+L1+L2 progressive retrieval.",
123 "inputSchema": {
124 "type": "object",
125 "properties": {
126 "palace": {"type": "string"},
127 "query": {"type": "string"},
128 "top_k": {"type": "integer", "default": 10}
129 },
130 "required": memory_recall_required,
131 }
132 },
133 {
134 "name": "memory_recall_deep",
135 "description": "Deep recall using L3 full HNSW search.",
136 "inputSchema": {
137 "type": "object",
138 "properties": {
139 "palace": {"type": "string"},
140 "query": {"type": "string"},
141 "top_k": {"type": "integer", "default": 10}
142 },
143 "required": memory_recall_required,
144 }
145 },
146 {
147 "name": "palace_create",
148 "description": "Create a new memory palace.",
149 "inputSchema": {
150 "type": "object",
151 "properties": {
152 "name": {"type": "string"},
153 "description": {"type": "string"}
154 },
155 "required": ["name"]
156 }
157 },
158 {
159 "name": "palace_list",
160 "description": "List all palaces on this machine.",
161 "inputSchema": {"type": "object", "properties": {}}
162 },
163 {
164 "name": "kg_assert",
165 "description": "Assert a fact in the temporal knowledge graph.",
166 "inputSchema": {
167 "type": "object",
168 "properties": {
169 "palace": {"type": "string"},
170 "subject": {"type": "string"},
171 "predicate": {"type": "string"},
172 "object": {"type": "string"},
173 "confidence": {"type": "number", "default": 1.0},
174 "provenance": {"type": "string"}
175 },
176 "required": kg_assert_required,
177 }
178 },
179 {
180 "name": "kg_query",
181 "description": "Query active knowledge-graph triples for a subject.",
182 "inputSchema": {
183 "type": "object",
184 "properties": {
185 "palace": {"type": "string"},
186 "subject": {"type": "string"}
187 },
188 "required": kg_query_required,
189 }
190 },
191 {
192 "name": "memory_list",
193 "description": "List drawers in a palace, optionally filtered by room type or tag.",
194 "inputSchema": {
195 "type": "object",
196 "properties": {
197 "palace": {"type": "string"},
198 "room": {"type": "string", "description": "Filter by room type (Frontend, Backend, Testing, Planning, Documentation, Research, Configuration, Meetings, General, or custom)"},
199 "tag": {"type": "string", "description": "Filter by tag"},
200 "limit": {"type": "integer", "description": "Max results (default 50)"}
201 },
202 "required": memory_list_required,
203 }
204 },
205 {
206 "name": "memory_forget",
207 "description": "Delete a drawer from a palace by its UUID.",
208 "inputSchema": {
209 "type": "object",
210 "properties": {
211 "palace": {"type": "string"},
212 "drawer_id": {"type": "string", "description": "UUID of the drawer to delete"}
213 },
214 "required": memory_forget_required,
215 }
216 },
217 {
218 "name": "palace_info",
219 "description": "Get metadata and stats for a single palace.",
220 "inputSchema": {
221 "type": "object",
222 "properties": {
223 "palace": {"type": "string"}
224 },
225 "required": palace_info_required,
226 }
227 },
228 {
229 "name": "palace_compact",
230 "description": "Remove orphaned vector index entries (vectors with no matching drawer row). See issue #49.",
231 "inputSchema": {
232 "type": "object",
233 "properties": {
234 "palace": {"type": "string"}
235 },
236 "required": palace_compact_required,
237 }
238 },
239 {
240 "name": "memory_recall_all",
241 "description": "Semantic search across ALL palaces simultaneously. Returns the top-k most relevant drawers ranked by similarity, regardless of which palace they belong to. Each result includes a `palace_id` field identifying its source.",
242 "inputSchema": {
243 "type": "object",
244 "properties": {
245 "q": {"type": "string", "description": "Free-text query"},
246 "top_k": {"type": "integer", "default": 10},
247 "deep": {"type": "boolean", "default": false}
248 },
249 "required": ["q"],
250 }
251 }
252 ]
253 })
254}
255
256fn parse_room(s: Option<&str>) -> RoomType {
264 match s.unwrap_or("General") {
265 "Frontend" => RoomType::Frontend,
266 "Backend" => RoomType::Backend,
267 "Testing" => RoomType::Testing,
268 "Planning" => RoomType::Planning,
269 "Documentation" => RoomType::Documentation,
270 "Research" => RoomType::Research,
271 "Configuration" => RoomType::Configuration,
272 "Meetings" => RoomType::Meetings,
273 "General" => RoomType::General,
274 other => RoomType::Custom(other.to_string()),
275 }
276}
277
278fn open_palace_handle(
280 state: &AppState,
281 palace_id: &str,
282) -> Result<std::sync::Arc<trusty_memory_core::PalaceHandle>> {
283 let pid = PalaceId::new(palace_id);
284 state
285 .registry
286 .open_palace(&state.data_root, &pid)
287 .with_context(|| format!("open palace {palace_id}"))
288}
289
290fn resolve_palace<'a>(state: &'a AppState, args: &'a Value, tool: &str) -> Result<String> {
302 if let Some(p) = args.get("palace").and_then(|v| v.as_str()) {
303 return Ok(p.to_string());
304 }
305 state
306 .default_palace
307 .clone()
308 .ok_or_else(|| anyhow!("{tool}: missing 'palace' (no --palace default configured)"))
309}
310
311pub async fn dispatch_tool(state: &AppState, name: &str, args: Value) -> Result<Value> {
321 match name {
322 "memory_remember" => {
323 let palace = resolve_palace(state, &args, "memory_remember")?;
324 let palace = palace.as_str();
325 let text = args
326 .get("text")
327 .and_then(|v| v.as_str())
328 .ok_or_else(|| anyhow!("memory_remember: missing 'text'"))?
329 .to_string();
330 let room = parse_room(args.get("room").and_then(|v| v.as_str()));
331 let tags: Vec<String> = args
332 .get("tags")
333 .and_then(|v| v.as_array())
334 .map(|arr| {
335 arr.iter()
336 .filter_map(|t| t.as_str().map(|s| s.to_string()))
337 .collect()
338 })
339 .unwrap_or_default();
340
341 let handle = open_palace_handle(state, palace)?;
342 let drawer_id = handle
343 .remember(text, room, tags, 0.5)
344 .await
345 .context("PalaceHandle::remember")?;
346 Ok(json!({
347 "drawer_id": drawer_id.to_string(),
348 "palace": palace,
349 "status": "stored",
350 }))
351 }
352 "memory_recall" => {
353 let palace = resolve_palace(state, &args, "memory_recall")?;
354 let query = args
355 .get("query")
356 .and_then(|v| v.as_str())
357 .ok_or_else(|| anyhow!("memory_recall: missing 'query'"))?;
358 let top_k = args.get("top_k").and_then(|v| v.as_u64()).unwrap_or(10) as usize;
359
360 let handle = open_palace_handle(state, &palace)?;
361 let embedder = state.embedder().await?;
362 let results = recall(&handle, embedder.as_ref(), query, top_k)
363 .await
364 .context("recall")?;
365 Ok(serialize_recall(&palace, query, results))
366 }
367 "memory_recall_deep" => {
368 let palace = resolve_palace(state, &args, "memory_recall_deep")?;
369 let query = args
370 .get("query")
371 .and_then(|v| v.as_str())
372 .ok_or_else(|| anyhow!("memory_recall_deep: missing 'query'"))?;
373 let top_k = args.get("top_k").and_then(|v| v.as_u64()).unwrap_or(10) as usize;
374
375 let handle = open_palace_handle(state, &palace)?;
376 let embedder = state.embedder().await?;
377 let results = recall_deep(&handle, embedder.as_ref(), query, top_k)
378 .await
379 .context("recall_deep")?;
380 Ok(serialize_recall(&palace, query, results))
381 }
382 "palace_create" => {
383 let palace_name = args
384 .get("name")
385 .and_then(|v| v.as_str())
386 .ok_or_else(|| anyhow!("palace_create: missing 'name'"))?;
387 let description = args
388 .get("description")
389 .and_then(|v| v.as_str())
390 .map(|s| s.to_string());
391 let palace = Palace {
392 id: PalaceId::new(palace_name),
393 name: palace_name.to_string(),
394 description,
395 created_at: chrono::Utc::now(),
396 data_dir: state.data_root.join(palace_name),
397 };
398 let _handle = state
399 .registry
400 .create_palace(&state.data_root, palace)
401 .context("create_palace")?;
402 Ok(json!({"palace_id": palace_name, "status": "created"}))
403 }
404 "palace_list" => {
405 let root = state.data_root.clone();
406 let palaces = tokio::task::spawn_blocking(move || {
407 trusty_memory_core::PalaceRegistry::list_palaces(&root)
408 })
409 .await
410 .context("join list_palaces")??;
411 let ids: Vec<String> = palaces.iter().map(|p| p.id.as_str().to_string()).collect();
412 Ok(json!({"palaces": ids}))
413 }
414 "kg_assert" => {
415 let palace = resolve_palace(state, &args, "kg_assert")?;
416 let palace = palace.as_str();
417 let subject = args
418 .get("subject")
419 .and_then(|v| v.as_str())
420 .ok_or_else(|| anyhow!("kg_assert: missing 'subject'"))?
421 .to_string();
422 let predicate = args
423 .get("predicate")
424 .and_then(|v| v.as_str())
425 .ok_or_else(|| anyhow!("kg_assert: missing 'predicate'"))?
426 .to_string();
427 let object = args
428 .get("object")
429 .and_then(|v| v.as_str())
430 .ok_or_else(|| anyhow!("kg_assert: missing 'object'"))?
431 .to_string();
432 let confidence = args
433 .get("confidence")
434 .and_then(|v| v.as_f64())
435 .map(|c| (c as f32).clamp(0.0, 1.0))
436 .unwrap_or(1.0);
437 let provenance = args
438 .get("provenance")
439 .and_then(|v| v.as_str())
440 .map(|s| s.to_string());
441
442 let handle = open_palace_handle(state, palace)?;
443 let triple = Triple {
444 subject,
445 predicate,
446 object,
447 valid_from: chrono::Utc::now(),
448 valid_to: None,
449 confidence,
450 provenance,
451 };
452 handle.kg.assert(triple).await.context("kg.assert")?;
453 Ok(json!({"status": "asserted"}))
454 }
455 "kg_query" => {
456 let palace = resolve_palace(state, &args, "kg_query")?;
457 let subject = args
458 .get("subject")
459 .and_then(|v| v.as_str())
460 .ok_or_else(|| anyhow!("kg_query: missing 'subject'"))?;
461 let handle = open_palace_handle(state, &palace)?;
462 let triples = handle
463 .kg
464 .query_active(subject)
465 .await
466 .context("kg.query_active")?;
467 let payload: Vec<Value> = triples
468 .iter()
469 .map(|t| {
470 json!({
471 "subject": t.subject,
472 "predicate": t.predicate,
473 "object": t.object,
474 "valid_from": t.valid_from.to_rfc3339(),
475 "valid_to": t.valid_to.as_ref().map(|d| d.to_rfc3339()),
476 "confidence": t.confidence,
477 "provenance": t.provenance,
478 })
479 })
480 .collect();
481 Ok(json!({"subject": subject, "triples": payload}))
482 }
483 "memory_list" => {
484 let palace = resolve_palace(state, &args, "memory_list")?;
485 let handle = open_palace_handle(state, &palace)?;
486 let room = args
487 .get("room")
488 .and_then(|v| v.as_str())
489 .map(|s| parse_room(Some(s)));
490 let tag = args
491 .get("tag")
492 .and_then(|v| v.as_str())
493 .map(|s| s.to_string());
494 let limit = args.get("limit").and_then(|v| v.as_u64()).unwrap_or(50) as usize;
495 let drawers = handle.list_drawers(room, tag, limit);
496 let payload: Vec<Value> = drawers
497 .iter()
498 .map(|d| {
499 json!({
500 "drawer_id": d.id.to_string(),
501 "content": d.content,
502 "importance": d.importance,
503 "tags": d.tags,
504 "created_at": d.created_at.to_rfc3339(),
505 })
506 })
507 .collect();
508 Ok(json!({"palace": palace, "drawers": payload}))
509 }
510 "memory_forget" => {
511 let palace = resolve_palace(state, &args, "memory_forget")?;
512 let drawer_id_str = args
513 .get("drawer_id")
514 .and_then(|v| v.as_str())
515 .ok_or_else(|| anyhow!("memory_forget: missing 'drawer_id'"))?;
516 let drawer_id = Uuid::parse_str(drawer_id_str)
517 .map_err(|e| anyhow!("memory_forget: invalid drawer_id UUID: {e}"))?;
518 let handle = open_palace_handle(state, &palace)?;
519 handle.forget(drawer_id).await.context("forget")?;
520 Ok(json!({"status": "deleted", "drawer_id": drawer_id_str, "palace": palace}))
521 }
522 "palace_info" => {
523 let palace = resolve_palace(state, &args, "palace_info")?;
524 let handle = open_palace_handle(state, &palace)?;
525 let drawer_count = handle.list_drawers(None, None, usize::MAX).len();
526 let data_dir = handle
527 .data_dir
528 .as_ref()
529 .map(|p| p.to_string_lossy().to_string());
530 Ok(json!({
531 "id": handle.id.as_str(),
532 "name": handle.id.as_str(),
533 "drawer_count": drawer_count,
534 "data_dir": data_dir,
535 }))
536 }
537 "palace_compact" => {
538 let palace = resolve_palace(state, &args, "palace_compact")?;
539 let handle = open_palace_handle(state, &palace)?;
540 let valid_ids: std::collections::HashSet<Uuid> =
544 handle.drawers.read().iter().map(|d| d.id).collect();
545 let vector_store = handle.vector_store.clone();
546 let res = tokio::task::spawn_blocking(move || vector_store.compact_orphans(&valid_ids))
547 .await
548 .context("join palace_compact")??;
549 Ok(json!({
550 "palace": palace,
551 "total_checked": res.total_checked,
552 "orphans_removed": res.orphans_removed,
553 "index_size_before": res.index_size_before,
554 "index_size_after": res.index_size_after,
555 }))
556 }
557 "memory_recall_all" => {
558 let query = args
559 .get("q")
560 .and_then(|v| v.as_str())
561 .ok_or_else(|| anyhow!("memory_recall_all: missing 'q'"))?;
562 let top_k = args.get("top_k").and_then(|v| v.as_u64()).unwrap_or(10) as usize;
563 let deep = args.get("deep").and_then(|v| v.as_bool()).unwrap_or(false);
564
565 let root = state.data_root.clone();
569 let palaces = tokio::task::spawn_blocking(move || {
570 trusty_memory_core::PalaceRegistry::list_palaces(&root)
571 })
572 .await
573 .context("join list_palaces")??;
574
575 let mut handles = Vec::with_capacity(palaces.len());
576 for p in &palaces {
577 match state.registry.open_palace(&state.data_root, &p.id) {
578 Ok(h) => handles.push(h),
579 Err(e) => {
580 tracing::warn!(palace = %p.id, "memory_recall_all: open failed: {e:#}")
581 }
582 }
583 }
584
585 let embedder = state.embedder().await?;
586 let erased: std::sync::Arc<dyn trusty_memory_core::embed::Embedder + Send + Sync> =
587 embedder;
588 let results = recall_across_palaces(&handles, &erased, query, top_k, deep)
589 .await
590 .context("recall_across_palaces")?;
591
592 let payload: Vec<Value> = results
593 .iter()
594 .map(|r| {
595 json!({
596 "palace_id": r.palace_id,
597 "drawer_id": r.result.drawer.id.to_string(),
598 "content": r.result.drawer.content,
599 "importance": r.result.drawer.importance,
600 "tags": r.result.drawer.tags,
601 "score": r.result.score,
602 "layer": r.result.layer,
603 })
604 })
605 .collect();
606 Ok(json!({ "query": query, "results": payload }))
607 }
608 other => anyhow::bail!("unknown tool: {other}"),
609 }
610}
611
612fn serialize_recall(
614 palace: &str,
615 query: &str,
616 results: Vec<trusty_memory_core::retrieval::RecallResult>,
617) -> Value {
618 let payload: Vec<Value> = results
619 .iter()
620 .map(|r| {
621 json!({
622 "drawer_id": r.drawer.id.to_string(),
623 "content": r.drawer.content,
624 "score": r.score,
625 "layer": r.layer,
626 "tags": r.drawer.tags,
627 "importance": r.drawer.importance,
628 })
629 })
630 .collect();
631 json!({
632 "palace": palace,
633 "query": query,
634 "results": payload,
635 })
636}
637
638#[cfg(test)]
639mod tests {
640 use super::*;
641 use crate::AppState;
642
643 fn test_state() -> AppState {
644 let tmp = tempfile::tempdir().expect("tempdir");
645 let root = tmp.path().to_path_buf();
646 std::mem::forget(tmp);
647 AppState::new(root)
648 }
649
650 #[test]
655 fn tool_definitions_drops_palace_required_when_default_set() {
656 let with_default = tool_definitions_with(true);
657 let without_default = tool_definitions_with(false);
658 for (name, palace_required_when_no_default) in [
659 ("memory_remember", true),
660 ("memory_recall", true),
661 ("memory_recall_deep", true),
662 ("memory_list", true),
663 ("memory_forget", true),
664 ("palace_info", true),
665 ("palace_compact", true),
666 ("kg_assert", true),
667 ("kg_query", true),
668 ] {
669 for (defs, has_default) in [(&with_default, true), (&without_default, false)] {
670 let tools = defs["tools"].as_array().unwrap();
671 let tool = tools.iter().find(|t| t["name"] == name).unwrap();
672 let required: Vec<&str> = tool["inputSchema"]["required"]
673 .as_array()
674 .unwrap()
675 .iter()
676 .filter_map(|v| v.as_str())
677 .collect();
678 let palace_required = required.contains(&"palace");
679 let expected = palace_required_when_no_default && !has_default;
680 assert_eq!(
681 palace_required, expected,
682 "tool={name} has_default={has_default} required={required:?}"
683 );
684 }
685 }
686 }
687
688 #[test]
689 fn tool_definitions_lists_all_tools() {
690 let defs = tool_definitions();
691 let tools = defs
692 .get("tools")
693 .and_then(|t| t.as_array())
694 .expect("tools array");
695 assert_eq!(tools.len(), 12);
696 let names: Vec<&str> = tools
697 .iter()
698 .filter_map(|t| t.get("name").and_then(|n| n.as_str()))
699 .collect();
700 for expected in [
701 "memory_remember",
702 "memory_recall",
703 "memory_recall_deep",
704 "memory_list",
705 "memory_forget",
706 "palace_create",
707 "palace_list",
708 "palace_info",
709 "palace_compact",
710 "kg_assert",
711 "kg_query",
712 "memory_recall_all",
713 ] {
714 assert!(names.contains(&expected), "missing tool: {expected}");
715 }
716 }
717
718 #[tokio::test]
721 async fn dispatch_palace_create_persists() {
722 let state = test_state();
723 let created = dispatch_tool(&state, "palace_create", json!({"name": "alpha"}))
724 .await
725 .expect("palace_create");
726 assert_eq!(created["palace_id"], "alpha");
727
728 let listed = dispatch_tool(&state, "palace_list", json!({}))
729 .await
730 .expect("palace_list");
731 let ids = listed["palaces"].as_array().expect("palaces array");
732 assert!(ids.iter().any(|v| v.as_str() == Some("alpha")));
733 }
734
735 #[tokio::test]
738 async fn dispatch_remember_then_recall() {
739 let state = test_state();
740 let _ = dispatch_tool(&state, "palace_create", json!({"name": "beta"}))
741 .await
742 .expect("palace_create");
743
744 let remembered = dispatch_tool(
745 &state,
746 "memory_remember",
747 json!({
748 "palace": "beta",
749 "text": "Quokkas are the happiest marsupials in Australia",
750 "room": "General",
751 "tags": ["wildlife"],
752 }),
753 )
754 .await
755 .expect("memory_remember");
756 assert!(remembered["drawer_id"].as_str().is_some());
757
758 let recalled = dispatch_tool(
759 &state,
760 "memory_recall",
761 json!({"palace": "beta", "query": "Quokkas marsupials Australia", "top_k": 5}),
762 )
763 .await
764 .expect("memory_recall");
765 let results = recalled["results"].as_array().expect("results");
766 assert!(
767 results
768 .iter()
769 .any(|r| r["content"].as_str().unwrap_or("").contains("Quokkas")),
770 "expected to recall the Quokkas drawer; got {results:?}"
771 );
772 }
773
774 #[tokio::test]
777 async fn dispatch_kg_assert_then_query() {
778 let state = test_state();
779 let _ = dispatch_tool(&state, "palace_create", json!({"name": "gamma"}))
780 .await
781 .expect("palace_create");
782
783 let _ = dispatch_tool(
784 &state,
785 "kg_assert",
786 json!({
787 "palace": "gamma",
788 "subject": "alice",
789 "predicate": "works_at",
790 "object": "Acme",
791 "confidence": 0.9,
792 "provenance": "test",
793 }),
794 )
795 .await
796 .expect("kg_assert");
797
798 let queried = dispatch_tool(
799 &state,
800 "kg_query",
801 json!({"palace": "gamma", "subject": "alice"}),
802 )
803 .await
804 .expect("kg_query");
805 let triples = queried["triples"].as_array().expect("triples array");
806 assert_eq!(triples.len(), 1);
807 assert_eq!(triples[0]["object"], "Acme");
808 assert_eq!(triples[0]["predicate"], "works_at");
809 }
810
811 #[tokio::test]
812 async fn dispatch_unknown_tool_errors() {
813 let state = test_state();
814 let err = dispatch_tool(&state, "does_not_exist", json!({}))
815 .await
816 .expect_err("should error");
817 assert!(err.to_string().contains("unknown tool"));
818 }
819}