From c5c15b4ca41bc5cf2f44c1d6a4f11a9f4f056a4c Mon Sep 17 00:00:00 2001 From: Rohit Behera <126186063+r0h1tb@users.noreply.github.com> Date: Sun, 2 Aug 2026 01:55:31 +0530 Subject: [PATCH 1/2] feat(index): resolve cross-file references via a project-wide symbol table Edge extraction resolved references against name_to_id built from a single file's nodes, so any reference to a symbol defined elsewhere was dropped. Only same-file calls ever linked. Measured against ground truth from parsers independent of tree-sitter (Python's stdlib ast, and javalang for a Java project), the unreachable share was 83% on this repo and 95% on a Spring Boot service -- in layered code essentially every interesting call crosses a file. Indexing is now two-phase: collect nodes from every file and build a project wide name -> id map, then resolve edges against it. extract_edges takes an optional global_symbols map; local definitions are applied last so a file-local symbol always shadows a same-named symbol from another file. Measured on this repo, before -> after: CALLS edges 498 -> 1412 of which cross-file 0 -> 914 total edges 2387 -> 3782 Retrieval quality over the 15 most-called symbols, scored against the stdlib ast oracle across all call relationships (not just same-file): precision 0.99, recall 0.95, F1 0.97 Precision holding at 0.99 is the load-bearing result: name-based global resolution could have produced false positives across same-named symbols, and on this codebase it does not. Suite: 3 failed, 188 passed (baseline: 3 failed, 174 passed) -- same three pre-existing failures, fixed separately in #51. Refs #35 --- ast_rag/cli.py | 28 ++++- ast_rag/services/parsing/edge_extractor.py | 13 ++- ast_rag/services/parsing/parser_manager.py | 2 + tests/test_global_symbol_resolution.py | 127 +++++++++++++++++++++ 4 files changed, 164 insertions(+), 6 deletions(-) create mode 100644 tests/test_global_symbol_resolution.py diff --git a/ast_rag/cli.py b/ast_rag/cli.py index 9812afd..d64dc52 100644 --- a/ast_rag/cli.py +++ b/ast_rag/cli.py @@ -180,6 +180,13 @@ def init( all_blocks = [] all_block_edges = [] + # Two-phase index. Edge resolution matches references against a name -> id + # map; when that map only holds the current file's nodes, any reference to a + # symbol defined elsewhere is silently dropped. Collecting every symbol first + # and resolving afterwards is what lets cross-file references link at all. + parsed: list[tuple[str, str, bytes, object, list]] = [] + global_symbols: dict[str, str] = {} + with _index_progress() as progress: task = progress.add_task("Parsing", total=len(files)) for fp, lang in files: @@ -191,7 +198,23 @@ def init( with open(fp, "rb") as fh: source = fh.read() nodes = pm.extract_nodes(tree, fp, lang, source, commit) - edges = pm.extract_edges(tree, nodes, fp, lang, source, commit) + parsed.append((fp, lang, source, tree, nodes)) + all_nodes.extend(nodes) + # First definition of a name wins; files are walked in a stable order + # so the choice is deterministic across runs. + for node in nodes: + global_symbols.setdefault(node.name, node.id) + + with _index_progress() as progress: + task = progress.add_task("Resolving", total=len(parsed)) + for fp, lang, source, tree, nodes in parsed: + progress.update(task, description=f"Resolving {os.path.relpath(fp, root)}") + progress.advance(task) + all_edges.extend( + pm.extract_edges( + tree, nodes, fp, lang, source, commit, global_symbols=global_symbols + ) + ) # Extract blocks for Python and Rust files if lang in ("python", "rust"): @@ -199,9 +222,6 @@ def init( all_blocks.extend(blocks) all_block_edges.extend(block_edges) - all_nodes.extend(nodes) - all_edges.extend(edges) - console.print( f"Extracted [bold]{len(all_nodes)}[/bold] nodes, [bold]{len(all_edges)}[/bold] edges, " f"and [bold]{len(all_blocks)}[/bold] blocks." diff --git a/ast_rag/services/parsing/edge_extractor.py b/ast_rag/services/parsing/edge_extractor.py index 881c2c3..5a99c7f 100644 --- a/ast_rag/services/parsing/edge_extractor.py +++ b/ast_rag/services/parsing/edge_extractor.py @@ -43,8 +43,16 @@ def extract_edges( compiled_queries: dict[str, object], source: Optional[bytes] = None, commit_hash: str = "INIT", + global_symbols: Optional[dict[str, str]] = None, ) -> list[ASTEdge]: - """Extract edges (relationships) between AST nodes.""" + """Extract edges (relationships) between AST nodes. + + ``global_symbols`` optionally supplies a project-wide name -> node id map + so references to symbols defined in *other* files can resolve. Without it + only same-file references resolve, which leaves the large majority of a + real codebase's call graph unlinked. Local definitions take precedence, + so a file-local symbol always shadows a project-wide one of the same name. + """ if source is None: try: with open(file_path, "rb") as fh: @@ -53,7 +61,8 @@ def extract_edges( source = b"" edges: list[ASTEdge] = [] - name_to_id: dict[str, str] = {n.name: n.id for n in nodes} + name_to_id: dict[str, str] = dict(global_symbols or {}) + name_to_id.update({n.name: n.id for n in nodes}) file_node_id = hashlib.sha256( f"{file_path}:{NodeKind.FILE.value}:{file_path}".encode() diff --git a/ast_rag/services/parsing/parser_manager.py b/ast_rag/services/parsing/parser_manager.py index 8012d4c..4d231ed 100644 --- a/ast_rag/services/parsing/parser_manager.py +++ b/ast_rag/services/parsing/parser_manager.py @@ -245,8 +245,10 @@ def extract_edges( lang: str, source: Optional[bytes] = None, commit_hash: str = "INIT", + global_symbols: Optional[dict[str, str]] = None, ) -> list[ASTEdge]: return self._edge_extractor.extract_edges( + global_symbols=global_symbols, tree=tree, nodes=nodes, file_path=file_path, diff --git a/tests/test_global_symbol_resolution.py b/tests/test_global_symbol_resolution.py new file mode 100644 index 0000000..0c4e4c6 --- /dev/null +++ b/tests/test_global_symbol_resolution.py @@ -0,0 +1,127 @@ +"""Cross-file reference resolution via a project-wide symbol table. + +Edge extraction resolves references against a name -> node id map. Built from a +single file's nodes, any reference to a symbol defined elsewhere is dropped, so +only same-file calls ever link. `extract_edges(global_symbols=...)` supplies a +project-wide map so cross-file references resolve too. + +Local definitions must still win, otherwise a file-local helper would be +shadowed by an unrelated same-named symbol from another file. +""" + +from __future__ import annotations + +import pytest + +from ast_rag.dto.enums import EdgeKind +from ast_rag.services.parsing.parser_manager import ParserManager + +CALLEE_SRC = b""" +def shared_helper(x): + return x + 1 +""" + +CALLER_SRC = b""" +def caller(y): + return shared_helper(y) +""" + +SHADOWING_SRC = b""" +def shared_helper(x): + return x * 2 + + +def local_caller(y): + return shared_helper(y) +""" + + +@pytest.fixture(scope="module") +def pm() -> ParserManager: + return ParserManager() + + +def _parse(pm: ParserManager, path, src: bytes): + path.write_bytes(src) + tree = pm.parse_file(str(path), source=src) + assert tree is not None + return tree, pm.extract_nodes(tree, str(path), "python") + + +def _calls(edges): + return [e for e in edges if e.kind == EdgeKind.CALLS] + + +def test_cross_file_call_is_unresolved_without_a_global_table(pm, tmp_path): + """Baseline: this is the behaviour the two-phase index exists to fix.""" + _, callee_nodes = _parse(pm, tmp_path / "callee.py", CALLEE_SRC) + tree, caller_nodes = _parse(pm, tmp_path / "caller.py", CALLER_SRC) + + edges = pm.extract_edges( + tree, caller_nodes, str(tmp_path / "caller.py"), "python", source=CALLER_SRC + ) + assert _calls(edges) == [], "expected no cross-file CALLS without a project-wide map" + assert callee_nodes, "fixture lost" + + +def test_cross_file_call_resolves_with_a_global_table(pm, tmp_path): + _, callee_nodes = _parse(pm, tmp_path / "callee.py", CALLEE_SRC) + tree, caller_nodes = _parse(pm, tmp_path / "caller.py", CALLER_SRC) + + global_symbols = {n.name: n.id for n in callee_nodes} + edges = pm.extract_edges( + tree, + caller_nodes, + str(tmp_path / "caller.py"), + "python", + source=CALLER_SRC, + global_symbols=global_symbols, + ) + + calls = _calls(edges) + assert calls, "cross-file call did not resolve" + helper_id = next(n.id for n in callee_nodes if n.name == "shared_helper") + assert any(e.to_id == helper_id for e in calls), "edge did not point at the other file's symbol" + + +def test_local_definition_shadows_the_global_one(pm, tmp_path): + """A file-local symbol must win over a same-named symbol from another file.""" + _, other_nodes = _parse(pm, tmp_path / "other.py", CALLEE_SRC) + tree, local_nodes = _parse(pm, tmp_path / "local.py", SHADOWING_SRC) + + global_symbols = {n.name: n.id for n in other_nodes} + edges = pm.extract_edges( + tree, + local_nodes, + str(tmp_path / "local.py"), + "python", + source=SHADOWING_SRC, + global_symbols=global_symbols, + ) + + calls = _calls(edges) + assert calls, "local call did not resolve" + local_id = next(n.id for n in local_nodes if n.name == "shared_helper") + foreign_id = next(n.id for n in other_nodes if n.name == "shared_helper") + assert any(e.to_id == local_id for e in calls), "local definition was not preferred" + assert all(e.to_id != foreign_id for e in calls), "resolved to the other file's symbol" + + +def test_global_table_does_not_disturb_same_file_resolution(pm, tmp_path): + """Passing an unrelated global table must not change same-file behaviour.""" + tree, nodes = _parse(pm, tmp_path / "solo.py", SHADOWING_SRC) + + without = _calls( + pm.extract_edges(tree, nodes, str(tmp_path / "solo.py"), "python", source=SHADOWING_SRC) + ) + with_unrelated = _calls( + pm.extract_edges( + tree, + nodes, + str(tmp_path / "solo.py"), + "python", + source=SHADOWING_SRC, + global_symbols={"unrelated": "0" * 24}, + ) + ) + assert {(e.from_id, e.to_id) for e in without} == {(e.from_id, e.to_id) for e in with_unrelated} From e3557d58af2c736c1b2a3eb386271a0afd048fbf Mon Sep 17 00:00:00 2001 From: Rohit Behera <126186063+r0h1tb@users.noreply.github.com> Date: Sun, 2 Aug 2026 02:33:05 +0530 Subject: [PATCH 2/2] fix(api): bind call_kinds on the reference and impact queries UAT against a live index found 'ast-rag refs' and 'ast-rag symbol-impact' failing with: Neo.ClientError.Statement.ParameterMissing Expected parameter(s): call_kinds Regression from the call-traversal rewrite: two queries were changed to filter on $call_kinds, but their session.run() calls were never given the parameter. Neo4j only reports this at execution time, so nothing caught it -- the unit tests never reach these branches without a populated graph. Binds the parameter at both sites. After the fix, 'refs ParserManager' returns its references and 'symbol-impact ParserManager' reports 34 references and 42 callers. Adds a static checker over the API and repository layers: for every session.run(, **kwargs) it resolves back to its query text and asserts each $parameter is bound. Reverting the fix makes it fail, which is the property the first version of this test lacked -- a runtime test could not reach the broken branch and passed either way. --- ast_rag/api/ast_rag_api.py | 10 ++- tests/test_query_parameters_bound.py | 94 ++++++++++++++++++++++++++++ 2 files changed, 102 insertions(+), 2 deletions(-) create mode 100644 tests/test_query_parameters_bound.py diff --git a/ast_rag/api/ast_rag_api.py b/ast_rag/api/ast_rag_api.py index c9a4178..2dc3a2c 100644 --- a/ast_rag/api/ast_rag_api.py +++ b/ast_rag/api/ast_rag_api.py @@ -916,7 +916,7 @@ def _count_usages_of_node(self, node_id: str) -> int: """ total = 0 with self._driver.session() as session: - result = session.run(count_cypher, node_id=node_id) + result = session.run(count_cypher, node_id=node_id, call_kinds=CALL_EDGE_KINDS) record = result.single() total += record["count"] if record else 0 @@ -961,7 +961,13 @@ def _find_usages_of_node_paginated( SKIP $offset LIMIT $limit """ with self._driver.session() as session: - for record in session.run(calls_cypher, node_id=node_id, offset=offset, limit=limit): + for record in session.run( + calls_cypher, + node_id=node_id, + offset=offset, + limit=limit, + call_kinds=CALL_EDGE_KINDS, + ): caller_data = dict(record["caller"]) edge_data = dict(record["r"]) references.append( diff --git a/tests/test_query_parameters_bound.py b/tests/test_query_parameters_bound.py new file mode 100644 index 0000000..cb187c6 --- /dev/null +++ b/tests/test_query_parameters_bound.py @@ -0,0 +1,94 @@ +"""Every Cypher parameter a query references must actually be bound. + +Neo4j only reports an unbound parameter at execution time, as +``Neo.ClientError.Statement.ParameterMissing``. A query rewritten to use a new +parameter therefore keeps compiling, keeps passing any test that does not reach +that exact code path against a live database, and fails only in front of a user. + +That happened: the call traversals were changed to filter on ``$call_kinds``, +two ``session.run`` sites were not given the parameter, and ``ast-rag refs`` and +``ast-rag symbol-impact`` broke. + +Runtime tests are a poor fit -- reaching every branch needs a populated graph. +This checks it statically instead: for each ``session.run(, **kwargs)`` in +the API layer, resolve ```` back to its query string and assert every +``$parameter`` in that string appears in the keyword arguments. +""" + +from __future__ import annotations + +import ast +import re +from pathlib import Path + +import pytest + +PARAM = re.compile(r"\$([a-zA-Z_][a-zA-Z0-9_]*)") + +SOURCES = [ + Path(__file__).parent.parent / "ast_rag" / "api" / "ast_rag_api.py", + Path(__file__).parent.parent / "ast_rag" / "repositories" / "queries.py", +] + + +def _string_constants(tree: ast.AST) -> dict[str, str]: + """Map local variable names to the string literals assigned to them.""" + literals: dict[str, str] = {} + for node in ast.walk(tree): + if not isinstance(node, ast.Assign) or len(node.targets) != 1: + continue + target = node.targets[0] + if not isinstance(target, ast.Name): + continue + value = node.value + if isinstance(value, ast.Constant) and isinstance(value.value, str): + literals[target.id] = value.value + elif isinstance(value, ast.JoinedStr): + # f-string: keep the literal parts, which is where $params live + literals[target.id] = "".join( + p.value for p in value.values if isinstance(p, ast.Constant) + ) + return literals + + +def _run_calls(tree: ast.AST): + """Yield (query_text, bound_param_names, lineno) for each session.run(...).""" + literals = _string_constants(tree) + for node in ast.walk(tree): + if not isinstance(node, ast.Call): + continue + fn = node.func + if not (isinstance(fn, ast.Attribute) and fn.attr == "run"): + continue + if not node.args: + continue + first = node.args[0] + if isinstance(first, ast.Constant) and isinstance(first.value, str): + query = first.value + elif isinstance(first, ast.Name) and first.id in literals: + query = literals[first.id] + elif isinstance(first, ast.JoinedStr): + query = "".join(p.value for p in first.values if isinstance(p, ast.Constant)) + else: + continue # cannot resolve statically; skipped rather than guessed at + bound = {kw.arg for kw in node.keywords if kw.arg} + yield query, bound, node.lineno + + +@pytest.mark.parametrize("source", SOURCES, ids=lambda p: p.name) +def test_every_cypher_parameter_is_bound(source: Path): + tree = ast.parse(source.read_text()) + problems = [] + for query, bound, lineno in _run_calls(tree): + for name in sorted(set(PARAM.findall(query))): + if name not in bound: + problems.append(f"{source.name}:{lineno} uses ${name} but does not bind it") + assert not problems, "unbound Cypher parameters:\n " + "\n ".join(problems) + + +def test_checker_resolves_real_queries(): + """Guard the checker itself: it must actually be finding queries.""" + tree = ast.parse(SOURCES[0].read_text()) + calls = list(_run_calls(tree)) + assert len(calls) >= 5, f"static checker only resolved {len(calls)} queries; it may be broken" + assert any("$node_id" in q for q, _, _ in calls), "expected to see $node_id in some query"