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>
This commit is contained in:
jp
2026-04-12 14:52:53 -07:00
parent 1657a79649
commit 8adf35a13c
+18 -9
View File
@@ -18,6 +18,7 @@ No external graph DB needed — built from ChromaDB metadata.
import hashlib import hashlib
import json import json
import os import os
import threading
import time import time
from collections import Counter, defaultdict from collections import Counter, defaultdict
from datetime import datetime, timezone from datetime import datetime, timezone
@@ -28,6 +29,7 @@ from .palace import mine_lock
# Module-level graph cache with TTL and write-invalidation. # Module-level graph cache with TTL and write-invalidation.
# Warm cache serves build_graph() in O(1); invalidate_graph_cache() clears on writes. # Warm cache serves build_graph() in O(1); invalidate_graph_cache() clears on writes.
_graph_cache_lock = threading.Lock()
_graph_cache_nodes = None _graph_cache_nodes = None
_graph_cache_edges = None _graph_cache_edges = None
_graph_cache_time = 0.0 _graph_cache_time = 0.0
@@ -37,9 +39,10 @@ _GRAPH_CACHE_TTL = 60.0 # seconds — graph changes less often than metadata
def invalidate_graph_cache(): def invalidate_graph_cache():
"""Clear the graph cache. Called from mcp_server.py on writes.""" """Clear the graph cache. Called from mcp_server.py on writes."""
global _graph_cache_nodes, _graph_cache_edges, _graph_cache_time global _graph_cache_nodes, _graph_cache_edges, _graph_cache_time
_graph_cache_nodes = None with _graph_cache_lock:
_graph_cache_edges = None _graph_cache_nodes = None
_graph_cache_time = 0.0 _graph_cache_edges = None
_graph_cache_time = 0.0
def _get_collection(config=None): def _get_collection(config=None):
@@ -59,7 +62,11 @@ def build_graph(col=None, config=None):
Build the palace graph from ChromaDB metadata. Build the palace graph from ChromaDB metadata.
Returns cached result if fresh (within TTL). Cache is invalidated Returns cached result if fresh (within TTL). Cache is invalidated
on writes via invalidate_graph_cache(). on writes via invalidate_graph_cache(). Thread-safe via _graph_cache_lock.
Note: warm cache ignores ``col`` and ``config`` arguments — this is
intentional for the MCP server's single-palace use case. Callers
switching collections should call ``invalidate_graph_cache()`` first.
Returns: Returns:
nodes: dict of {room: {wings: set, halls: set, count: int}} nodes: dict of {room: {wings: set, halls: set, count: int}}
@@ -69,8 +76,9 @@ def build_graph(col=None, config=None):
now = time.time() now = time.time()
# NOTE: warm cache ignores col/config args — intentional for the MCP server's # NOTE: warm cache ignores col/config args — intentional for the MCP server's
# single-palace use case. Callers switching collections must invalidate first. # single-palace use case. Callers switching collections must invalidate first.
if _graph_cache_nodes is not None and (now - _graph_cache_time) < _GRAPH_CACHE_TTL: with _graph_cache_lock:
return _graph_cache_nodes, _graph_cache_edges 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: if col is None:
col = _get_collection(config) col = _get_collection(config)
@@ -130,9 +138,10 @@ def build_graph(col=None, config=None):
# Only cache non-empty graphs so new data is picked up immediately # Only cache non-empty graphs so new data is picked up immediately
# when the palace is first populated. # when the palace is first populated.
if nodes: if nodes:
_graph_cache_nodes = nodes with _graph_cache_lock:
_graph_cache_edges = edges _graph_cache_nodes = nodes
_graph_cache_time = time.time() _graph_cache_edges = edges
_graph_cache_time = time.time()
return nodes, edges return nodes, edges