From 2e9ed9925d642133ffd0039f4047d20301bd53b4 Mon Sep 17 00:00:00 2001 From: Kartik Gupta Date: Tue, 28 Jul 2026 23:07:32 +0530 Subject: [PATCH] fix(serve): bound multi-project graph contexts (#2268) --- README.md | 1 + graphify/serve.py | 138 +++++++++++++++++++++++++++------------ tests/test_serve_http.py | 66 ++++++++++++++++++- 3 files changed, 161 insertions(+), 44 deletions(-) diff --git a/README.md b/README.md index a4f188f46..a5918cad1 100644 --- a/README.md +++ b/README.md @@ -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 | --- diff --git a/graphify/serve.py b/graphify/serve.py index f32a91673..025fda41c 100644 --- a/graphify/serve.py +++ b/graphify/serve.py @@ -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 _select_graph(project_path) -> None: 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") diff --git a/tests/test_serve_http.py b/tests/test_serve_http.py index 7d542901b..7893a0f4e 100644 --- a/tests/test_serve_http.py +++ b/tests/test_serve_http.py @@ -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 ``/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: