From 77852e07422a1f3eeb0e943318dcb098612dab11 Mon Sep 17 00:00:00 2001 From: Wolfvin Date: Tue, 14 Jul 2026 19:02:07 +0700 Subject: [PATCH 1/2] feat(context): optional LSP-backed find-references under --deep (closes #255) trace-up under --deep + active LSP uses lsp_client.find_references as precision source (annotated trace_source=lsp); zero-config/graph path unchanged (trace_source=graph). Reuses HybridEngine find_references/_find_symbol_definition/_find_symbol_char + _filter_external_references; no new LSP infra. Adds HybridEngine.find_references_for_symbol + commands.trace._apply_lsp_trace_up + tests. --- scripts/commands/trace.py | 61 ++++++++ scripts/hybrid_engine.py | 50 +++++++ tests/test_issue255_lsp_references.py | 191 ++++++++++++++++++++++++++ 3 files changed, 302 insertions(+) create mode 100644 tests/test_issue255_lsp_references.py diff --git a/scripts/commands/trace.py b/scripts/commands/trace.py index 762eace3..68e5e142 100644 --- a/scripts/commands/trace.py +++ b/scripts/commands/trace.py @@ -64,6 +64,16 @@ def execute(args, workspace): max_results=args.max_results, use_graph=use_graph, ) + # Issue #255: opt-in LSP-backed find-references for trace-up precision. + # Only when --deep is active AND an LSP server is available AND we are + # tracing callers (up/both). Otherwise the graph path above is used + # unchanged (zero-config, no regression, no hang). + if isinstance(result, dict) and getattr(args, "deep", False) \ + and args.direction in ("up", "both"): + _apply_lsp_trace_up(args.name, workspace, result) + else: + if isinstance(result, dict): + result.setdefault("trace_source", "graph") # Apply pagination to chains.up and chains.down (issue #17). if isinstance(result, dict) and isinstance(result.get("chains"), dict): chains = result["chains"] @@ -81,4 +91,55 @@ def execute(args, workspace): result["limit"] = limit return result +def _apply_lsp_trace_up(name, workspace, result): + """Replace ``result['chains']['up']`` with LSP-derived references when a + language server is available (issue #255). + + Annotates ``result['trace_source']`` as ``"lsp"`` on success or ``"graph"`` + when LSP is unavailable / cannot resolve the symbol, so consumers know the + precision source. Falls back to the graph chains (leaves them untouched) on + any failure — LSP is a precision enhancement, never a hard dependency. + """ + graph_up = result.get("chains", {}).get("up", []) if isinstance(result.get("chains"), dict) else [] + try: + from hybrid_engine import create_hybrid_engine + engine = create_hybrid_engine(workspace, deep=True) + except Exception: + result["trace_source"] = "graph" + return + try: + if not engine.lsp_active: + result["trace_source"] = "graph" + result.setdefault("lsp_available", False) + return + refs = engine.find_references_for_symbol(name) + finally: + engine.cleanup() + + result["lsp_available"] = True + if refs is None: + # LSP active but symbol unresolved / no references list — keep graph. + result["trace_source"] = "graph" + return + + lsp_up = [] + for ref in refs: + lsp_up.append({ + "fn": "", + "file": ref.get("file", ""), + "line": ref.get("line", 0), + "depth": 1, + "source": "lsp", + }) + if isinstance(result.get("chains"), dict): + result["chains"]["up"] = lsp_up + else: + result["chains"] = {"up": lsp_up, "down": []} + result["trace_source"] = "lsp" + result["graph_callers_found"] = len(graph_up) + result["lsp_callers_found"] = len(lsp_up) + stats = result.setdefault("stats", {}) + stats["callers_found"] = len(lsp_up) + + # Issue #199: deprecated "trace" alias registration removed; this module is now an implementation module imported by the "context" umbrella command. diff --git a/scripts/hybrid_engine.py b/scripts/hybrid_engine.py index f69856d4..b2dc4840 100644 --- a/scripts/hybrid_engine.py +++ b/scripts/hybrid_engine.py @@ -400,6 +400,56 @@ def _find_symbol_definition(self, symbol_name: str) -> Tuple[Optional[str], Opti pass return None, None + def find_references_for_symbol(self, symbol_name: str) -> Optional[List[Dict]]: + """Resolve ``symbol_name`` to its definition, then ask the LSP server + for its references (issue #255 — LSP-backed trace-up precision). + + Reuses the existing ``lsp_client.find_references`` + + ``_find_symbol_definition`` + ``_find_symbol_char`` machinery — no new + LSP infrastructure. Returns a list of reference dicts:: + + {"file": , "line": <1-indexed>, "character": } + + excluding the definition site itself (the caller wants callers, not the + declaration). Returns ``None`` when LSP is not active or the symbol + cannot be resolved/located, so the caller can distinguish "no LSP path" + from "LSP ran and found zero references" (empty list). Never raises. + """ + if not self.lsp_active: + return None + def_file, def_line = self._find_symbol_definition(symbol_name) + if not def_file or not def_line: + return None + abs_def = def_file if os.path.isabs(def_file) else os.path.join(self.workspace, def_file) + if not os.path.exists(abs_def): + return None + self.open_file_for_lsp(abs_def) + client = self.get_lsp_client(abs_def) + if not client: + return None + char = self._find_symbol_char(abs_def, def_line, symbol_name) + if char is None: + char = 0 + lsp_line = max(0, def_line - 1) + try: + raw = client.find_references(abs_def, lsp_line, char, include_declaration=False) + except Exception: + return None + if raw is None: + return None + external = self._filter_external_references(raw, abs_def, lsp_line, char) + out: List[Dict] = [] + for ref in external: + ref_uri = ref.get("uri", "") + ref_path = _uri_to_path(ref_uri) if ref_uri else "" + start = ref.get("range", {}).get("start", {}) + out.append({ + "file": ref_path, + "line": start.get("line", 0) + 1, # LSP 0-indexed -> report 1-indexed + "character": start.get("character", 0), + }) + return out + def _paths_match(path_a: str, path_b: str) -> bool: """Compare two file paths for equality using normalized absolute paths. diff --git a/tests/test_issue255_lsp_references.py b/tests/test_issue255_lsp_references.py new file mode 100644 index 00000000..110057cf --- /dev/null +++ b/tests/test_issue255_lsp_references.py @@ -0,0 +1,191 @@ +# @WHO: tests/test_issue255_lsp_references.py +# @WHAT: Tests for optional LSP-backed find-references in trace-up (issue #255) +# @PART: tests +"""Tests for LSP-backed trace-up precision (issue #255). + +Feature: when ``--deep`` is active AND an LSP server is available, +``context --check trace --direction up`` uses LSP ``textDocument/references`` +(via ``lsp_client.find_references``) as the precision source; otherwise it +falls back to the existing graph path unchanged. + +Two halves, mirroring #253's split (live env can't verify the LSP happy path): + +1. Graceful degradation — LIVE. A real CLI-equivalent trace with no ``--deep`` + (and with ``--deep`` when the symbol/LSP can't resolve) must keep working + via the graph path, never hang, never error, and be annotated + ``trace_source == "graph"``. + +2. LSP happy path — MOCKED. ``lsp_client.find_references`` and the hybrid + engine are mocked so the precision path can be exercised without a live + language server (rust-analyzer, the only server installed in the dev env, + does not respond to ``initialize`` within 60s — same limitation as #253). +""" + +import os +import sys +import tempfile +from unittest import mock + +import pytest + +SCRIPT_DIR = os.path.join( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "scripts" +) +if SCRIPT_DIR not in sys.path: + sys.path.insert(0, SCRIPT_DIR) + +from commands import trace as trace_cmd # noqa: E402 +import hybrid_engine # noqa: E402 + + +# ─── Half 1: graceful degradation (LIVE, no LSP dependency) ────────────── + +def _make_scanned_ws(): + """Create + scan a tiny workspace; return its path (real graph tables).""" + ws = tempfile.mkdtemp(prefix="cl255_") + with open(os.path.join(ws, "mod.py"), "w", encoding="utf-8") as f: + f.write( + "def helper(x):\n" + " return x + 1\n\n" + "def caller_a():\n" + " return helper(1)\n\n" + "def caller_b():\n" + " return helper(2)\n" + ) + import codelens # noqa: F401 (ensures scripts importable) + from commands import scan as scan_cmd + import argparse + ns = argparse.Namespace(workspace=ws, format="json") + scan_cmd.execute(ns, ws) + return ws + + +def _trace_ns(name, ws, direction="up", deep=False): + import argparse + return argparse.Namespace( + name=name, workspace=ws, direction=direction, depth=10, + domain="auto", max_results=1000, limit=20, offset=0, + use_graph=True, deep=deep, + ) + + +def test_no_deep_uses_graph_path_live(): + ws = _make_scanned_ws() + result = trace_cmd.execute(_trace_ns("helper", ws, deep=False), ws) + assert result["trace_source"] == "graph" + # graph path still finds the callers — no regression + assert result["stats"]["callers_found"] >= 1 + assert "lsp_available" not in result # LSP path never touched + + +def test_deep_but_symbol_unresolved_falls_back_to_graph_live(): + """--deep on, but a nonexistent symbol can't be resolved for LSP → + engine returns None refs → graph path retained, no error, no hang.""" + ws = _make_scanned_ws() + result = trace_cmd.execute(_trace_ns("helper", ws, deep=True), ws) + # Whatever the live LSP does, output must be well-formed and never crash. + assert result["status"] == "ok" + assert result["trace_source"] in ("graph", "lsp") + assert isinstance(result["chains"]["up"], list) + + +# ─── Half 2: LSP happy path (MOCKED) ───────────────────────────────────── + +def test_apply_lsp_trace_up_rewrites_chains_when_lsp_active(): + """With a mocked hybrid engine reporting LSP references, trace-up is + rewritten from the graph chains to the LSP chains and annotated 'lsp'.""" + result = { + "status": "ok", "symbol": "helper", + "chains": {"up": [{"fn": "caller_a", "file": "mod.py", "line": 5}], "down": []}, + "stats": {"callers_found": 1}, + } + fake_engine = mock.Mock() + fake_engine.lsp_active = True + fake_engine.find_references_for_symbol.return_value = [ + {"file": "/abs/mod.py", "line": 5, "character": 11}, + {"file": "/abs/mod.py", "line": 8, "character": 11}, + {"file": "/abs/other.py", "line": 3, "character": 4}, + ] + with mock.patch("hybrid_engine.create_hybrid_engine", return_value=fake_engine): + trace_cmd._apply_lsp_trace_up("helper", "/ws", result) + + assert result["trace_source"] == "lsp" + assert result["lsp_available"] is True + assert result["stats"]["callers_found"] == 3 + assert len(result["chains"]["up"]) == 3 + assert all(e["source"] == "lsp" for e in result["chains"]["up"]) + assert result["graph_callers_found"] == 1 + assert result["lsp_callers_found"] == 3 + fake_engine.cleanup.assert_called_once() + + +def test_apply_lsp_trace_up_keeps_graph_when_lsp_inactive(): + result = { + "chains": {"up": [{"fn": "caller_a", "file": "mod.py", "line": 5}], "down": []}, + "stats": {"callers_found": 1}, + } + fake_engine = mock.Mock() + fake_engine.lsp_active = False + with mock.patch("hybrid_engine.create_hybrid_engine", return_value=fake_engine): + trace_cmd._apply_lsp_trace_up("helper", "/ws", result) + assert result["trace_source"] == "graph" + assert result["lsp_available"] is False + assert result["chains"]["up"] == [{"fn": "caller_a", "file": "mod.py", "line": 5}] + + +def test_apply_lsp_trace_up_keeps_graph_when_refs_none(): + """LSP active but symbol can't be resolved (refs None) → keep graph.""" + result = {"chains": {"up": [{"fn": "caller_a"}], "down": []}, "stats": {}} + fake_engine = mock.Mock() + fake_engine.lsp_active = True + fake_engine.find_references_for_symbol.return_value = None + with mock.patch("hybrid_engine.create_hybrid_engine", return_value=fake_engine): + trace_cmd._apply_lsp_trace_up("helper", "/ws", result) + assert result["trace_source"] == "graph" + assert result["chains"]["up"] == [{"fn": "caller_a"}] + + +def test_apply_lsp_trace_up_engine_creation_failure_falls_back(): + result = {"chains": {"up": [], "down": []}, "stats": {}} + with mock.patch("hybrid_engine.create_hybrid_engine", side_effect=RuntimeError("boom")): + trace_cmd._apply_lsp_trace_up("helper", "/ws", result) + assert result["trace_source"] == "graph" + + +def test_find_references_for_symbol_resolves_and_filters(): + """HybridEngine.find_references_for_symbol: mock the LSP client's + find_references + symbol definition resolution; verify def-site excluded, + 0→1-indexed conversion, and shape.""" + eng = hybrid_engine.HybridEngine.__new__(hybrid_engine.HybridEngine) + eng.deep = True + eng._lsp_available = True # lsp_active property = deep and _lsp_available + eng.workspace = os.getcwd() + + with tempfile.TemporaryDirectory() as d: + fpath = os.path.join(d, "mod.py") + with open(fpath, "w", encoding="utf-8") as f: + f.write("def helper(x):\n return x\n\ncaller = helper\n") + + from lsp_client import _path_to_uri + fake_client = mock.Mock() + # LSP returns one real reference at line 3 (0-indexed). + fake_client.find_references.return_value = [ + {"uri": _path_to_uri(fpath), + "range": {"start": {"line": 3, "character": 9}}}, + ] + eng._find_symbol_definition = lambda n: (fpath, 1) + eng.open_file_for_lsp = lambda p: None + eng.get_lsp_client = lambda p: fake_client + + refs = eng.find_references_for_symbol("helper") + assert refs is not None + assert len(refs) == 1 + assert refs[0]["line"] == 4 # 0-indexed 3 -> 1-indexed 4 + assert refs[0]["character"] == 9 + + +def test_find_references_for_symbol_none_when_lsp_inactive(): + eng = hybrid_engine.HybridEngine.__new__(hybrid_engine.HybridEngine) + eng.deep = False + eng._lsp_available = False + assert eng.find_references_for_symbol("helper") is None From d0f5c7d286eb34794e9c99ab69beda0234589e44 Mon Sep 17 00:00:00 2001 From: Wolfvin Date: Tue, 14 Jul 2026 19:04:01 +0700 Subject: [PATCH 2/2] docs(context): design doc for LSP-backed find-references (#255) --- docs/design/0255-lsp-find-references.md | 116 ++++++++++++++++++++++++ 1 file changed, 116 insertions(+) create mode 100644 docs/design/0255-lsp-find-references.md diff --git a/docs/design/0255-lsp-find-references.md b/docs/design/0255-lsp-find-references.md new file mode 100644 index 00000000..0bc79a12 --- /dev/null +++ b/docs/design/0255-lsp-find-references.md @@ -0,0 +1,116 @@ +# Design Doc: Optional LSP-backed find-references for trace-up precision + +> **Status:** Accepted +> **Date:** 2026-07-14 +> **Author:** Claude (direct implementation, no worker — user directive) +> **Related issues:** #255 +> **Related PRs:** (this PR) + +--- + +## Problem + +Gap-analysis vs Serena MCP: Serena's find-references uses the language +server (`textDocument/references`) — a real AST/symbol table, high precision, +no missed references. CodeLens's caller/reference discovery +(`context --check trace --direction up`) uses a home-grown call graph that is +an *approximation*. Evidence: a run of ref-count/trace edge-case bugs were +found and fixed across the project (#210, #219, #222, #223 module-level +callers). The graph will always have edge cases; LSP `textDocument/references` +does not. + +CodeLens already had the LSP capability: `lsp_client.py:350` +`find_references(file, line, character)` issues `textDocument/references`, and +`hybrid_engine.py` already used it internally to *verify* dead-code and +enhance impact (`_filter_external_references`). But that precision was never +exposed as a navigation path for agents. + +## Goal + +When `--deep` is active **and** an LSP server is available, +`context --check trace --direction up` (and `--direction both`) uses LSP +`textDocument/references` as the precision source for callers, annotating the +result `trace_source: "lsp"`. Without `--deep`, or without a live LSP server, +or when the symbol can't be resolved/located — the existing graph path is used +unchanged (`trace_source: "graph"`). Zero-config keeps working with no +regression and no LSP dependency. + +## Changes + +### Modified Files +- `scripts/hybrid_engine.py` — new + `HybridEngine.find_references_for_symbol(symbol_name)`. Reuses existing + machinery only: `_find_symbol_definition` (registry lookup) to resolve the + symbol → `(file, line)`, `_find_symbol_char` to locate the column, then + `lsp_client.find_references(..., include_declaration=False)`, then + `_filter_external_references` to drop the definition site. Converts LSP + 0-indexed lines to 1-indexed. Returns `None` (not `[]`) when LSP is + inactive or the symbol can't be resolved, so the caller can distinguish + "no LSP path" from "LSP ran, found zero references". Never raises. +- `scripts/commands/trace.py`: + - `execute()` — after the graph `trace_symbol` call, when `args.deep` is + truthy and `direction in ("up", "both")`, calls the new + `_apply_lsp_trace_up`; otherwise annotates `trace_source: "graph"`. + - `_apply_lsp_trace_up(name, workspace, result)` — creates a hybrid engine + with `deep=True`, and **only if `engine.lsp_active`** replaces + `result["chains"]["up"]` with LSP-derived caller entries + (`source: "lsp"`), sets `trace_source: "lsp"`, and records + `graph_callers_found` / `lsp_callers_found` for A/B comparison. On engine + creation failure, inactive LSP, or `None` refs, it leaves the graph + chains untouched and annotates `trace_source: "graph"`. Always calls + `engine.cleanup()`. + +### No new LSP infrastructure +Per the issue constraint, this reuses `lsp_client.find_references` and the +existing `hybrid_engine` resolution/filter helpers. `find_references_for_symbol` +is orchestration over those, not new LSP plumbing. LSP is never made a hard +dependency — the graph path is the default and the fallback. + +### Placement rationale +The precision upgrade lives at the command boundary (`commands/trace.py`), +not in `trace_engine.py`. `trace_engine` stays a pure graph/flat backend with +an unchanged output shape; the opt-in LSP overlay is applied on top only when +`--deep` + LSP are present. This keeps the zero-config trace path completely +untouched and easy to reason about. + +## Testing + +`tests/test_issue255_lsp_references.py` (8 tests): + +**Graceful degradation — live (real scan + CLI-equivalent trace):** +- no `--deep` → `trace_source: "graph"`, callers still found, LSP path never + touched (no `lsp_available` key). +- `--deep` on a real scanned workspace → well-formed output, `status: ok`, + no crash, no hang, `trace_source` in `{graph, lsp}`. + +**LSP happy path — mocked** (`create_hybrid_engine` / `find_references` +mocked, mirroring #253): +- LSP active + refs → chains.up rewritten to LSP entries, `trace_source: lsp`, + stats + `graph/lsp_callers_found` updated, `cleanup()` called. +- LSP inactive → graph retained, `lsp_available: false`. +- refs `None` (symbol unresolved) → graph retained. +- engine creation raises → graph retained. +- `find_references_for_symbol` resolves def site, excludes it, converts + 0→1-indexed; returns `None` when LSP inactive. + +**Live verification (real CLI, this environment):** +- `codelens context --check trace --name helper --direction up` → + `trace_source: graph`, callers found — zero-config unaffected. +- same with `--deep` → `lsp_available: true`, but the live server did not + return usable references for the symbol, so it degraded to + `trace_source: graph` — no hang, no error, exit 0. + +**LSP happy-path live limitation (honest):** the LSP happy path +(`trace_source: lsp` with real references) could **not** be verified against a +live server in the dev environment — rust-analyzer (the only installed +server) does not respond to `initialize` within 60s (pre-existing, same +limitation documented in #253). The happy path is covered by the mocked tests +above; only the graph fallback and graceful-degradation paths are +live-verified. + +## Backward compatibility + +Zero-config (`context --check trace ...` without `--deep`) is byte-for-byte +unchanged — the graph result only gains a `trace_source: "graph"` annotation. +No behavior change to `trace_engine.py`. `--direction down` is never touched +by this feature (callees are not references).