From c935cec4af744510286b235bb1fd60a22ec618b5 Mon Sep 17 00:00:00 2001
From: Dustin <dustin.hanusch@mailbox.tu-dresden.de>
Date: Fri, 10 Apr 2026 16:45:47 +0200
Subject: [PATCH] fix: KG path mismatch between MCP server and CLI
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The MCP server's KnowledgeGraph was initialized with either
palace_path/knowledge_graph.sqlite3 (with --palace) or
DEFAULT_KG_PATH (~/.mempalace/knowledge_graph.sqlite3) (without).
The CLI reads from palace_path/knowledge.db.

This caused KG facts written via MCP tools (mempalace_kg_add) to be
invisible to CLI commands (mempalace search, mempalace_kg_query in
new sessions) — a silent data loss bug for any user running the MCP
server without --palace.

Fix: Always derive KG path from MempalaceConfig.palace_path using
the canonical filename "knowledge.db", matching the CLI.

Also adds an atexit shutdown hook that checkpoints the SQLite WAL
and releases the ChromaDB client before process exit, preventing
data loss when the stdio transport closes.
---
 mempalace/mcp_server.py | 40 ++++++++++++++++++++++++++++++++++++----
 1 file changed, 36 insertions(+), 4 deletions(-)

diff --git a/mempalace/mcp_server.py b/mempalace/mcp_server.py
index bffd3b2f2d..a7776e0d29 100644
--- a/mempalace/mcp_server.py
+++ b/mempalace/mcp_server.py
@@ -57,10 +57,11 @@ def _parse_args():
     os.environ["MEMPALACE_PALACE_PATH"] = os.path.abspath(_args.palace)
 
 _config = MempalaceConfig()
-if _args.palace:
-    _kg = KnowledgeGraph(db_path=os.path.join(_config.palace_path, "knowledge_graph.sqlite3"))
-else:
-    _kg = KnowledgeGraph()
+# Always derive KG path from palace_path to match CLI expectations.
+# Previously, the no-arg fallback used DEFAULT_KG_PATH (~/.mempalace/knowledge_graph.sqlite3)
+# while CLI reads from palace_path/knowledge.db — causing KG writes via MCP to be
+# invisible to CLI search/query commands.
+_kg = KnowledgeGraph(db_path=os.path.join(_config.palace_path, "knowledge.db"))
 
 
 _client_cache = None
@@ -921,6 +922,37 @@ def handle_request(request):
     }
 
 
+def _shutdown():
+    """Flush KG WAL and release ChromaDB client before exit.
+
+    Prevents data loss when the stdio transport closes — ChromaDB's
+    PersistentClient uses an in-memory WAL that may not flush if the
+    process exits before the background thread completes.
+
+    Note: atexit handlers do not run on SIGKILL. For maximum robustness,
+    consider PRAGMA journal_mode=DELETE at startup (trades write
+    concurrency for crash safety). WAL checkpoint on clean exit covers
+    the common case (SIGTERM, SIGINT, normal exit).
+    """
+    try:
+        if _kg and _kg._connection:
+            _kg._connection.execute("PRAGMA wal_checkpoint(TRUNCATE)")
+            _kg._connection.close()
+            logger.info("KG WAL checkpointed and connection closed")
+    except Exception:
+        pass
+    try:
+        if _client_cache:
+            del _client_cache
+            logger.info("ChromaDB client released")
+    except Exception:
+        pass
+
+
+import atexit
+atexit.register(_shutdown)
+
+
 def main():
     logger.info("MemPalace MCP Server starting...")
     while True: