diff --git a/graphify/serve.py b/graphify/serve.py index c97da5ed0..19e0c3923 100644 --- a/graphify/serve.py +++ b/graphify/serve.py @@ -1345,17 +1345,24 @@ def _traversal_view(G: nx.Graph) -> nx.Graph: fold there too; letting networkx assign the keys keeps both, and nothing downstream reads the key. - A fresh copy per query rather than a cached one: `_filter_graph_by_context` - already copies per query when a filter applies, and the copy shares node - data dicts with `G`, so only the edge dicts are duplicated. + Built once per graph object and cached on `G.graph`, like the trigram index. + It used to be rebuilt on every query, copying every edge of the graph, which + was most of a warm `query_graph` call. Sharing it is safe because nothing + downstream writes to it: `_filter_graph_by_context` builds a new graph when + a filter applies, and the traversals and the renderer only read. A hot + reload swaps in a fresh `G`, so a cached view never outlives its graph. """ if not G.is_directed(): return G + cached = G.graph.get("_traversal_view") + if cached is not None: + return cached H = nx.MultiGraph() if G.is_multigraph() else nx.Graph() H.graph.update(G.graph) H.add_nodes_from(G.nodes(data=True)) for u, v, d in G.edges(data=True): H.add_edge(u, v, **{**d, "_src": d.get("_src", u), "_tgt": d.get("_tgt", v)}) + G.graph["_traversal_view"] = H return H def _query_graph_text( @@ -1727,6 +1734,55 @@ def _resolve_path_endpoint( return _pick_scored_endpoint(G, scored, query), scored, None +def _path_search_graph(G: nx.Graph, undirected: bool) -> nx.Graph: + """The sorted, materialized graph `_shortest_path_text` searches. + + Deterministic path (#2074): the hash-seeded undirected view picked an + arbitrary route among equal-length paths, so the search runs on a graph + built from sorted node and edge lists, which makes the chosen path + canonical. Serve's shared G is left untouched (its degree feeds query-seed + tie-breaks). + + Built once per graph object and direction mode, and cached on `G.graph` + like the trigram index: rebuilding it from every edge was nearly all of a + `shortest_path` call. Only `nx.shortest_path` reads it, and a hot reload + swaps in a fresh `G`, so a cached graph never outlives its source. + """ + cache = G.graph.setdefault("_path_search_graphs", {}) + cached = cache.get(undirected) + if cached is not None: + return cached + if undirected: + H = nx.Graph() + H.add_nodes_from(sorted(G.nodes)) + H.add_edges_from(sorted((min(u, v), max(u, v)) for u, v in G.edges())) + else: + # Directed by default (#2487). True direction is NOT raw arc + # order: legacy canonicalized files persist a flipped arc with + # _src/_tgt markers (#2309), so build the digraph from _src/_tgt + # (falling back to the loaded arc) rather than to_directed(). + H = nx.DiGraph() + H.add_nodes_from(sorted(G.nodes)) + H.add_edges_from(sorted( + (d.get("_src", u), d.get("_tgt", v)) for u, v, d in G.edges(data=True) + )) + # A `contains` edge only runs file -> symbol; there is no stored + # edge back out to the containing file, so a route that reaches a + # symbol (via an `imports`/`calls`/`references` hop) can never + # continue on to the file that defines it, and a file-to-file + # dependency routed through a shared symbol finds no path at all + # even though both halves exist (#3878). Add the implied reverse hop + # for traversal only — the printed segment still recovers the real + # stored `contains` edge and its true direction from G below. + H.add_edges_from(sorted( + (d.get("_tgt", v), d.get("_src", u)) + for u, v, d in G.edges(data=True) + if d.get("relation") == "contains" + )) + cache[undirected] = H + return H + + def _shortest_path_text(G: nx.Graph, arguments: dict) -> str: """Body of the `shortest_path` MCP tool (module-level so tests can call it without an mcp install). @@ -1769,39 +1825,7 @@ def _shortest_path_text(G: nx.Graph, arguments: dict) -> str: max_hops = int(arguments.get("max_hops", 8)) undirected = bool(arguments.get("undirected", False)) try: - # Deterministic path (#2074): the hash-seeded undirected view picked an - # arbitrary route among equal-length paths. Build a sorted, materialized - # graph so the chosen path is canonical. Serve's shared G is left - # untouched (its degree feeds query-seed tie-breaks). - if undirected: - _und = nx.Graph() - _und.add_nodes_from(sorted(G.nodes)) - _und.add_edges_from(sorted((min(u, v), max(u, v)) for u, v in G.edges())) - path_nodes = nx.shortest_path(_und, src_nid, tgt_nid) - else: - # Directed by default (#2487). True direction is NOT raw arc - # order: legacy canonicalized files persist a flipped arc with - # _src/_tgt markers (#2309), so build the digraph from _src/_tgt - # (falling back to the loaded arc) rather than to_directed(). - _dg = nx.DiGraph() - _dg.add_nodes_from(sorted(G.nodes)) - _dg.add_edges_from(sorted( - (d.get("_src", u), d.get("_tgt", v)) for u, v, d in G.edges(data=True) - )) - # A `contains` edge only runs file -> symbol; there is no stored - # edge back out to the containing file, so a route that reaches a - # symbol (via an `imports`/`calls`/`references` hop) can never - # continue on to the file that defines it, and a file-to-file - # dependency routed through a shared symbol finds no path at all - # even though both halves exist (#3878). Add the implied reverse hop - # for traversal only — the printed segment still recovers the real - # stored `contains` edge and its true direction from G below. - _dg.add_edges_from(sorted( - (d.get("_tgt", v), d.get("_src", u)) - for u, v, d in G.edges(data=True) - if d.get("relation") == "contains" - )) - path_nodes = nx.shortest_path(_dg, src_nid, tgt_nid) + path_nodes = nx.shortest_path(_path_search_graph(G, undirected), src_nid, tgt_nid) except (nx.NetworkXNoPath, nx.NodeNotFound): src_label = G.nodes[src_nid].get("label", src_nid) tgt_label = G.nodes[tgt_nid].get("label", tgt_nid) diff --git a/tests/test_mcp_graph_view_cache.py b/tests/test_mcp_graph_view_cache.py new file mode 100644 index 000000000..b4c2c72e4 --- /dev/null +++ b/tests/test_mcp_graph_view_cache.py @@ -0,0 +1,124 @@ +"""The MCP server's traversal view and shortest-path graphs are built once per graph. + +`query_graph` used to copy every edge of the loaded graph into an undirected +view on every call, and `shortest_path` rebuilt a sorted DiGraph of every edge +on every call. Both are now cached on `G.graph`, like the trigram index. These +tests pin what makes that safe: the cached structure is reused, answers equal a +freshly built one in any order of calls, and a reloaded graph (a new object) +never sees the previous graph's structures. +""" +import json + +import networkx as nx +from networkx.readwrite import json_graph + +from graphify.serve import ( + _load_graph, + _path_search_graph, + _query_graph_text, + _shortest_path_text, + _traversal_view, +) + + +def _write(tmp_path, edges, name="graph.json"): + G = nx.Graph() + for n in sorted({n for e in edges for n in e[:2]}): + G.add_node(n, label=f"{n}_fn", source_file=f"{n}.py", source_location="L1", community=0) + data = json_graph.node_link_data(G, edges="links") + data["links"] = [ + {"source": u, "target": v, "relation": rel, "confidence": "EXTRACTED", "context": ctx} + for u, v, rel, ctx in edges + ] + path = tmp_path / name + path.write_text(json.dumps(data)) + return path + + +EDGES = [ + ("main", "parse", "calls", "call"), + ("parse", "lexer", "calls", "call"), + ("parse", "ast", "imports", "import"), + ("main", "config", "imports", "import"), + ("config", "lexer", "calls", "call"), + ("report", "main", "calls", "call"), +] +QUERIES = [ + ("parse_fn", "bfs", None), + ("lexer_fn", "dfs", None), + ("what calls lexer_fn", "bfs", None), + ("config_fn", "bfs", ["import"]), + ("main_fn", "dfs", ["call"]), +] + + +def _fresh(G): + G.graph.pop("_traversal_view", None) + G.graph.pop("_path_search_graphs", None) + + +def test_traversal_view_is_built_once_per_graph(tmp_path): + G = _load_graph(str(_write(tmp_path, EDGES))) + assert G.is_directed() + assert _traversal_view(G) is _traversal_view(G) + + +def test_cached_view_answers_like_a_fresh_one_in_any_order(tmp_path): + G = _load_graph(str(_write(tmp_path, EDGES))) + expected = [] + for q, mode, ctx in QUERIES: + _fresh(G) + expected.append(_query_graph_text(G, q, mode=mode, depth=3, context_filters=ctx)) + _fresh(G) + for order in (QUERIES, list(reversed(QUERIES)), QUERIES): + for q, mode, ctx in order: + got = _query_graph_text(G, q, mode=mode, depth=3, context_filters=ctx) + assert got == expected[QUERIES.index((q, mode, ctx))], (q, mode, ctx) + + +def test_reloaded_graph_gets_its_own_view(tmp_path): + path = _write(tmp_path, EDGES) + old = _load_graph(str(path)) + before = _query_graph_text(old, "ast_fn", depth=1) + assert "NODE extra_fn" not in before + _write(tmp_path, EDGES + [("extra", "ast", "calls", "call")]) + new = _load_graph(str(path)) + assert _traversal_view(new) is not _traversal_view(old) + assert "NODE extra_fn" in _query_graph_text(new, "ast_fn", depth=1) + assert _query_graph_text(old, "ast_fn", depth=1) == before + + +def test_path_search_graphs_are_built_once_per_direction(tmp_path): + G = _load_graph(str(_write(tmp_path, EDGES))) + directed = _path_search_graph(G, False) + undirected = _path_search_graph(G, True) + assert directed is _path_search_graph(G, False) + assert undirected is _path_search_graph(G, True) + assert directed.is_directed() and not undirected.is_directed() + + +def test_cached_shortest_path_answers_like_a_fresh_one(tmp_path): + G = _load_graph(str(_write(tmp_path, EDGES))) + cases = [ + {"source": "report_fn", "target": "lexer_fn"}, + {"source": "lexer_fn", "target": "report_fn"}, + {"source": "lexer_fn", "target": "report_fn", "undirected": True}, + {"source": "ast_fn", "target": "config_fn", "undirected": True}, + ] + expected = [] + for c in cases: + _fresh(G) + expected.append(_shortest_path_text(G, dict(c))) + _fresh(G) + for _ in range(2): + assert [_shortest_path_text(G, dict(c)) for c in cases] == expected + + +def test_reloaded_graph_gets_its_own_path_graph(tmp_path): + path = _write(tmp_path, EDGES) + old = _load_graph(str(path)) + assert "No directed path" in _shortest_path_text(old, {"source": "lexer_fn", "target": "ast_fn"}) + _write(tmp_path, EDGES + [("lexer", "ast", "calls", "call")]) + new = _load_graph(str(path)) + assert "Shortest path (1 hops)" in _shortest_path_text(new, {"source": "lexer_fn", "target": "ast_fn"}) + assert "No directed path" in _shortest_path_text(old, {"source": "lexer_fn", "target": "ast_fn"})