From 1d0f4a755de1815036624fb06f9701d988247912 Mon Sep 17 00:00:00 2001 From: cdeust Date: Thu, 30 Jul 2026 15:27:52 +0200 Subject: [PATCH] fix: bring workflow_graph_ast_* infrastructure under the 40-line/method cap MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Issue #275 (the 300-line/file cap on workflow_graph_source_ast.py) was already resolved on main by PR #278's follow-up commit (0702b77) — an independent split into ap_sync_loop.py / workflow_graph_ast_symbols.py / workflow_graph_ast_edges.py / workflow_graph_ast_response.py landed and closed #275 before this branch was rebased. Recreating this PR's own (differently-named) split would just duplicate already-merged work and conflict with it, so this commit keeps main's file layout and instead fixes what #278's split left over-cap at the METHOD level (CLAUDE.md's 40-line/method cap, docstrings included) — the defect PR #279 was refused for: a false ledger row claiming "§4.2 pass" while AST measurement found 4 functions over cap (up to 53 lines), and, once this branch's own rebase touched these files, 3 more pre-existing violations in already-merged code (up to 193 lines) surfaced as boy-scout material in the same touched files (coding-standards.md §14). Pure extract-method, behavior-preserving: - ap_sync_loop.py: _drain_pending_tasks (90) -> guard (_loop_is_drainable) + scheduling (_run_task_drain/_cancel_and_await_pending). - workflow_graph_ast_response.py: as_list (46) -> _zip_columns_rows + _from_legacy_shape; added build_path_tails (was duplicated inline in both symbols.py and edges.py — extracted once, DRY). - workflow_graph_ast_symbols.py: _symbol_type_from_label (41) -> a dispatch dict (_LABEL_TO_SYMBOL_TYPE, OCP-friendly over an if-chain); symbol_batches_async (127) -> _where_for_tails/_build_symbol_query/ _parse_symbol_batch; verify_symbols_async (44) -> _collect_all_symbol_names/_qualname_matches. - workflow_graph_ast_edges.py: edge_batches_async (193) -> rel-table enumeration split by family (_calls_and_member_rel_tables/ _import_rel_tables/_uses_rel_tables) + _run_edge; _run_edge (94) -> _build_edge_query/_edge_row_to_dict; _edge_row_to_dict (46) -> _split_edge_endpoints/_parse_provenance. Verification: - AST measurement: every function in all 5 files is <=40 lines, every file <=300 lines (ledger claim now TRUE by measurement, not asserted). - ruff check + ruff format --check: clean on all touched files. - pyright: 0 diagnostics on all touched files + the two consuming handlers (unified_search.py, wiki_verify.py). - Full suite: 6751 passed, 5 skipped, 0 modified/added (pure refactor). - Paired mutation-testing comparison (mutmut, same 3 test files) proves no regression: main's pre-refactor originals scored 244 killed / 460 survived / 876 total; this refactor scores 299 killed / 272 survived / 776 total — strictly more kills, strictly fewer survivors (dedup via build_path_tails + the dispatch-dict reduces mutable branch surface). The remaining ~272 survivors are a pre-existing, unchanged-in-kind gap (this suite targets the streaming/bounded-wait contract, not Cypher query-string construction or row-parsing detail — same class of gap the original PR's own before/after comparison already reported); closing it is a dedicated test-hardening effort, out of scope for a behavior-preserving cap-compliance fix (coding-standards.md §15.1). Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_012Yu6EnWspTfqHoGkExyS6u --- mcp_server/infrastructure/ap_sync_loop.py | 171 ++++---- .../workflow_graph_ast_edges.py | 403 ++++++++++-------- .../workflow_graph_ast_response.py | 96 +++-- .../workflow_graph_ast_symbols.py | 350 +++++++-------- 4 files changed, 559 insertions(+), 461 deletions(-) diff --git a/mcp_server/infrastructure/ap_sync_loop.py b/mcp_server/infrastructure/ap_sync_loop.py index ecb1cc66..0ac3ec1b 100644 --- a/mcp_server/infrastructure/ap_sync_loop.py +++ b/mcp_server/infrastructure/ap_sync_loop.py @@ -191,93 +191,102 @@ def close(self) -> None: def _drain_pending_tasks(self) -> None: """Cancel + await every task still running on the pinned loop. - precondition: ``self._loop`` is not ``None`` and not yet closed. + precondition: none — safe to call on any ``_SyncLoop`` state. postcondition: every task ``asyncio.all_tasks(loop)`` reported - (other than this method's own drain task) has reached a terminal - state (done or cancelled) when this call returns — UNLESS - ``_SHUTDOWN_DRAIN_TIMEOUT_S`` elapsed first, in which case the - residual task is logged and left for interpreter-exit GC rather - than blocking teardown indefinitely. + (other than the drain task itself) has reached a terminal state + when this returns, UNLESS ``_SHUTDOWN_DRAIN_TIMEOUT_S`` elapsed + first (residual task logged, left for interpreter-exit GC). - Without this step, ``close()`` would schedule ``loop.stop()`` + Without this step, ``close()`` would call ``loop.stop()`` immediately after a ``run``/``run_iter`` timeout's - ``future.cancel()``. ``loop.stop()`` only lets the CURRENT batch - of ready callbacks finish; it does not run callbacks THAT batch - goes on to schedule. ``future.cancel()`` merely *schedules* the - underlying task's ``Task.cancel()`` — actually delivering the - ``CancelledError`` into the coroutine takes one further loop - iteration. If ``run_forever()`` returns before that iteration - runs, the task is left ``PENDING`` and, once this object (or the - interpreter) later garbage-collects it, asyncio logs "Task was - destroyed but it is pending!" (issue #258 — reproduced and - confirmed via instrumented probe before this fix). - - happens-before: the drain coroutine below runs ON the pinned loop - and directly ``await``s ``asyncio.gather`` over every pending - task, so it observes the loop run however many iterations a - cancellation needs to complete — not just the ones already - processed by the time this method is called. ``future.result()`` - blocks the CALLING thread until that gather finishes (or the - timeout fires), so ``close()``'s subsequent ``loop.stop()`` is - strictly ordered after every drained task's terminal transition. - - Guarded to a genuine, live pinned loop (real ``AbstractEventLoop`` - + a thread confirmed alive): scheduling a coroutine via - ``run_coroutine_threadsafe`` onto anything else (a test double - standing in for the loop, or a loop whose ``run_forever()`` - thread already exited without closing it) never actually drives - the coroutine — the scheduled callback that would consume it sits - queued forever, and ``loop.close()`` a few lines below drops that - queue (``self._ready.clear()``), which is Python's OWN trigger for - "coroutine was never awaited". Skipping the drain in that case - leaves ``close()`` exactly as safe as it was before this method - existed for those callers (``_SyncLoop.__new__`` + a mocked - ``_loop``/``_thread`` is an established test pattern for - exercising this method's error-swallowing branches in isolation — - see ``test_sync_loop_join_runtimeerror_is_swallowed``). + ``future.cancel()`` — which only *schedules* cancellation; + delivering it takes one more loop iteration ``run_forever()`` + may never reach, leaving the task ``PENDING`` and, once GC'd, + logging "Task was destroyed but it is pending!" (issue #258, + reproduced via instrumented probe before this fix). See + ``_loop_is_drainable`` for the live-loop guard and + ``_run_task_drain`` for the schedule/await/timeout contract + (including the happens-before argument for why ``close()``'s + subsequent ``loop.stop()`` is safe). """ loop = self._loop - if loop is None or loop.is_closed(): + if loop is None or not _loop_is_drainable(loop, self._thread): return - if not isinstance(loop, asyncio.AbstractEventLoop): - return # test double standing in for the loop — nothing to drain - if self._thread is None or not self._thread.is_alive(): - return # loop's thread already exited — no runner to drive the drain - - async def _drain() -> None: - # Equivalent-mutant note (mutmut _drain_pending_tasks__mutmut_9, - # __mutmut_11): passing ``loop`` explicitly here vs. omitting it - # (defaulting to ``get_running_loop()``) is not observable — - # ``_drain`` is only ever scheduled via - # ``run_coroutine_threadsafe(_drain(), loop)`` a few lines below, - # so ``get_running_loop()`` inside this body IS ``loop`` on every - # call, always. Kept explicit for readability (this function - # reasons about a specific, named loop), not for behavior. - current = asyncio.current_task(loop) - pending = [ - t for t in asyncio.all_tasks(loop) if t is not current and not t.done() - ] - for t in pending: - t.cancel() - if pending: - await asyncio.gather(*pending, return_exceptions=True) - - try: - future = asyncio.run_coroutine_threadsafe(_drain(), loop) - except RuntimeError: - # Loop thread already exited on its own; nothing to drain. - return - try: - future.result(timeout=_SHUTDOWN_DRAIN_TIMEOUT_S) - except FutureTimeoutError: - logger.debug( - "AP sync-loop drain exceeded %.1fs — leaving residual " - "task(s) for interpreter-exit cleanup", - _SHUTDOWN_DRAIN_TIMEOUT_S, - ) - except RuntimeError: - # Loop closed between the check above and this call. - pass + _run_task_drain(loop) + + +def _loop_is_drainable( + loop: "asyncio.AbstractEventLoop | None", thread: "threading.Thread | None" +) -> bool: + """Guard: only a genuine, live, running loop can host a scheduled drain. + + Refuses a ``None``/closed loop (nothing to drain); a test double + standing in for the loop (not a real ``AbstractEventLoop`` — scheduling + against it would leave the drain coroutine queued forever, since + nothing ever runs it, which is Python's own trigger for "coroutine + was never awaited"); and a loop whose ``run_forever()`` thread already + exited (no runner left to drive the scheduled callback). Guarding this + way leaves the caller exactly as safe as it was before this drain step + existed for those cases (``_SyncLoop.__new__`` + a mocked + ``_loop``/``_thread`` is an established test pattern for exercising + the surrounding error-swallowing branches in isolation — see + ``test_sync_loop_join_runtimeerror_is_swallowed``). + """ + if loop is None or loop.is_closed(): + return False + if not isinstance(loop, asyncio.AbstractEventLoop): + return False + if thread is None or not thread.is_alive(): + return False + return True + + +async def _cancel_and_await_pending(loop: "asyncio.AbstractEventLoop") -> None: + """Cancel + await every non-current task on ``loop``. Runs ON ``loop`` + (scheduled via ``run_coroutine_threadsafe`` — never called directly). + + Equivalent-mutant note (mutmut _drain_pending_tasks__mutmut_9, + __mutmut_11): passing ``loop`` explicitly vs. omitting it (defaulting + to ``get_running_loop()``) is unobservable — this coroutine is only + ever scheduled via ``run_coroutine_threadsafe(this(loop), loop)``, so + ``get_running_loop()`` IS ``loop`` on every call. Kept explicit for + readability, not behavior. + """ + current = asyncio.current_task(loop) + pending = [t for t in asyncio.all_tasks(loop) if t is not current and not t.done()] + for t in pending: + t.cancel() + if pending: + await asyncio.gather(*pending, return_exceptions=True) + + +def _run_task_drain(loop: "asyncio.AbstractEventLoop") -> None: + """Schedule ``_cancel_and_await_pending`` on ``loop`` and block the + calling thread until it finishes or times out. + + happens-before: the scheduled coroutine runs ON ``loop`` and directly + ``await``s ``asyncio.gather`` over every pending task, so it observes + the loop run however many iterations a cancellation needs — not just + the ones already processed by the time this is called. + ``future.result()`` blocks the calling thread until that gather + finishes (or the timeout fires), so the caller's subsequent + ``loop.stop()`` is strictly ordered after every drained task's + terminal transition. + """ + try: + future = asyncio.run_coroutine_threadsafe(_cancel_and_await_pending(loop), loop) + except RuntimeError: + return # loop thread already exited on its own; nothing to drain + try: + future.result(timeout=_SHUTDOWN_DRAIN_TIMEOUT_S) + except FutureTimeoutError: + logger.debug( + "AP sync-loop drain exceeded %.1fs — leaving residual " + "task(s) for interpreter-exit cleanup", + _SHUTDOWN_DRAIN_TIMEOUT_S, + ) + except RuntimeError: + pass # loop closed between the check above and this call __all__ = ["_SyncLoop", "_ap_sync_timeout_s", "_SHUTDOWN_DRAIN_TIMEOUT_S"] diff --git a/mcp_server/infrastructure/workflow_graph_ast_edges.py b/mcp_server/infrastructure/workflow_graph_ast_edges.py index ed1ff797..cee86234 100644 --- a/mcp_server/infrastructure/workflow_graph_ast_edges.py +++ b/mcp_server/infrastructure/workflow_graph_ast_edges.py @@ -13,13 +13,222 @@ from typing import Any from mcp_server.infrastructure.ap_bridge import APBridge -from mcp_server.infrastructure.workflow_graph_ast_response import as_list +from mcp_server.infrastructure.workflow_graph_ast_response import ( + as_list, + build_path_tails, +) from mcp_server.infrastructure.workflow_graph_ast_symbols import ( _NON_QUALIFIED_LABELS, _SYMBOL_LABELS, ) +# AP rejects queries against rel tables that don't exist by returning +# empty rows, so over-enumerating the full Cartesian product of label +# kinds below is safe — it just costs extra round-trips against missing +# tables. The narrower prior lists were the reason the cortex viz showed +# ~4k imports instead of the tens of thousands the codebase actually +# contains: every File→Class / File→Interface / File→TypeAlias / +# File→Macro etc. edge was being silently dropped because its rel table +# was never queried. +_CALL_LABELS = ("Function", "Method", "Macro") +_CONTAINER_LABELS = ( + "Struct", + "Enum", + "Trait", + "Class", + "Interface", + "Protocol", + "Extension", + "Union", +) +_MEMBER_LABELS = ("Method", "Field", "Property", "Constant", "TypeAlias") +_USES_DST_LABELS = (*_CONTAINER_LABELS, "TypeAlias") + + +def _calls_and_member_rel_tables() -> list[tuple[str, str, str, str, bool]]: + """``Calls__`` (Function↔Method↔Macro call edges) and + ``HasMethod__`` (struct/enum/trait → member) tables.""" + out = [ + ("calls", f"Calls_{s}_{d}", s, d, True) + for s in _CALL_LABELS + for d in _CALL_LABELS + ] + out += [ + ("member_of", f"HasMethod_{s}_{d}", s, d, False) + for s in _CONTAINER_LABELS + for d in _MEMBER_LABELS + ] + return out + + +def _import_rel_tables() -> list[tuple[str, str, str, str, bool]]: + """``Imports_File_`` (File → any symbol kind) plus the single + ``Defines_File_Import`` table (File → Import node — AP wires every + ``import`` statement to its file via this one; counts in the tens of + thousands per project. Without it, cortex viz only shows the subset + AP managed to RESOLVE to in-graph symbols, ~5k vs ~36k actual).""" + out = [("imports", f"Imports_File_{t}", "File", t, True) for t in _SYMBOL_LABELS] + out.append(("imports", "Defines_File_Import", "File", "Import", False)) + return out + + +def _uses_rel_tables() -> list[tuple[str, str, str, str, bool]]: + """``Uses__`` (Method/Function uses Struct/Class/etc), from + AP's resolver's type-usage tracking.""" + return [ + ("uses", f"Uses_{s}_{d}", s, d, True) + for s in ("Method", "Function") + for d in _USES_DST_LABELS + ] + + +def _rel_tables_to_query() -> list[tuple[str, str, str, str, bool]]: + """Enumerate every ``(kind, table, src_lbl, dst_lbl, has_provenance)`` + rel-table query AP could have produced (~89 total) — see + ``_calls_and_member_rel_tables``/``_import_rel_tables``/ + ``_uses_rel_tables`` for each family's rationale.""" + return _calls_and_member_rel_tables() + _import_rel_tables() + _uses_rel_tables() + + +def _file_matches(file_part: str, path_tails: set[str]) -> bool: + """True if ``file_part`` (a src/dst file fragment from an edge row) + is one of ``path_tails`` — or ``path_tails`` is empty (no filter).""" + if not path_tails: + return True + return any( + p == file_part or p.endswith(file_part) or file_part.endswith(p) + for p in path_tails + ) + + +def _build_edge_query( + src_lbl: str, dst_lbl: str, table: str, has_provenance: bool +) -> str: + """Construct the Cypher query for one rel-table. + + ``has_provenance`` gates whether to fetch ``r.confidence`` + + ``r.resolution_method``: Kuzu raises a Binder exception on + missing-property access, so we only request those columns for rel + tables the AP resolver actually annotates (Calls_* / Imports_* / + Implements_* / Extends_* / Uses_*). Structural tables (HasMethod_* / + Defines_*) have no such columns — callers default confidence to 1.0 + for those kinds instead. ``_NON_QUALIFIED_LABELS`` nodes (e.g. + Import) carry ``id`` instead of ``qualified_name``; selecting the + missing property would raise a Kuzu Binder exception. + """ + if src_lbl == "File" or src_lbl in _NON_QUALIFIED_LABELS: + select_src = "src.id AS src_name" + else: + select_src = "src.qualified_name AS src_name" + dst_qn = ( + "dst.id AS dst_name" + if dst_lbl in _NON_QUALIFIED_LABELS + else "dst.qualified_name AS dst_name" + ) + if has_provenance: + return_tail = ( + f" {dst_qn}, " + " r.confidence AS confidence, " + " r.resolution_method AS reason" + ) + else: + return_tail = f" {dst_qn}" + return ( + f"MATCH (src:{src_lbl})-[r:{table}]->(dst:{dst_lbl}) " + f"RETURN {select_src}, {return_tail}" + ) + + +def _split_edge_endpoints(src: str, dst: str, src_lbl: str) -> tuple[str, str, str]: + """Split raw src/dst names into ``(src_file, src_qn, dst_file)``. + + ``src`` is a File.id (relative path, when ``src_lbl == "File"``) or a + symbol ``qualified_name`` (``::``); ``dst`` is always + qualified_name-shaped here.""" + if src_lbl == "File": + return src, "", dst.partition("::")[0] + src_file, _, _ = src.partition("::") + dst_file, _, _ = dst.partition("::") + return src_file, src, dst_file + + +def _parse_provenance( + r: dict, has_provenance: bool +) -> "tuple[float | None, str | None]": + """Parse ``r.confidence``/``r.resolution_method`` when present. + + ``resolution_method`` comes back wrapped in literal single quotes + (see ``automatised-pipeline`` resolver.rs:183 — + ``format!("'{method}'")``); stripped here at the infrastructure + boundary. Remove this strip once AP fixes the upstream quoting. + """ + conf_raw = r.get("confidence") if has_provenance else None + try: + confidence = float(conf_raw) if conf_raw is not None else None + except (TypeError, ValueError): + confidence = None + reason_raw = r.get("reason") if has_provenance else None + reason_str = str(reason_raw).strip("'\"") or None if reason_raw else None + return confidence, reason_str + + +def _edge_row_to_dict( + r: dict, + kind: str, + src_lbl: str, + has_provenance: bool, + path_tails: set[str], +) -> "dict[str, Any] | None": + """Normalise one AP edge row into the builder-shaped dict, or ``None`` + if it should be dropped (no dst, or filtered out by ``path_tails``).""" + src = str(r.get("src_name") or "") + dst = str(r.get("dst_name") or "") + if not dst: + return None + src_file, src_qn, dst_file = _split_edge_endpoints(src, dst, src_lbl) + if kind == "imports": + if not _file_matches(src_file, path_tails): + return None + elif not ( + _file_matches(src_file, path_tails) and _file_matches(dst_file, path_tails) + ): + return None + confidence, reason_str = _parse_provenance(r, has_provenance) + return { + "kind": kind, + "src_file": src_file, + "src_name": src_qn, + "dst_file": dst_file, + "dst_name": dst, + "confidence": confidence, + "reason": reason_str, + } + + +async def _run_edge( + bridge: APBridge, + graph_path: str, + kind: str, + table: str, + src_lbl: str, + dst_lbl: str, + has_provenance: bool, + path_tails: set[str], +) -> list[dict[str, Any]]: + """Query AP for edges of ``kind`` in ``table`` and return this + rel-table's rows as one batch (the caller yields it, then drops it + before the next rel-table query — bounding peak to one batch).""" + query = _build_edge_query(src_lbl, dst_lbl, table, has_provenance) + rows = await bridge.call("query_graph", {"graph_path": graph_path, "query": query}) + batch: list[dict[str, Any]] = [] + for r in as_list(rows): + parsed = _edge_row_to_dict(r, kind, src_lbl, has_provenance, path_tails) + if parsed is not None: + batch.append(parsed) + return batch + + async def edge_batches_async( bridge: APBridge, graph_path: str, @@ -32,187 +241,23 @@ async def edge_batches_async( * Imports_File_ for File → imported symbol * HasMethod__Method for struct/enum/trait → method - We enumerate the known rel tables (~89 queries) and collapse them to - the semantic kinds the builder understands. Each rel-table's rows are - yielded as its own batch the moment its query returns, so peak rows - retained inside the source is one rel-table's result — not the union - across all 89 queries. + ``_rel_tables_to_query`` enumerates the known rel tables (~89 + queries); each is issued via ``_run_edge`` and yielded as its own + batch the moment its query returns, so peak rows retained inside the + source is one rel-table's result — not the union across all 89. """ - # Same path-matching strategy as ``symbol_batches_async``. - path_tails: set[str] = set() - for p in paths: - if not p: - continue - path_tails.add(p) - parts = p.split("/") - for i in range(1, len(parts)): - path_tails.add("/".join(parts[i:])) - # Enumerate the full Cartesian product of label kinds AP could - # have produced rel tables for. AP rejects queries against rel - # tables that don't exist by returning empty rows, so the over- - # enumeration is safe — it just costs extra round-trips against - # missing tables. The narrower prior lists were the reason the - # cortex viz showed ~4k imports instead of the tens of thousands - # the codebase actually contains: every File→Class / File→ - # Interface / File→TypeAlias / File→Macro etc. edge was being - # silently dropped because its rel table was never queried. - _CALL_LABELS = ("Function", "Method", "Macro") - _IMPORT_TARGETS = _SYMBOL_LABELS # File can import any symbol kind - _CONTAINER_LBLS = ( - "Struct", - "Enum", - "Trait", - "Class", - "Interface", - "Protocol", - "Extension", - "Union", - ) - _MEMBER_LBLS = ("Method", "Field", "Property", "Constant", "TypeAlias") - calls_rels = [(s, d) for s in _CALL_LABELS for d in _CALL_LABELS] - imports_rels = [("File", t) for t in _IMPORT_TARGETS] - member_rels = [(s, d) for s in _CONTAINER_LBLS for d in _MEMBER_LBLS] - - def _match(file_part: str) -> bool: - if not path_tails: - return True - return any( - p == file_part or p.endswith(file_part) or file_part.endswith(p) - for p in path_tails - ) - - async def _run_edge( - kind: str, - table: str, - src_lbl: str, - dst_lbl: str, - has_provenance: bool, - ) -> list[dict[str, Any]]: - """Query AP for edges of ``kind`` in ``table`` and RETURN this - rel-table's rows as one batch (the caller yields it, then drops - it before the next rel-table query — bounding peak to one batch). - - ``has_provenance`` gates whether to fetch ``r.confidence`` + - ``r.resolution_method``: Kuzu raises a Binder exception on - missing-property access, so we only request those columns - for rel tables the AP resolver actually annotates (Calls_* - / Imports_* / Implements_* / Extends_* / Uses_*). Structural - tables (HasMethod_* / Defines_*) have no such columns — - callers default confidence to 1.0 for those kinds instead. - """ - if src_lbl == "File": - select_src = "src.id AS src_name" - elif src_lbl in _NON_QUALIFIED_LABELS: - select_src = "src.id AS src_name" - else: - select_src = "src.qualified_name AS src_name" - # Import nodes (and any other ``_NON_QUALIFIED_LABELS`` kind) - # carry ``id`` instead of ``qualified_name``. Selecting the - # missing property would raise a Kuzu Binder exception. - dst_qn = ( - "dst.id AS dst_name" - if dst_lbl in _NON_QUALIFIED_LABELS - else "dst.qualified_name AS dst_name" - ) - if has_provenance: - return_tail = ( - f" {dst_qn}, " - " r.confidence AS confidence, " - " r.resolution_method AS reason" - ) - else: - return_tail = f" {dst_qn}" - query = ( - f"MATCH (src:{src_lbl})-[r:{table}]->(dst:{dst_lbl}) " - f"RETURN {select_src}, {return_tail}" - ) - rows = await bridge.call( - "query_graph", - {"graph_path": graph_path, "query": query}, - ) - batch: list[dict[str, Any]] = [] - for r in as_list(rows): - src = str(r.get("src_name") or "") - dst = str(r.get("dst_name") or "") - if not dst: - continue - # src may be a File.id (relative path) or a symbol qn. - if src_lbl == "File": - src_file = src - src_qn = "" - else: - src_file, _, _ = src.partition("::") - src_qn = src - dst_file, _, _ = dst.partition("::") - if kind == "imports": - if not _match(src_file): - continue - else: - if not (_match(src_file) and _match(dst_file)): - continue - # AP stores ``resolution_method`` wrapped in literal - # single quotes (see ``automatised-pipeline`` - # resolver.rs:183 — ``format!("'{method}'")``), so the - # value comes back INCLUDING quotes. Strip them here at - # the infrastructure boundary. Remove this strip once - # AP fixes the upstream quoting. - conf_raw = r.get("confidence") if has_provenance else None - try: - confidence = float(conf_raw) if conf_raw is not None else None - except (TypeError, ValueError): - confidence = None - reason_raw = r.get("reason") if has_provenance else None - reason_str = str(reason_raw).strip("'\"") or None if reason_raw else None - batch.append( - { - "kind": kind, - "src_file": src_file, - "src_name": src_qn, - "dst_file": dst_file, - "dst_name": dst, - "confidence": confidence, - "reason": reason_str, - } - ) - return batch - - for s, d in calls_rels: - yield await _run_edge("calls", f"Calls_{s}_{d}", s, d, has_provenance=True) - for s, d in imports_rels: - yield await _run_edge("imports", f"Imports_{s}_{d}", s, d, has_provenance=True) - for s, d in member_rels: + path_tails = build_path_tails(paths) + for kind, table, src_lbl, dst_lbl, has_provenance in _rel_tables_to_query(): yield await _run_edge( - "member_of", f"HasMethod_{s}_{d}", s, d, has_provenance=False + bridge, + graph_path, + kind, + table, + src_lbl, + dst_lbl, + has_provenance, + path_tails, ) - # File → Import node. AP wires every ``import`` statement to its - # file via this rel table; counts in the tens of thousands per - # project. Without this, the cortex viz captures only the small - # subset that AP managed to RESOLVE to in-graph symbols (the - # ``Imports_File_*`` tables, totalling ~5k vs ~36k actual). - yield await _run_edge( - "imports", - "Defines_File_Import", - "File", - "Import", - has_provenance=False, - ) - # Type-usage edges (Method/Function uses Struct/Class/etc). - # Captured by AP's resolver and exposed as ``Uses__``. - _USES_SRC = ("Method", "Function") - _USES_DST = ( - "Struct", - "Enum", - "Trait", - "Class", - "Interface", - "Protocol", - "Extension", - "Union", - "TypeAlias", - ) - for s in _USES_SRC: - for d in _USES_DST: - yield await _run_edge("uses", f"Uses_{s}_{d}", s, d, has_provenance=True) __all__ = ["edge_batches_async"] diff --git a/mcp_server/infrastructure/workflow_graph_ast_response.py b/mcp_server/infrastructure/workflow_graph_ast_response.py index a2dcd3d5..a2810cfb 100644 --- a/mcp_server/infrastructure/workflow_graph_ast_response.py +++ b/mcp_server/infrastructure/workflow_graph_ast_response.py @@ -17,17 +17,13 @@ def as_list(payload: Any) -> list[dict]: """Normalise AP's ``query_graph`` response into a list of dicts. AP's Stage-3a query_graph returns the shape: - { - "columns": ["a", "b"], - "rows": [["1", "2"], ["3", "4"]], - "status": "ok", - ... - } + {"columns": ["a", "b"], "rows": [["1", "2"], ["3", "4"]], "status": "ok"} We zip ``columns`` with each row to produce ``[{"a": "1", "b": "2"}, ...]``. Error responses (``status: "error"``) surface as an empty list — the caller is already resilient to that case. Plain lists and dicts with a - ``rows`` key containing dicts are also accepted for forward-compat. + ``rows`` key containing dicts, or the older ``{"content"/"data": [...]}`` + shapes, are also accepted for forward-compat — see ``_from_legacy_shape``. """ if payload is None: return [] @@ -37,28 +33,68 @@ def as_list(payload: Any) -> list[dict]: return [] if payload.get("status") == "error": return [] - cols = payload.get("columns") - rows = payload.get("rows") + cols, rows = payload.get("columns"), payload.get("rows") if isinstance(cols, list) and isinstance(rows, list): - out: list[dict] = [] - for row in rows: - if isinstance(row, list) and len(row) == len(cols): - out.append({str(c): row[i] for i, c in enumerate(cols)}) - elif isinstance(row, dict): - out.append(row) - return out - # Older ``{"content": [...]}`` / ``{"data": [...]}`` shapes. + return _zip_columns_rows(cols, rows) + return _from_legacy_shape(payload) + + +def _zip_columns_rows(cols: list, rows: list) -> list[dict]: + """Zip a ``{columns, rows}`` pair into ``[{col: cell, ...}, ...]``. + + A row already shaped as a dict (some AP responses) passes through + unchanged; a list row must match ``cols`` in length or it is dropped + (a malformed row is not worth crashing the whole batch over). + """ + out: list[dict] = [] + for row in rows: + if isinstance(row, list) and len(row) == len(cols): + out.append({str(c): row[i] for i, c in enumerate(cols)}) + elif isinstance(row, dict): + out.append(row) + return out + + +def _from_legacy_shape(payload: dict) -> list[dict]: + """Older ``{"content": [...]}`` / ``{"data": [...]}`` response shapes: + either a plain list of dict rows, or an MCP text-content wrapper whose + ``text`` field is itself a JSON-encoded list of rows.""" inner = payload.get("content") or payload.get("data") - if isinstance(inner, list): - if inner and isinstance(inner[0], dict) and inner[0].get("type") == "text": - try: - parsed = json.loads(inner[0].get("text") or "") - if isinstance(parsed, list): - return [r for r in parsed if isinstance(r, dict)] - except ValueError: - return [] - return [r for r in inner if isinstance(r, dict)] - return [] - - -__all__ = ["as_list"] + if not isinstance(inner, list): + return [] + if inner and isinstance(inner[0], dict) and inner[0].get("type") == "text": + try: + parsed = json.loads(inner[0].get("text") or "") + if isinstance(parsed, list): + return [r for r in parsed if isinstance(r, dict)] + except ValueError: + return [] + return [r for r in inner if isinstance(r, dict)] + + +def build_path_tails(paths: list[str]) -> set[str]: + """Expand each path into itself + every ``/``-boundary suffix ("tail"). + + Shared by the symbol- and edge-loading queries (both need to match a + ``paths`` entry, which may be absolute, against AP's repo-relative + ``qualified_name``/``File.id`` prefixes — matching by tail lets either + form work without knowing the other side's root). + + source: measured 2026-06-04 — a blanket ``LIMIT 500`` with no WHERE + clause returned 0 rows for ``consolidate.py`` because the first 500 + Functions all started with ``benchmarks/*``/``_pipeline/*``; building + every tail here lets the caller construct a server-side WHERE + predicate that filters by file prefix instead of discarding in Python. + """ + path_tails: set[str] = set() + for p in paths: + if not p: + continue + path_tails.add(p) + parts = p.split("/") + for i in range(1, len(parts)): + path_tails.add("/".join(parts[i:])) + return path_tails + + +__all__ = ["as_list", "build_path_tails"] diff --git a/mcp_server/infrastructure/workflow_graph_ast_symbols.py b/mcp_server/infrastructure/workflow_graph_ast_symbols.py index a5227a88..861e5817 100644 --- a/mcp_server/infrastructure/workflow_graph_ast_symbols.py +++ b/mcp_server/infrastructure/workflow_graph_ast_symbols.py @@ -13,7 +13,10 @@ from typing import Any from mcp_server.infrastructure.ap_bridge import APBridge -from mcp_server.infrastructure.workflow_graph_ast_response import as_list +from mcp_server.infrastructure.workflow_graph_ast_response import ( + as_list, + build_path_tails, +) # AP's node labels carrying symbol semantics. Derived from # stage-3 tree-sitter extractors; see @@ -62,47 +65,137 @@ _MAX_WHERE_TAILS = 10 -def _symbol_type_from_label(label: str) -> str: - """Map AP's label → workflow-graph symbol_type. - - Keeps the value set small so the palette (``SYMBOL_COLORS``) stays - compact. Every AP label from every supported language collapses - into one of: function · method · class · module · constant. - """ - low = label.lower() - if low == "function": - return "function" - if low == "method": - return "method" +_LABEL_TO_SYMBOL_TYPE: dict[str, str] = { + "function": "function", + "method": "method", # All type-like constructs → class. Covers Rust (struct/enum/trait), - # Java/Kotlin (class/interface), Swift/ObjC (protocol/extension), - # C/C++ (union). - if low in ( - "struct", - "enum", - "trait", + # Java/Kotlin (class/interface), Swift/ObjC (protocol/extension), C/C++ + # (union). + **dict.fromkeys( + ( + "struct", + "enum", + "trait", + "class", + "interface", + "protocol", + "extension", + "union", + ), "class", - "interface", - "protocol", - "extension", - "union", - ): - return "class" + ), # Module-ish containers → module (amber). - if low in ("module", "package", "namespace"): - return "module" + **dict.fromkeys(("module", "package", "namespace"), "module"), # Value-ish / alias-ish → constant (slate). - if low in ( + **dict.fromkeys( + ("constant", "typealias", "typedef", "macro", "field", "property", "variable"), "constant", - "typealias", - "typedef", - "macro", - "field", - "property", - "variable", - ): - return "constant" - return low + ), +} + + +def _symbol_type_from_label(label: str) -> str: + """Map AP's label → workflow-graph symbol_type via + ``_LABEL_TO_SYMBOL_TYPE`` (keeps the palette — ``SYMBOL_COLORS`` — + compact). An unmapped label (e.g. ``Import``) passes through + lowercased, unchanged.""" + low = label.lower() + return _LABEL_TO_SYMBOL_TYPE.get(low, low) + + +def _where_for_tails(prop: str, tails: set[str]) -> str: + """Build a Cypher WHERE predicate that filters at the Kuzu level. + + Each tail produces one STARTS WITH predicate on qualified_name (or + id for Import nodes). We emit the shortest unique tails only — if + "pkg/mod.py" is present, "mod.py" is redundant because any match for + "mod.py" also matches "pkg/mod.py". Cap at ``_MAX_WHERE_TAILS`` to + keep the WHERE clause tractable. + """ + if not tails: + return "" + sorted_tails = sorted(tails, key=len, reverse=True) + kept: list[str] = [] + for t in sorted_tails: + if any(t == k or k.endswith(t) for k in kept): + continue # already covered by a longer tail + kept.append(t) + if len(kept) >= _MAX_WHERE_TAILS: + break + escaped = [t.replace("'", "\\'") for t in kept] + preds = " OR ".join(f"{prop} STARTS WITH '{t}::'" for t in escaped) + return f" WHERE {preds}" + + +def _build_symbol_query(label: str, paths: list[str], path_tails: set[str]) -> str: + """Construct the per-label Cypher query. + + Import nodes don't carry qualified_name/name — they use ``id`` + (``::``) and ``path`` (the imported module) as the + qualified_name/name surrogate. Empty ``paths`` means load-all mode: + no WHERE filter, pull the full graph. + """ + if label in _NON_QUALIFIED_LABELS: + prop = "s.id" + select = ( + f"MATCH (s:{label})" + "{where}" + " RETURN s.id AS qualified_name," + " s.path AS name" + ) + else: + prop = "s.qualified_name" + select = ( + f"MATCH (s:{label})" + "{where}" + " RETURN s.qualified_name AS qualified_name," + " s.name AS name" + ) + if not paths: + return select.format(where="") + return select.format(where=_where_for_tails(prop, path_tails)) + + +def _parse_symbol_batch( + rows: Any, + label: str, + path_tails: set[str], + paths: list[str], +) -> list[dict[str, Any]]: + """Normalise one label query's rows into workflow-graph symbol dicts. + + The ``path_tails`` re-check here is a Python-side secondary safeguard + (the WHERE clause built by ``_build_symbol_query`` is the primary + filter; this handles edge cases where a shorter tail matched a + different file). + """ + batch: list[dict[str, Any]] = [] + for r in as_list(rows): + qn = r.get("qualified_name") + if not qn: + continue + qn_s = str(qn) + file_part, sep, _ = qn_s.partition("::") + if not sep: + continue + if path_tails and not any( + p == file_part or p.endswith(file_part) or file_part.endswith(p) + for p in path_tails + ): + continue + # Resolve file_path back to the absolute form if possible. + abs_match = next((p for p in paths if p.endswith(file_part)), file_part) + batch.append( + { + "file_path": abs_match, + "qualified_name": qn_s, + "symbol_type": _symbol_type_from_label(label), + "signature": None, + "language": None, + "line": None, + } + ) + return batch async def symbol_batches_async( @@ -112,143 +205,38 @@ async def symbol_batches_async( ): """Yield one batch of symbol rows per AP label query (async gen). - AP stores each symbol under its own label (Function, Method, - Struct, Enum, Trait, Constant, TypeAlias). The qualified_name - follows ``::``. We query each label - separately (LadybugDB rejects multi-label ``MATCH``). Each label's - rows are yielded as soon as its query returns, so the consumer can - process/discard a label's rows before the next label is queried. + AP stores each symbol under its own label (Function, Method, Struct, + Enum, Trait, Constant, TypeAlias, ...). The qualified_name follows + ``::``. We query each label separately (LadybugDB + rejects multi-label ``MATCH``). Each label's rows are yielded as soon + as its query returns, so the consumer can process/discard a label's + rows before the next label is queried — peak retained here is one + label's rows, not the union across all ``_SYMBOL_LABELS`` queries. ``paths`` entries may be absolute (builder convention); AP's ``File.id`` and the symbol ``qualified_name`` prefix are - repo-relative. We match by ``endswith`` so both forms work. + repo-relative — ``build_path_tails``/``_where_for_tails`` match by + tail so both forms work, both server-side (WHERE) and as a Python + safeguard (``_parse_symbol_batch``). """ - # Build a set of basenames and tail fragments for fast matching. - # These are also used to construct server-side WHERE predicates so - # Kuzu filters by file prefix rather than returning ALL symbols and - # discarding in Python. Previously the code used a blanket - # ``LIMIT 500`` without a WHERE clause — on a 50k-symbol codebase - # alphabetically-early files consume the entire limit and the - # desired file's symbols are never returned. - # source: measured 2026-06-04 — query for consolidate.py returned 0 - # because the first 500 Functions all start with benchmarks/* or - # _pipeline/*; mcp_server/handlers/* never appeared. - path_tails: set[str] = set() - for p in paths: - if not p: - continue - path_tails.add(p) - # e.g. /abs/root/pkg/mod.py → pkg/mod.py, mod.py - parts = p.split("/") - for i in range(1, len(parts)): - path_tails.add("/".join(parts[i:])) - - # Build a Cypher WHERE predicate that filters at the Kuzu level. - # Each tail produces one STARTS WITH predicate on qualified_name - # (or id for Import nodes). We emit the shortest unique tails only - # — if "pkg/mod.py" is present, "mod.py" is redundant because any - # match for "mod.py" also matches "pkg/mod.py". Cap at 10 tails - # to keep the WHERE clause tractable. - def _where_for_tails(prop: str, tails: set[str]) -> str: - if not tails: - return "" - # Sort longest-first so shorter redundant tails are skipped. - sorted_tails = sorted(tails, key=len, reverse=True) - kept: list[str] = [] - for t in sorted_tails: - if any(t == k or k.endswith(t) for k in kept): - continue # already covered by a longer tail - kept.append(t) - if len(kept) >= _MAX_WHERE_TAILS: - break - escaped = [t.replace("'", "\\'") for t in kept] - preds = " OR ".join(f"{prop} STARTS WITH '{t}::'" for t in escaped) - return f" WHERE {preds}" - + path_tails = build_path_tails(paths) for label in _SYMBOL_LABELS: - # Import nodes don't carry qualified_name / name — they use - # ``id`` (``::``) and ``path`` (the imported - # module). Use those as the qualified_name / name surrogate. - if label in _NON_QUALIFIED_LABELS: - prop = "s.id" - select = ( - f"MATCH (s:{label})" - "{where}" - " RETURN s.id AS qualified_name," - " s.path AS name" - ) - else: - prop = "s.qualified_name" - select = ( - f"MATCH (s:{label})" - "{where}" - " RETURN s.qualified_name AS qualified_name," - " s.name AS name" - ) - if paths: - where = _where_for_tails(prop, path_tails) - query = select.format(where=where) - else: - # Load-all mode: no filter, no limit — pull the full graph. - query = select.format(where="") + query = _build_symbol_query(label, paths, path_tails) rows = await bridge.call( "query_graph", {"graph_path": graph_path, "query": query}, ) - # Per-label batch: built, yielded, then dropped before the next - # label's query runs — peak retained inside the source is one - # label's rows, not the union across all _SYMBOL_LABELS queries. - batch: list[dict[str, Any]] = [] - for r in as_list(rows): - qn = r.get("qualified_name") - if not qn: - continue - qn_s = str(qn) - file_part, sep, _ = qn_s.partition("::") - if not sep: - continue - # Python-side match as a secondary safeguard (the WHERE - # clause is the primary filter; this handles edge cases - # where a shorter tail matched a different file). - if path_tails and not any( - p == file_part or p.endswith(file_part) or file_part.endswith(p) - for p in path_tails - ): - continue - # Resolve file_path back to the absolute form if possible. - abs_match = next( - (p for p in paths if p.endswith(file_part)), - file_part, - ) - batch.append( - { - "file_path": abs_match, - "qualified_name": qn_s, - "symbol_type": _symbol_type_from_label(label), - "signature": None, - "language": None, - "line": None, - } - ) + batch = _parse_symbol_batch(rows, label, path_tails, paths) if batch: yield batch -async def verify_symbols_async( - bridge: APBridge, - graph_path: str, - qualnames: list[str], -) -> dict[str, bool]: - """Batch verification across every AP symbol label. - - AP has no unified ``Symbol`` label — we iterate the known set - (Function, Method, Struct, ...). Wiki references are usually - bare names (``WorkflowGraphBuilder``), so we widen the match: - a qualname counts as found if any AP symbol name equals it, - its name equals the tail, or the qualified_name endswith the - tail (``::tail`` or ``.tail``). - """ - out: dict[str, bool] = {q: False for q in qualnames} +async def _collect_all_symbol_names( + bridge: APBridge, graph_path: str +) -> tuple[list[str], list[str]]: + """Query every AP symbol label once; return ``(qualified_names, + short_names)`` across all of them — the corpus ``verify_symbols_async`` + matches candidate qualnames against.""" all_names: list[str] = [] all_short: list[str] = [] for label in _SYMBOL_LABELS: @@ -268,16 +256,36 @@ async def verify_symbols_async( all_names.append(qn) if nm: all_short.append(nm) - for q in qualnames: - tail = q.rsplit(".", 1)[-1] - if tail in all_short: - out[q] = True - continue - for qn in all_names: - if qn == q or qn.endswith(f"::{tail}") or qn.endswith(f".{tail}"): - out[q] = True - break - return out + return all_names, all_short + + +def _qualname_matches(q: str, all_names: list[str], all_short: list[str]) -> bool: + """Widened match: a qualname counts as found if any AP symbol name + equals it, its short name equals the tail, or the qualified_name + endswith the tail (``::tail`` or ``.tail``) — wiki references are + usually bare names (``WorkflowGraphBuilder``).""" + tail = q.rsplit(".", 1)[-1] + if tail in all_short: + return True + return any( + qn == q or qn.endswith(f"::{tail}") or qn.endswith(f".{tail}") + for qn in all_names + ) + + +async def verify_symbols_async( + bridge: APBridge, + graph_path: str, + qualnames: list[str], +) -> dict[str, bool]: + """Batch verification across every AP symbol label. + + AP has no unified ``Symbol`` label — ``_collect_all_symbol_names`` + iterates the known set once and ``_qualname_matches`` checks each + candidate against that corpus (see its docstring for the match rule). + """ + all_names, all_short = await _collect_all_symbol_names(bridge, graph_path) + return {q: _qualname_matches(q, all_names, all_short) for q in qualnames} __all__ = [