1use crate::AppState;
23use anyhow::{anyhow, Context, Result};
24use serde_json::{json, Value};
25use trusty_common::memory_core::palace::{Palace, PalaceId, RoomType};
26use trusty_common::memory_core::retrieval::{recall, recall_across_palaces, recall_deep};
27use trusty_common::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_common::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_common::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_common::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<
587 dyn trusty_common::memory_core::embed::Embedder + Send + Sync,
588 > = embedder;
589 let results = recall_across_palaces(&handles, &erased, query, top_k, deep)
590 .await
591 .context("recall_across_palaces")?;
592
593 let payload: Vec<Value> = results
594 .iter()
595 .map(|r| {
596 json!({
597 "palace_id": r.palace_id,
598 "drawer_id": r.result.drawer.id.to_string(),
599 "content": r.result.drawer.content,
600 "importance": r.result.drawer.importance,
601 "tags": r.result.drawer.tags,
602 "score": r.result.score,
603 "layer": r.result.layer,
604 })
605 })
606 .collect();
607 Ok(json!({ "query": query, "results": payload }))
608 }
609 other => anyhow::bail!("unknown tool: {other}"),
610 }
611}
612
613fn serialize_recall(
615 palace: &str,
616 query: &str,
617 results: Vec<trusty_common::memory_core::retrieval::RecallResult>,
618) -> Value {
619 let payload: Vec<Value> = results
620 .iter()
621 .map(|r| {
622 json!({
623 "drawer_id": r.drawer.id.to_string(),
624 "content": r.drawer.content,
625 "score": r.score,
626 "layer": r.layer,
627 "tags": r.drawer.tags,
628 "importance": r.drawer.importance,
629 })
630 })
631 .collect();
632 json!({
633 "palace": palace,
634 "query": query,
635 "results": payload,
636 })
637}
638
639#[cfg(test)]
640mod tests {
641 use super::*;
642 use crate::AppState;
643
644 fn test_state() -> AppState {
645 let tmp = tempfile::tempdir().expect("tempdir");
646 let root = tmp.path().to_path_buf();
647 std::mem::forget(tmp);
648 AppState::new(root)
649 }
650
651 #[test]
656 fn tool_definitions_drops_palace_required_when_default_set() {
657 let with_default = tool_definitions_with(true);
658 let without_default = tool_definitions_with(false);
659 for (name, palace_required_when_no_default) in [
660 ("memory_remember", true),
661 ("memory_recall", true),
662 ("memory_recall_deep", true),
663 ("memory_list", true),
664 ("memory_forget", true),
665 ("palace_info", true),
666 ("palace_compact", true),
667 ("kg_assert", true),
668 ("kg_query", true),
669 ] {
670 for (defs, has_default) in [(&with_default, true), (&without_default, false)] {
671 let tools = defs["tools"].as_array().unwrap();
672 let tool = tools.iter().find(|t| t["name"] == name).unwrap();
673 let required: Vec<&str> = tool["inputSchema"]["required"]
674 .as_array()
675 .unwrap()
676 .iter()
677 .filter_map(|v| v.as_str())
678 .collect();
679 let palace_required = required.contains(&"palace");
680 let expected = palace_required_when_no_default && !has_default;
681 assert_eq!(
682 palace_required, expected,
683 "tool={name} has_default={has_default} required={required:?}"
684 );
685 }
686 }
687 }
688
689 #[test]
690 fn tool_definitions_lists_all_tools() {
691 let defs = tool_definitions();
692 let tools = defs
693 .get("tools")
694 .and_then(|t| t.as_array())
695 .expect("tools array");
696 assert_eq!(tools.len(), 12);
697 let names: Vec<&str> = tools
698 .iter()
699 .filter_map(|t| t.get("name").and_then(|n| n.as_str()))
700 .collect();
701 for expected in [
702 "memory_remember",
703 "memory_recall",
704 "memory_recall_deep",
705 "memory_list",
706 "memory_forget",
707 "palace_create",
708 "palace_list",
709 "palace_info",
710 "palace_compact",
711 "kg_assert",
712 "kg_query",
713 "memory_recall_all",
714 ] {
715 assert!(names.contains(&expected), "missing tool: {expected}");
716 }
717 }
718
719 #[tokio::test]
722 async fn dispatch_palace_create_persists() {
723 let state = test_state();
724 let created = dispatch_tool(&state, "palace_create", json!({"name": "alpha"}))
725 .await
726 .expect("palace_create");
727 assert_eq!(created["palace_id"], "alpha");
728
729 let listed = dispatch_tool(&state, "palace_list", json!({}))
730 .await
731 .expect("palace_list");
732 let ids = listed["palaces"].as_array().expect("palaces array");
733 assert!(ids.iter().any(|v| v.as_str() == Some("alpha")));
734 }
735
736 #[tokio::test]
739 async fn dispatch_remember_then_recall() {
740 let state = test_state();
741 let _ = dispatch_tool(&state, "palace_create", json!({"name": "beta"}))
742 .await
743 .expect("palace_create");
744
745 let remembered = dispatch_tool(
746 &state,
747 "memory_remember",
748 json!({
749 "palace": "beta",
750 "text": "Quokkas are the happiest marsupials in Australia",
751 "room": "General",
752 "tags": ["wildlife"],
753 }),
754 )
755 .await
756 .expect("memory_remember");
757 assert!(remembered["drawer_id"].as_str().is_some());
758
759 let recalled = dispatch_tool(
760 &state,
761 "memory_recall",
762 json!({"palace": "beta", "query": "Quokkas marsupials Australia", "top_k": 5}),
763 )
764 .await
765 .expect("memory_recall");
766 let results = recalled["results"].as_array().expect("results");
767 assert!(
768 results
769 .iter()
770 .any(|r| r["content"].as_str().unwrap_or("").contains("Quokkas")),
771 "expected to recall the Quokkas drawer; got {results:?}"
772 );
773 }
774
775 #[tokio::test]
778 async fn dispatch_kg_assert_then_query() {
779 let state = test_state();
780 let _ = dispatch_tool(&state, "palace_create", json!({"name": "gamma"}))
781 .await
782 .expect("palace_create");
783
784 let _ = dispatch_tool(
785 &state,
786 "kg_assert",
787 json!({
788 "palace": "gamma",
789 "subject": "alice",
790 "predicate": "works_at",
791 "object": "Acme",
792 "confidence": 0.9,
793 "provenance": "test",
794 }),
795 )
796 .await
797 .expect("kg_assert");
798
799 let queried = dispatch_tool(
800 &state,
801 "kg_query",
802 json!({"palace": "gamma", "subject": "alice"}),
803 )
804 .await
805 .expect("kg_query");
806 let triples = queried["triples"].as_array().expect("triples array");
807 assert_eq!(triples.len(), 1);
808 assert_eq!(triples[0]["object"], "Acme");
809 assert_eq!(triples[0]["predicate"], "works_at");
810 }
811
812 #[tokio::test]
813 async fn dispatch_unknown_tool_errors() {
814 let state = test_state();
815 let err = dispatch_tool(&state, "does_not_exist", json!({}))
816 .await
817 .expect_err("should error");
818 assert!(err.to_string().contains("unknown tool"));
819 }
820}