From 84e2aa16e42d86b96cd74f48cc70ced259e77fec Mon Sep 17 00:00:00 2001
From: jp <jp@jphein.com>
Date: Sat, 11 Apr 2026 19:44:49 -0700
Subject: [PATCH] perf: graph cache with write-invalidation in build_graph()
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

build_graph() scans every drawer's metadata in 1000-item batches on
every call — O(n) per graph build with no caching. At 50K+ drawers
this costs several seconds per MCP tool call (traverse, find_tunnels,
graph_stats all call build_graph on every invocation).

Add a module-level cache (nodes + edges + timestamp) with a 60-second
TTL. Cache is invalidated via invalidate_graph_cache(), exported for
write operations to call. Tests updated with setup_method cache resets
and two new tests verifying cache hit and invalidation behaviour.
---
 mempalace/palace_graph.py  | 27 +++++++++++++++++++++++++++
 tests/test_palace_graph.py | 37 +++++++++++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+)

diff --git a/mempalace/palace_graph.py b/mempalace/palace_graph.py
index 71cad89..faeedd3 100644
--- a/mempalace/palace_graph.py
+++ b/mempalace/palace_graph.py
@@ -18,6 +18,7 @@ No external graph DB needed — built from ChromaDB metadata.
 import hashlib
 import json
 import os
+import time
 from collections import Counter, defaultdict
 from datetime import datetime, timezone
 
@@ -25,6 +26,20 @@ from .config import MempalaceConfig
 from .palace import get_collection as _get_palace_collection
 from .palace import mine_lock
 
+# Module-level graph cache — mirrors _metadata_cache pattern in mcp_server.py
+_graph_cache_nodes = None
+_graph_cache_edges = None
+_graph_cache_time = 0.0
+_GRAPH_CACHE_TTL = 60.0  # seconds — graph changes less often than metadata
+
+
+def invalidate_graph_cache():
+    """Clear the graph cache. Called from mcp_server.py on writes."""
+    global _graph_cache_nodes, _graph_cache_edges, _graph_cache_time
+    _graph_cache_nodes = None
+    _graph_cache_edges = None
+    _graph_cache_time = 0.0
+
 
 def _get_collection(config=None):
     config = config or MempalaceConfig()
@@ -42,10 +57,18 @@ def build_graph(col=None, config=None):
     """
     Build the palace graph from ChromaDB metadata.
 
+    Returns cached result if fresh (within TTL). Cache is invalidated
+    on writes via invalidate_graph_cache().
+
     Returns:
         nodes: dict of {room: {wings: set, halls: set, count: int}}
         edges: list of {room, wing_a, wing_b, hall} — one per tunnel crossing
     """
+    global _graph_cache_nodes, _graph_cache_edges, _graph_cache_time
+    now = time.time()
+    if _graph_cache_nodes is not None and (now - _graph_cache_time) < _GRAPH_CACHE_TTL:
+        return _graph_cache_nodes, _graph_cache_edges
+
     if col is None:
         col = _get_collection(config)
     if not col:
@@ -101,6 +124,10 @@ def build_graph(col=None, config=None):
             "dates": sorted(data["dates"])[-5:] if data["dates"] else [],
         }
 
+    _graph_cache_nodes = nodes
+    _graph_cache_edges = edges
+    _graph_cache_time = time.time()
+
     return nodes, edges
 
 
diff --git a/tests/test_palace_graph.py b/tests/test_palace_graph.py
index ddda272..37497c1 100644
--- a/tests/test_palace_graph.py
+++ b/tests/test_palace_graph.py
@@ -30,6 +30,7 @@ with patch.dict("sys.modules", {"chromadb": MagicMock()}):
         build_graph,
         find_tunnels,
         graph_stats,
+        invalidate_graph_cache,
         traverse,
     )
 
@@ -38,6 +39,9 @@ with patch.dict("sys.modules", {"chromadb": MagicMock()}):
 
 
 class TestBuildGraph:
+    def setup_method(self):
+        invalidate_graph_cache()
+
     def test_empty_collection(self):
         col = _make_fake_collection([])
         nodes, edges = build_graph(col=col)
@@ -114,11 +118,38 @@ class TestBuildGraph:
         nodes, _ = build_graph(col=col)
         assert len(nodes["busy"]["dates"]) <= 5
 
+    def test_cache_returns_same_result(self):
+        """Second call within TTL returns cached nodes without re-scanning."""
+        col = _make_fake_collection(
+            [{"room": "auth", "wing": "wing_code", "hall": "security", "date": "2026-01-01"}]
+        )
+        nodes1, edges1 = build_graph(col=col)
+        # Second call with a *different* collection — should still return cached result
+        col2 = _make_fake_collection([])
+        nodes2, edges2 = build_graph(col=col2)
+        assert nodes1 == nodes2
+        assert edges1 == edges2
+
+    def test_invalidate_clears_cache(self):
+        """invalidate_graph_cache() forces a fresh scan on next call."""
+        col = _make_fake_collection(
+            [{"room": "auth", "wing": "wing_code", "hall": "security", "date": "2026-01-01"}]
+        )
+        build_graph(col=col)
+        invalidate_graph_cache()
+        col_empty = _make_fake_collection([])
+        nodes, edges = build_graph(col=col_empty)
+        assert nodes == {}
+        assert edges == []
+
 
 # --- traverse ---
 
 
 class TestTraverse:
+    def setup_method(self):
+        invalidate_graph_cache()
+
     def _build_col(self):
         return _make_fake_collection(
             [
@@ -156,6 +187,9 @@ class TestTraverse:
 
 
 class TestFindTunnels:
+    def setup_method(self):
+        invalidate_graph_cache()
+
     def _build_tunnel_col(self):
         return _make_fake_collection(
             [
@@ -192,6 +226,9 @@ class TestFindTunnels:
 
 
 class TestGraphStats:
+    def setup_method(self):
+        invalidate_graph_cache()
+
     def test_empty_graph(self):
         col = _make_fake_collection([])
         stats = graph_stats(col=col)