Skip to content

Commit 509120c

Browse files
jpheinclaude
andcommitted
fix: add threading lock to graph cache, expand docstring
Address review feedback from @bensig: 1. Wrap cache reads/writes in threading.Lock for thread safety 2. Promote the col-arg caveat from inline comment to docstring Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
1 parent c6f8ab3 commit 509120c

1 file changed

Lines changed: 18 additions & 9 deletions

File tree

mempalace/palace_graph.py

Lines changed: 18 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
No external graph DB needed — built from ChromaDB metadata.
1616
"""
1717

18+
import threading
1819
import time
1920
from collections import defaultdict, Counter
2021

@@ -23,6 +24,7 @@
2324

2425
# Module-level graph cache with TTL and write-invalidation.
2526
# Warm cache serves build_graph() in O(1); invalidate_graph_cache() clears on writes.
27+
_graph_cache_lock = threading.Lock()
2628
_graph_cache_nodes = None
2729
_graph_cache_edges = None
2830
_graph_cache_time = 0.0
@@ -32,9 +34,10 @@
3234
def invalidate_graph_cache():
3335
"""Clear the graph cache. Called from mcp_server.py on writes."""
3436
global _graph_cache_nodes, _graph_cache_edges, _graph_cache_time
35-
_graph_cache_nodes = None
36-
_graph_cache_edges = None
37-
_graph_cache_time = 0.0
37+
with _graph_cache_lock:
38+
_graph_cache_nodes = None
39+
_graph_cache_edges = None
40+
_graph_cache_time = 0.0
3841

3942

4043
def _get_collection(config=None):
@@ -54,7 +57,11 @@ def build_graph(col=None, config=None):
5457
Build the palace graph from ChromaDB metadata.
5558
5659
Returns cached result if fresh (within TTL). Cache is invalidated
57-
on writes via invalidate_graph_cache().
60+
on writes via invalidate_graph_cache(). Thread-safe via _graph_cache_lock.
61+
62+
Note: warm cache ignores ``col`` and ``config`` arguments — this is
63+
intentional for the MCP server's single-palace use case. Callers
64+
switching collections should call ``invalidate_graph_cache()`` first.
5865
5966
Returns:
6067
nodes: dict of {room: {wings: set, halls: set, count: int}}
@@ -64,8 +71,9 @@ def build_graph(col=None, config=None):
6471
now = time.time()
6572
# NOTE: warm cache ignores col/config args — intentional for the MCP server's
6673
# single-palace use case. Callers switching collections must invalidate first.
67-
if _graph_cache_nodes is not None and (now - _graph_cache_time) < _GRAPH_CACHE_TTL:
68-
return _graph_cache_nodes, _graph_cache_edges
74+
with _graph_cache_lock:
75+
if _graph_cache_nodes is not None and (now - _graph_cache_time) < _GRAPH_CACHE_TTL:
76+
return _graph_cache_nodes, _graph_cache_edges
6977

7078
if col is None:
7179
col = _get_collection(config)
@@ -125,9 +133,10 @@ def build_graph(col=None, config=None):
125133
# Only cache non-empty graphs so new data is picked up immediately
126134
# when the palace is first populated.
127135
if nodes:
128-
_graph_cache_nodes = nodes
129-
_graph_cache_edges = edges
130-
_graph_cache_time = time.time()
136+
with _graph_cache_lock:
137+
_graph_cache_nodes = nodes
138+
_graph_cache_edges = edges
139+
_graph_cache_time = time.time()
131140

132141
return nodes, edges
133142

0 commit comments

Comments
 (0)