fix(context): keep the printed slice within its budget, and build it quickly

codebase_rag/constants/cli.py

@@ -1101,6 +1101,14 @@ class DiffMarker:
 )
 # context slice (issue #1536).
 CONTEXT_WHY_TARGET = "target"
+# How many of the candidates left out of a context slice it names; the rest
+# are only counted in `omitted_count` (issue #3243).
+CONTEXT_OMITTED_SAMPLE = 10
+# How deep a context slice prints each piece: inside the slice and its
+# `pieces` list, each level one MCP_JSON_INDENT more.
+CONTEXT_PIECE_DEPTH = 2
+# Files whose decoded lines a context slice keeps between quotes of them.
+CONTEXT_SOURCE_CACHE_FILES = 256
 CONTEXT_WHY_CALLER = "direct caller: the call line"
 CONTEXT_WHY_CALLEE = "direct callee: its signature"
 CONTEXT_WHY_RETURNS = "type it returns"

codebase_rag/context_slice.py

@@ -16,8 +16,11 @@
 
 from __future__ import annotations
 
+import json
 import re
+import textwrap
 from collections.abc import Callable, Iterable
+from functools import lru_cache, partial
 from pathlib import Path
 from typing import NamedTuple, TypedDict
 
@@ -27,7 +30,7 @@
 from .graph_query import QueryFn, ReachIndex
 from .types_defs import ResultRow, SemanticSearchResult
 from .utils.path_utils import cached_resolve_posix
-from .utils.source_extraction import extract_source_lines
+from .utils.source_extraction import read_source_lines, slice_source_lines
 from .utils.token_utils import count_tokens
 
 Searcher = Callable[[str], list[SemanticSearchResult]]
@@ -53,6 +56,8 @@ class Piece(TypedDict):
     span: list[int]
     why_included: str
     source: str
+    # What the piece costs against the budget as printed: its source and
+    # every field beside it (issue #3243).
     tokens: int
 
 
@@ -62,7 +67,10 @@ class ContextSlice(TypedDict):
     budget_tokens: int
     used_tokens: int
     pieces: list[Piece]
+    # The best-ranked candidates that did not fit, at most
+    # CONTEXT_OMITTED_SAMPLE of them; `omitted_count` counts them all.
     omitted: list[str]
+    omitted_count: int
     truncated: bool
 
 
@@ -75,6 +83,9 @@ class _Candidate(NamedTuple):
     span: tuple[int, int]
     why: str
     source: str
+    # Reads the source when the piece might still fit: a central symbol
+    # has thousands of callers and tests, nearly all of them left out.
+    loader: Callable[[], str] | None = None
 
     @property
     def order(self) -> tuple[int, int, float, str]:
@@ -104,6 +115,15 @@ def _normalise(text: str) -> str:
     return text.replace("\r\n", "\n").replace("\r", "\n")
 
 
+@lru_cache(maxsize=cs.CONTEXT_SOURCE_CACHE_FILES)
+def _file_lines(path: Path, stamp: tuple[int, int, int, int]) -> tuple[str, ...]:
+    # Keyed by modification time and size as well as path, so an edited
+    # file is read again; a slice quotes one test file once per test in it.
+    # The inode and change time catch an edit or atomic replacement that
+    # keeps the old size and modification time.
+    return tuple(read_source_lines(path))
+
+
 def _lines(repo_root: Path | None, path: str | None, start: int, end: int) -> str:
     if repo_root is None or not path or start < 1 or end < start:
         return ""
@@ -113,7 +133,15 @@ def _lines(repo_root: Path | None, path: str | None, start: int, end: int) -> st
     candidate = (root / path).resolve()
     if not candidate.is_relative_to(root):
         return ""
-    return _normalise(extract_source_lines(candidate, start, end) or "")
+    try:
+        stat = candidate.stat()
+        lines = _file_lines(
+            candidate,
+            (stat.st_mtime_ns, stat.st_size, stat.st_ctime_ns, stat.st_ino),
+        )
+    except (OSError, ValueError, LookupError):
+        return ""
+    return _normalise(slice_source_lines(lines, start, end, candidate) or "")
 
 
 def _line(repo_root: Path | None, path: str | None, line: int | None) -> str:
@@ -210,7 +238,8 @@ def _caller_pieces(
     out: list[_Candidate] = []
     for row in graph_query.callers(fetch_all, project, qn):
         line = row["line"]
-        source = _line(repo_root, row["path"], line)
+        source = ""
+        loader: Callable[[], str] | None = partial(_line, repo_root, row["path"], line)
         if line is None:
             definition = graph_query.definition(
                 fetch_all, project, row["qualified_name"], repo_root
@@ -219,6 +248,7 @@ def _caller_pieces(
             source = _header(
                 repo_root, row["path"], definition["start_line"], definition["end_line"]
             )
+            loader = None
         out.append(
             _Candidate(
                 1,
@@ -229,6 +259,7 @@ def _caller_pieces(
                 (line or 0, line or 0),
                 cs.CONTEXT_WHY_CALLER,
                 source,
+                loader,
             )
         )
     return out
@@ -306,10 +337,9 @@ def _test_pieces(
     out: list[_Candidate] = []
     reach = ReachIndex.build(fetch_all, project)
     for row in reach.tests_reaching(qn):
-        definition = graph_query.definition(
-            fetch_all, project, row["qualified_name"], repo_root
-        )
-        start, end = definition["start_line"] or 0, definition["end_line"] or 0
+        # The span is in the reach index already; the source is read only
+        # if the test can still fit, never one graph query per test.
+        start, end = reach.span(row["qualified_name"])
         out.append(
             _Candidate(
                 row["depth"] + 1,
@@ -319,7 +349,8 @@ def _test_pieces(
                 row["path"],
                 (start, end),
                 cs.CONTEXT_WHY_TEST.format(depth=row["depth"], through=row["through"]),
-                _normalise(definition["source"] or "").rstrip("\n"),
+                "",
+                partial(_lines, repo_root, row["path"], start, end),
             )
         )
     return out
@@ -391,55 +422,155 @@ def _doc_section(repo_root: Path, doc_qn: str, row: ResultRow) -> _Candidate | N
 # --- the slice ---------------------------------------------------------------------
 
 
-def _fit(
-    candidates: Iterable[_Candidate], budget: int
-) -> tuple[list[Piece], list[str], int, bool]:
+def _piece_tokens(
+    qualified_name: str, file: str | None, span: list[int], why: str, source: str
+) -> int:
+    """What a piece costs as printed: its source and every field beside it.
+
+    Counting the source alone let a slice of one-line call sites print ten
+    times its budget in names, paths and reasons (issue #3243).
+    """
+    piece = Piece(
+        qualified_name=qualified_name,
+        file=file,
+        span=span,
+        why_included=why,
+        source=source,
+        tokens=0,
+    )
+    # The piece prints with its own count where it was measured with 0, and
+    # a count of four digits takes more tokens than "0" does: measure again
+    # with the count in place until it stops growing. It is measured as
+    # the slice prints it, every line indented to its depth there: measured
+    # flush left, a 10,000 budget printed 10,478 tokens of pieces and fields.
+    margin = " " * (cs.CONTEXT_PIECE_DEPTH * cs.MCP_JSON_INDENT)
+    while (
+        cost := count_tokens(
+            textwrap.indent(json.dumps(piece, indent=cs.MCP_JSON_INDENT), margin)
+        )
+    ) > piece["tokens"]:
+        piece["tokens"] = cost
+    return piece["tokens"]
+
+
+class _Fitted(NamedTuple):
+    pieces: list[Piece]
+    omitted: list[str]
+    omitted_count: int
+    used: int
+    truncated: bool
+
+
+class _Omitted:
+    # The best-ranked candidates that did not fit: a sample to print, and
+    # the count of all of them.
+    __slots__ = ("count", "sample")
+
+    def __init__(self) -> None:
+        self.sample: list[str] = []
+        self.count = 0
+
+    def add(self, candidate: _Candidate) -> None:
+        self.count += 1
+        if len(self.sample) < cs.CONTEXT_OMITTED_SAMPLE:
+            self.sample.append(f"{candidate.qualified_name} ({candidate.why})")
+
+
+def _fit(candidates: Iterable[_Candidate], budget: int) -> _Fitted:
     pieces: list[Piece] = []
-    omitted: list[str] = []
+    omitted = _Omitted()
     used = 0
     truncated = False
     seen: set[tuple[str, str, tuple[int, int]]] = set()
     for candidate in sorted(candidates, key=lambda c: c.order):
         key = (candidate.qualified_name, candidate.why, candidate.span)
-        if key in seen or not candidate.source:
+        if key in seen:
             continue
         seen.add(key)
-        source = candidate.source
-        tokens = count_tokens(source)
-        if used + tokens > budget:
-            if candidate.distance == 0 and not pieces:
-                # The target itself must fit: keep as many of its lines as
-                # the budget allows and say so. A budget too small for even
-                # its first line leaves it out rather than padding with
-                # nothing.
-                source, tokens = _trim(source, budget)
-                truncated = True
-                if not source:
-                    omitted.append(f"{candidate.qualified_name} ({candidate.why})")
-                    continue
-            else:
-                omitted.append(f"{candidate.qualified_name} ({candidate.why})")
-                continue
+        span = [candidate.span[0], candidate.span[1]]
+        is_target = candidate.distance == 0 and not pieces
+        if (sized := _sized(candidate, span, budget - used, is_target)) is None:
+            omitted.add(candidate)
+            continue
+        source, tokens, trimmed = sized
+        if trimmed:
+            truncated = True
+        if not source:
+            # No source to quote, or a target too big for even its first
+            # line: only the latter is a candidate the budget left out.
+            if trimmed:
+                omitted.add(candidate)
+            continue
         used += tokens
         pieces.append(
             Piece(
                 qualified_name=candidate.qualified_name,
                 file=candidate.file,
-                span=[candidate.span[0], candidate.span[1]],
+                span=span,
                 why_included=candidate.why,
                 source=source,
                 tokens=tokens,
             )
         )
-    return pieces, omitted, used, truncated
+    return _Fitted(pieces, omitted.sample, omitted.count, used, truncated)
+
+
+def _sized(
+    candidate: _Candidate, span: list[int], room: int, is_target: bool
+) -> tuple[str, int, bool] | None:
+    # (source, tokens, trimmed) of the candidate within `room` tokens, or
+    # None when it does not fit. Its fields alone are a floor on what it
+    # costs: past it, the source is never read.
+    floor = _piece_tokens(
+        candidate.qualified_name, candidate.file, span, candidate.why, ""
+    )
+    if floor > room and not is_target:
+        return None
+    source = candidate.source or (candidate.loader() if candidate.loader else "")
+    if not source:
+        return "", 0, False
+    tokens = _piece_tokens(
+        candidate.qualified_name, candidate.file, span, candidate.why, source
+    )
+    if tokens <= room:
+        return source, tokens, False
+    if not is_target:
+        return None
+    # The target itself must fit: keep as many of its lines as the budget
+    # allows and say so. A budget too small for even its first line leaves
+    # it out rather than padding with nothing.
+    source, tokens = _trim(candidate, span, source, room)
+    return source, tokens, True
+
 
+def _trim(
+    candidate: _Candidate, span: list[int], source: str, room: int
+) -> tuple[str, int]:
+    """The longest run of the source's leading lines whose piece fits `room`.
 
-def _trim(source: str, budget: int) -> tuple[str, int]:
+    Bisected on the line count: a few dozen encodes where trimming a line
+    at a time re-encoded the whole remaining text once per line, minutes
+    for a class of thousands of lines (issue #3243).
+    """
     lines = source.split("\n")
-    while lines and count_tokens("\n".join(lines)) > budget:
-        lines.pop()
-    text = "\n".join(lines)
-    return text, count_tokens(text) if text else 0
+
+    def cost(count: int) -> int:
+        return _piece_tokens(
+            candidate.qualified_name,
+            candidate.file,
+            span,
+            candidate.why,
+            "\n".join(lines[:count]),
+        )
+
+    low, high = 0, len(lines)
+    while low < high:
+        middle = (low + high + 1) // 2
+        if cost(middle) <= room:
+            low = middle
+        else:
+            high = middle - 1
+    return "\n".join(lines[:low]), cost(low)
 
 
 def context(
@@ -461,6 +592,7 @@ def context(
             used_tokens=0,
             pieces=[],
             omitted=[],
+            omitted_count=0,
             truncated=False,
         )
     definition = graph_query.definition(fetch_all, project, qn, repo_root)
@@ -472,6 +604,7 @@ def context(
             used_tokens=0,
             pieces=[],
             omitted=[],
+            omitted_count=0,
             truncated=False,
         )
     hotness = _hotness(fetch_all, project, qn)
@@ -488,13 +621,14 @@ def context(
         candidates = [
             c._replace(similarity=scores.get(c.qualified_name, 0.0)) for c in candidates
         ]
-    pieces, omitted, used, truncated = _fit(candidates, budget)
+    fitted = _fit(candidates, budget)
     return ContextSlice(
         target=target,
         resolved=qn,
         budget_tokens=budget,
-        used_tokens=used,
-        pieces=pieces,
-        omitted=omitted,
-        truncated=truncated,
+        used_tokens=fitted.used,
+        pieces=fitted.pieces,
+        omitted=fitted.omitted,
+        omitted_count=fitted.omitted_count,
+        truncated=fitted.truncated,
     )

codebase_rag/cypher_queries.py

@@ -902,6 +902,7 @@ def _any_label(var: str, labels: str) -> str:
 RETURN labels(n)[0] AS label, n.qualified_name AS qualified_name,
        n.name AS name, n.path AS path, n.namespace AS namespace,
        n.start_line AS start_line, n.end_line AS end_line,
+       n.decorated_start_line AS decorated_start_line,
        n.decorators AS decorators, n.is_exported AS is_exported,
        n.overrides_external AS overrides_external,
        n.is_program_main AS is_program_main,

codebase_rag/dead_code.py

@@ -1408,6 +1408,9 @@ def _node_props(row: ResultRow) -> PropertyDict:
         # non-int values coalesce to 0 so the containment checks skip them.
         cs.KEY_START_LINE: _as_line(row.get(cs.KEY_START_LINE)),
         cs.KEY_END_LINE: _as_line(row.get(cs.KEY_END_LINE)),
+        # The first decorator's line: a reaching test is quoted from it, as
+        # `definition()` quotes a decorated def (issue #2428).
+        cs.KEY_DECORATED_START_LINE: _as_line(row.get(cs.KEY_DECORATED_START_LINE)),
         cs.KEY_RUST_CFG_TEST_MODS: _as_str_list(row.get(cs.KEY_RUST_CFG_TEST_MODS)),
         cs.KEY_RUST_UNGATED_MODS: _as_str_list(row.get(cs.KEY_RUST_UNGATED_MODS)),
         # The signature the C# entry-point rule checks (issue #2471).