beyhan/graphify
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(serve): bound multi-project graph contexts (#2268)

Browse Source
This commit is contained in:
Kartik Gupta
2026-07-29 17:02:00 +01:00
committed by safishamsi
parent 6d6b674b20
commit b4865ffcbf
3 changed files with 161 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+1
View File
@@ -524,6 +524,7 @@ These are only needed for **headless / CI extraction** (`graphify extract`). Whe
| `GRAPHIFY_QUERY_LOG_DISABLE` | Set to `1` to force the query log off (wins over the enable vars) | optional |
| `GRAPHIFY_QUERY_LOG_RESPONSES` | When the log is enabled, also record full subgraph responses (off by default) | optional |
| `GRAPHIFY_MAX_GRAPH_BYTES` | Override the 512 MiB graph.json size cap — e.g. `700MB`, `2GB`, or plain bytes | optional — useful for very large corpora |
| `GRAPHIFY_MAX_CONTEXTS` | Maximum number of non-default project graphs retained by one multi-project MCP server | optional — default: `8`; invalid values use `8`, and values below `1` use `1` |
| `GRAPHIFY_LLM_TEMPERATURE` | Override LLM temperature for semantic extraction — e.g. `0.7`, or `none` to omit | optional — auto-omitted for o1/o3/o4/gpt-5 reasoning models |
---
graphify/serve.py
+96 -42
View File
@@ -2,10 +2,13 @@
from __future__ import annotations
import json
import math
import os
import re
import sys
from array import array
from collections import OrderedDict
from pathlib import Path
import threading
from typing import NamedTuple
import networkx as nx
from networkx.readwrite import json_graph
@@ -73,6 +76,85 @@ def _communities_from_graph(G: nx.Graph) -> dict[int, list[str]]:
return communities
def _max_server_contexts() -> int:
"""Return the project-context LRU capacity (default 8, minimum 1).
``GRAPHIFY_MAX_CONTEXTS`` overrides the default. Invalid or blank values
use 8; zero and negative values clamp to 1, since each request needs a
graph context. The server's configured default graph is pinned separately
and does not count against this limit.
"""
raw = os.environ.get("GRAPHIFY_MAX_CONTEXTS", "").strip()
if not raw:
return 8
try:
return max(1, int(raw))
except ValueError:
return 8
class _GraphContextCache:
"""Thread-safe graph contexts: one pinned default plus an LRU of projects."""
def __init__(self, max_contexts: int):
self._max_contexts = max_contexts
self._entries: OrderedDict[str, dict] = OrderedDict()
self._pinned: dict[str, dict] = {}
self._lock = threading.Lock()
def _load_entry(self, resolved_path: str, key: tuple[int, int]) -> dict:
"""Build one entry for an already-resolved path and known file key.
``_load_graph`` is also used by the CLI, where invalid input terminates
the process. A client-supplied ``project_path`` must instead become a
tool error, so the shared MCP server can continue serving other graphs.
"""
try:
graph = _load_graph(resolved_path)
except SystemExit as exc:
raise RuntimeError(f"could not load graph.json at {resolved_path}") from exc
# Warm the index before exposing the graph so its first query does not
# pay the expensive build cost.
_get_trigram_index(graph)
communities = _communities_from_graph(graph)
entry = {
"key": key,
"G": graph,
"communities": communities,
}
return entry
def load(self, resolved_path: str, *, pinned: bool = False) -> tuple[nx.Graph, dict[int, list[str]]]:
"""Return a fresh context, retaining project contexts by LRU order.
``resolved_path`` is resolved by the caller, making this method the
sole owner of file statting and cache-key construction.
``pinned=True`` is reserved for the server's configured default graph;
it remains warm without consuming a project-cache slot.
"""
with self._lock:
try:
stat_result = Path(resolved_path).stat()
except FileNotFoundError:
raise FileNotFoundError(f"graph.json not found: {resolved_path}") from None
key = (stat_result.st_mtime_ns, stat_result.st_size)
entries = self._pinned if pinned else self._entries
entry = entries.get(resolved_path)
if entry is not None and entry["key"] == key:
if not pinned:
self._entries.move_to_end(resolved_path)
return entry["G"], entry["communities"]
entry = self._load_entry(resolved_path, key)
entries[resolved_path] = entry
if not pinned:
self._entries.move_to_end(resolved_path)
while len(self._entries) > self._max_contexts:
self._entries.popitem(last=False)
return entry["G"], entry["communities"]
def _strip_diacritics(text: str | None) -> str:
import unicodedata
if not isinstance(text, str):
@@ -1063,9 +1145,6 @@ def _filter_blank_stdin() -> None:
JSONRPCMessage, so a bare newline triggers a Pydantic ValidationError.
This installs an OS-level pipe that relays stdin while dropping blanks.
"""
import os
import threading
r_fd, w_fd = os.pipe()
saved_fd = os.dup(sys.stdin.fileno())
@@ -1108,8 +1187,6 @@ def _build_server(graph_path: str):
Streamable HTTP) and runs it. Hot-reload of graph.json works the same way
regardless of transport, since reloads happen inside the tool handlers.
"""
import threading
try:
from mcp.server import Server
from mcp import types
@@ -1119,44 +1196,21 @@ def _build_server(graph_path: str):
from graphify import paths as _paths
# Per-graph context cache: resolved graph.json path -> {key, G, communities}.
# The server's default graph is just the first entry; a tool call carrying a
# project_path adds its own. Routing every graph through one cache means the
# eager trigram index and the mtime+size hot-reload behave identically for
# the default graph and for any project graph.
_default_graph_path = graph_path
_ctx_lock = threading.Lock()
_ctx_cache: dict[str, dict] = {}
# Graph contexts comprise one pinned configured default plus a bounded LRU
# of project_path graphs. This preserves the configured graph's warm index
# while preventing a shared server from retaining every project it serves.
_default_graph_path = str(Path(graph_path).resolve())
_ctx_cache = _GraphContextCache(_max_server_contexts())
def _load_ctx(path: str):
"""Return (G, communities) for a graph.json path, reusing a cached
context until the file's (mtime, size) changes and then transparently
rebuilding it. Unlike ``_load_graph`` it never exits the process on a
missing/corrupt file — it raises, so a bad project_path surfaces as a
tool error instead of killing a server that is happily serving other
projects."""
try:
s = Path(path).stat()
key = (s.st_mtime_ns, s.st_size)
except FileNotFoundError:
raise FileNotFoundError(f"graph.json not found: {path}")
ent = _ctx_cache.get(path)
if ent is not None and ent["key"] == key:
return ent["G"], ent["communities"]
with _ctx_lock:
ent = _ctx_cache.get(path)
if ent is not None and ent["key"] == key:
return ent["G"], ent["communities"] # another thread built it
try:
new_G = _load_graph(path)
except SystemExit as e: # _load_graph exits on missing/corrupt file
raise RuntimeError(f"could not load graph.json at {path}") from e
# Warm the trigram index before exposing the graph so the first query
# against it is fast (same rationale as the original startup warm-up).
_get_trigram_index(new_G)
comm = _communities_from_graph(new_G)
_ctx_cache[path] = {"key": key, "G": new_G, "communities": comm}
return new_G, comm
"""Return the current default or project graph context as a tool error.
Unlike ``_load_graph``, this never lets a missing or corrupt client
graph terminate the MCP process; it raises so other projects remain
available on the same server.
"""
resolved_path = str(Path(path).resolve())
return _ctx_cache.load(resolved_path, pinned=resolved_path == _default_graph_path)
def _resolve_graph_path(project_path) -> str:
"""Map an optional project_path to a concrete graph.json path. ``None``
@@ -1185,7 +1239,7 @@ def _build_server(graph_path: str):
nonlocal G, communities, active_graph_path
path = _resolve_graph_path(project_path)
G, communities = _load_ctx(path)
active_graph_path = path
active_graph_path = str(Path(path).resolve())
server = Server("graphify")
tests/test_serve_http.py
+64 -2
View File
@@ -170,9 +170,9 @@ def test_tools_list_over_http(tmp_path):
assert {"query_graph", "get_node", "graph_stats"} <= names
def _project_with_graph(tmp_path, node_count: int) -> str:
def _project_with_graph(tmp_path, node_count: int, name: str = "proj") -> str:
"""Create ``<proj>/graphify-out/graph.json`` and return the project dir."""
proj = tmp_path / "proj"
proj = tmp_path / name
(proj / "graphify-out").mkdir(parents=True)
graph = {
"directed": True,
@@ -227,6 +227,56 @@ def test_project_path_routes_to_that_projects_graph(tmp_path):
assert "Nodes: 2" in _call_tool(client, headers, "graph_stats", {}, rid=4)
@pytest.mark.parametrize(
("value", "expected"),
[(None, 8), ("", 8), ("bad", 8), ("0", 1), ("-4", 1), ("3", 3)],
)
def test_max_server_contexts_parsing(monkeypatch, value, expected):
if value is None:
monkeypatch.delenv("GRAPHIFY_MAX_CONTEXTS", raising=False)
else:
monkeypatch.setenv("GRAPHIFY_MAX_CONTEXTS", value)
assert serve_mod._max_server_contexts() == expected
def test_project_context_cache_is_lru_and_pins_default_graph(tmp_path, monkeypatch):
"""Project contexts hit, promote, and evict without evicting the default."""
monkeypatch.setenv("GRAPHIFY_MAX_CONTEXTS", "2")
original_load = serve_mod._load_graph
loads: dict[str, int] = {}
def counting_load(path: str):
resolved = str(Path(path).resolve())
loads[resolved] = loads.get(resolved, 0) + 1
return original_load(path)
monkeypatch.setattr(serve_mod, "_load_graph", counting_load)
projects = [
_project_with_graph(tmp_path, node_count=i + 3, name=f"project-{i}")
for i in range(3)
]
default_graph = _graph_file(tmp_path)
app = serve_mod._build_http_app(default_graph, json_response=True)
with _client(app) as client:
headers = _init_session(client)
assert "Nodes: 3" in _call_tool(client, headers, "graph_stats", {"project_path": projects[0]}, rid=2)
assert "Nodes: 4" in _call_tool(client, headers, "graph_stats", {"project_path": projects[1]}, rid=3)
# A cache hit promotes project-0 above project-1 in LRU recency.
assert "Nodes: 3" in _call_tool(client, headers, "graph_stats", {"project_path": projects[0]}, rid=4)
assert "Nodes: 5" in _call_tool(client, headers, "graph_stats", {"project_path": projects[2]}, rid=5)
# project-1, not the re-touched project-0, was evicted.
assert "Nodes: 4" in _call_tool(client, headers, "graph_stats", {"project_path": projects[1]}, rid=6)
# The configured default graph stays warm even when project capacity is full.
assert "Nodes: 2" in _call_tool(client, headers, "graph_stats", {}, rid=7)
first_graph = str((Path(projects[0]) / "graphify-out" / "graph.json").resolve())
second_graph = str((Path(projects[1]) / "graphify-out" / "graph.json").resolve())
default_graph = str(Path(default_graph).resolve())
assert loads[first_graph] == 1
assert loads[second_graph] == 2
assert loads[default_graph] == 1
def test_bad_project_path_errors_without_killing_server(tmp_path):
"""A missing project graph is a tool error, not a process exit — the server
keeps serving the default graph."""
@@ -239,6 +289,18 @@ def test_bad_project_path_errors_without_killing_server(tmp_path):
assert "Nodes: 2" in _call_tool(client, headers, "graph_stats", {}, rid=3)
def test_corrupt_project_graph_is_a_tool_error_without_killing_server(tmp_path):
"""A CLI-style SystemExit from a client graph cannot stop the MCP server."""
project = Path(_project_with_graph(tmp_path, node_count=3))
(project / "graphify-out" / "graph.json").write_text("{not json", encoding="utf-8")
app = serve_mod._build_http_app(_graph_file(tmp_path), json_response=True)
with _client(app) as client:
headers = _init_session(client)
bad = _call_tool(client, headers, "graph_stats", {"project_path": str(project)}, rid=2)
assert "could not load graph.json" in bad
assert "Nodes: 2" in _call_tool(client, headers, "graph_stats", {}, rid=3)
def test_stateless_mode_initialize(tmp_path):
app = serve_mod._build_http_app(_graph_file(tmp_path), stateless=True, json_response=True)
with _client(app) as client: