From bf43d3b82e983f98a3c0ca9d2db3bb3e19495852 Mon Sep 17 00:00:00 2001 From: vitali87 Date: Sat, 10 Oct 2026 00:06:36 +0000 Subject: [PATCH 1/6] context: keep the printed slice within its budget, and build it quickly `context` counted only each piece's source against the budget, listed every candidate that did not fit in `omitted`, fetched every reaching test's definition before budgeting, and trimmed an oversized target one line and one full re-encode at a time. A 2-line function reached by 2,000 tests printed 163,019 tokens for a 4,000 budget after 21 s, and a 9,000-line class took 87 s to trim. - A piece is charged for what it prints: its source and every field beside it (`tokens` and `used_tokens` count both). - `omitted` names the best-ranked ten candidates left out, and the new `omitted_count` counts them all. - Callers' and tests' sources are read only for a candidate whose fields still fit; a test's span comes from the reach index already built instead of one definition query per test, and each quoted file is decoded once (`read_source_lines`, factored out of `extract_source_lines`, behind a cache keyed by mtime and size). - An oversized target is trimmed by bisecting its line count: a few dozen encodes instead of one per line, with the same lines kept. The context-slice tests that pinned a 25-token slice of the target now size their budgets from the target's own fields, and the normalisation test drives the test piece through its loader. Fixes #3243 Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01EbsCWQcSTGQXv3g5sk418m --- codebase_rag/constants/cli.py | 5 + codebase_rag/context_slice.py | 171 +++++++++--- codebase_rag/graph_query.py | 12 + .../tests/test_context_budget_output.py | 244 ++++++++++++++++++ codebase_rag/tests/test_context_slice.py | 16 +- .../tests/test_context_slice_normalisation.py | 23 +- codebase_rag/utils/source_extraction.py | 71 ++--- docs/architecture/context-slice.md | 28 +- 8 files changed, 476 insertions(+), 94 deletions(-) create mode 100644 codebase_rag/tests/test_context_budget_output.py diff --git a/codebase_rag/constants/cli.py b/codebase_rag/constants/cli.py index cd82b076a2..a5b6308e67 100644 --- a/codebase_rag/constants/cli.py +++ b/codebase_rag/constants/cli.py @@ -910,6 +910,11 @@ 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 +# 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" diff --git a/codebase_rag/context_slice.py b/codebase_rag/context_slice.py index e76796f5ba..8d7f2b616a 100644 --- a/codebase_rag/context_slice.py +++ b/codebase_rag/context_slice.py @@ -16,8 +16,10 @@ from __future__ import annotations +import json import re from collections.abc import Callable, Iterable +from functools import lru_cache, partial from pathlib import Path from typing import NamedTuple, TypedDict @@ -27,7 +29,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 +55,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 +66,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 +82,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 +114,13 @@ 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]) -> 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. + 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 +130,12 @@ 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)) + 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 +232,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 +242,7 @@ def _caller_pieces( source = _header( repo_root, row["path"], definition["start_line"], definition["end_line"] ) + loader = None out.append( _Candidate( 1, @@ -229,6 +253,7 @@ def _caller_pieces( (line or 0, line or 0), cs.CONTEXT_WHY_CALLER, source, + loader, ) ) return out @@ -306,10 +331,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 +343,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 +416,122 @@ 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, + ) + return count_tokens(json.dumps(piece, indent=cs.MCP_JSON_INDENT)) + + +class _Fitted(NamedTuple): + pieces: list[Piece] + omitted: list[str] + omitted_count: int + used: int + truncated: bool + + +def _fit(candidates: Iterable[_Candidate], budget: int) -> _Fitted: pieces: list[Piece] = [] omitted: list[str] = [] + omitted_count = 0 used = 0 truncated = False seen: set[tuple[str, str, tuple[int, int]]] = set() + + def omit(candidate: _Candidate) -> None: + nonlocal omitted_count + omitted_count += 1 + if len(omitted) < cs.CONTEXT_OMITTED_SAMPLE: + omitted.append(f"{candidate.qualified_name} ({candidate.why})") + 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) + span = [candidate.span[0], candidate.span[1]] + is_target = candidate.distance == 0 and not pieces + # 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 used + floor > budget and not is_target: + omit(candidate) + continue + source = candidate.source or (candidate.loader() if candidate.loader else "") + if not source: + continue + tokens = _piece_tokens( + candidate.qualified_name, candidate.file, span, candidate.why, 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})") + if not is_target: + omit(candidate) + continue + # 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, budget - used) + truncated = True + if not source: + omit(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, omitted_count, used, truncated) + +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 +553,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 +565,7 @@ def context( used_tokens=0, pieces=[], omitted=[], + omitted_count=0, truncated=False, ) hotness = _hotness(fetch_all, project, qn) @@ -488,13 +582,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, ) diff --git a/codebase_rag/graph_query.py b/codebase_rag/graph_query.py index d7aed4ca6e..5e03386e50 100644 --- a/codebase_rag/graph_query.py +++ b/codebase_rag/graph_query.py @@ -625,6 +625,18 @@ def _walk(self, qualified_name: str) -> tuple[dict[str, int], dict[str, str]]: frontier = next_frontier return depth_of, through_of + def span(self, qualified_name: str) -> tuple[int, int]: + """The recorded (start, end) lines of a node, 0 where unknown.""" + entry = self._by_qn.get(qualified_name) + if entry is None: + return 0, 0 + props = entry[1] + start, end = props.get(cs.KEY_START_LINE), props.get(cs.KEY_END_LINE) + return ( + start if isinstance(start, int) else 0, + end if isinstance(end, int) else 0, + ) + def tests_reaching(self, qualified_name: str) -> list[TestReachRow]: depth_of, through_of = self._walk(qualified_name) out: list[TestReachRow] = [] diff --git a/codebase_rag/tests/test_context_budget_output.py b/codebase_rag/tests/test_context_budget_output.py new file mode 100644 index 0000000000..c62090cb5f --- /dev/null +++ b/codebase_rag/tests/test_context_budget_output.py @@ -0,0 +1,244 @@ +"""The context slice's printed output stays within its token budget, fast. + +`context` counted only each piece's source against the budget, listed every +candidate that did not fit in `omitted`, fetched every reaching test's +definition before budgeting, and trimmed an oversized target one line and +one full re-encode at a time. A 2-line function reached by 2,000 tests +printed 163,019 tokens for a 4,000 budget after 21 s, and a 9,000-line class +took 87 s to trim (issue #3243). +""" + +from __future__ import annotations + +import json +import time +from pathlib import Path + +import pytest + +from codebase_rag import constants as cs +from codebase_rag import context_slice +from codebase_rag.context_slice import context +from codebase_rag.graph_query import DefinitionRow, QueryFn +from codebase_rag.graph_updater import GraphUpdater +from codebase_rag.parser_loader import load_parsers +from codebase_rag.utils.token_utils import count_tokens +from evals.cgr_graph import _StatefulIngestor + +PROJECT = "ctxb" +TEST_FILES = 6 +TESTS_PER_FILE = 50 +# The slice's own fields beside its pieces: target, resolved, the counts and +# a capped sample of what was left out. +ENVELOPE = 400 + + +def _repo(root: Path, files: dict[str, str]) -> _StatefulIngestor: + for rel, text in files.items(): + (root / rel).parent.mkdir(parents=True, exist_ok=True) + (root / rel).write_text(text, encoding="utf-8") + store = _StatefulIngestor() + parsers, queries = load_parsers() + GraphUpdater( + ingestor=store, + repo_path=root, + parsers=parsers, + queries=queries, + project_name=PROJECT, + ).run(force=True) + return store + + +def _central(root: Path) -> _StatefulIngestor: + files = {"app.py": "def core(x):\n return x + 1\n", "tests/__init__.py": ""} + for f in range(TEST_FILES): + body = "from app import core\n\n\n" + "".join( + f"def test_{f}_{i}():\n assert core({i}) == {i + 1}\n\n\n" + for i in range(TESTS_PER_FILE) + ) + files[f"tests/test_{f}.py"] = body + return _repo(root, files) + + +def _printed(slice_: context_slice.ContextSlice) -> int: + return count_tokens(json.dumps(slice_, indent=cs.MCP_JSON_INDENT)) + + +@pytest.mark.parametrize("budget", [4000, 1000, 300]) +def test_the_printed_slice_fits_the_budget(temp_repo: Path, budget: int) -> None: + store = _central(temp_repo) + + slice_ = context(store.fetch_all, PROJECT, f"{PROJECT}.app.core", budget, temp_repo) + + total = TEST_FILES * TESTS_PER_FILE + assert slice_["used_tokens"] <= budget + assert slice_["used_tokens"] == sum(p["tokens"] for p in slice_["pieces"]) + assert _printed(slice_) <= budget + ENVELOPE, _printed(slice_) + # What was left out is counted, and named only in a short sample. + assert len(slice_["omitted"]) <= cs.CONTEXT_OMITTED_SAMPLE + assert slice_["omitted_count"] > 0 + assert slice_["omitted_count"] + len(slice_["pieces"]) >= total + + +def test_a_piece_is_charged_for_its_fields(temp_repo: Path) -> None: + store = _central(temp_repo) + + slice_ = context(store.fetch_all, PROJECT, f"{PROJECT}.app.core", 4000, temp_repo) + + for piece in slice_["pieces"]: + assert piece["tokens"] > count_tokens(piece["source"]), piece + + +def test_each_quoted_file_is_read_once_and_no_test_is_fetched_by_name( + temp_repo: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + store = _central(temp_repo) + reads: list[Path] = [] + real = context_slice.read_source_lines + + def counting(path: Path) -> list[str]: + reads.append(path) + return real(path) + + definitions: list[str] = [] + real_definition = context_slice.graph_query.definition + + def counting_definition( + fetch_all: QueryFn, project: str, qualified_name: str, root: Path | None + ) -> DefinitionRow: + definitions.append(qualified_name) + return real_definition(fetch_all, project, qualified_name, root) + + monkeypatch.setattr(context_slice, "read_source_lines", counting) + monkeypatch.setattr(context_slice.graph_query, "definition", counting_definition) + + slice_ = context(store.fetch_all, PROJECT, f"{PROJECT}.app.core", 4000, temp_repo) + + # 300 tests and 300 call sites in six files reach `core`: each file is + # decoded once, and only the target's definition comes from the graph. + assert len(reads) == len(set(reads)) <= TEST_FILES, reads + assert definitions == [f"{PROJECT}.app.core"] + assert slice_["pieces"] + + +def test_nothing_is_read_once_the_budget_is_spent( + temp_repo: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + # A budget the target fills leaves every caller and test unread. + store = _central(temp_repo) + full = context(store.fetch_all, PROJECT, f"{PROJECT}.app.core", 4000, temp_repo) + target_tokens = full["pieces"][0]["tokens"] + loads: list[tuple[str | None, int]] = [] + real = context_slice._lines + + def counting(root: Path | None, path: str | None, start: int, end: int) -> str: + loads.append((path, start)) + return real(root, path, start, end) + + monkeypatch.setattr(context_slice, "_lines", counting) + + slice_ = context( + store.fetch_all, PROJECT, f"{PROJECT}.app.core", target_tokens, temp_repo + ) + + assert [p["why_included"] for p in slice_["pieces"]] == [cs.CONTEXT_WHY_TARGET] + assert loads == [] + assert slice_["omitted_count"] == 2 * TEST_FILES * TESTS_PER_FILE + + +def test_a_large_target_trims_without_re_encoding_every_line( + temp_repo: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + body = "class Big:\n" + "".join( + f" def m{i}(self, x):\n return x * {i}\n\n" for i in range(1500) + ) + store = _repo(temp_repo, {"big.py": body}) + calls = 0 + real = context_slice.count_tokens + + def counting(text: str) -> int: + nonlocal calls + calls += 1 + return real(text) + + monkeypatch.setattr(context_slice, "count_tokens", counting) + started = time.perf_counter() + + slice_ = context(store.fetch_all, PROJECT, f"{PROJECT}.big.Big", 500, temp_repo) + + elapsed = time.perf_counter() - started + assert slice_["truncated"] + (target,) = slice_["pieces"] + assert target["source"].startswith("class Big:\n def m0(self, x):") + assert slice_["used_tokens"] <= 500 + # Bisecting the line count: a few dozen encodes, not one per line. + assert calls < 100, calls + assert elapsed < 5, elapsed + + +def test_the_trim_keeps_the_most_whole_lines_that_fit(temp_repo: Path) -> None: + # The bisected trim keeps exactly the lines a line-by-line trim kept. + body = "class Big:\n" + "".join( + f" def m{i}(self, x):\n return x * {i}\n\n" for i in range(60) + ) + store = _repo(temp_repo, {"big.py": body}) + + slice_ = context(store.fetch_all, PROJECT, f"{PROJECT}.big.Big", 300, temp_repo) + + (target,) = slice_["pieces"] + lines = body.rstrip("\n").split("\n") + kept = target["source"].split("\n") + assert kept == lines[: len(kept)] + one_more = "\n".join(lines[: len(kept) + 1]) + assert ( + context_slice._piece_tokens( + f"{PROJECT}.big.Big", + "big.py", + [1, len(lines)], + cs.CONTEXT_WHY_TARGET, + one_more, + ) + > 300 + ) + + +def test_a_budget_that_exactly_fits_some_lines_keeps_them(temp_repo: Path) -> None: + # A prefix whose piece costs exactly the budget fits: the trim's bound + # is inclusive. + body = "class Big:\n" + "".join( + f" def m{i}(self, x):\n return x * {i}\n\n" for i in range(60) + ) + store = _repo(temp_repo, {"big.py": body}) + lines = body.rstrip("\n").split("\n") + full = context(store.fetch_all, PROJECT, f"{PROJECT}.big.Big", 10_000, temp_repo) + span = full["pieces"][0]["span"] + budget = context_slice._piece_tokens( + f"{PROJECT}.big.Big", + "big.py", + span, + cs.CONTEXT_WHY_TARGET, + "\n".join(lines[:7]), + ) + + slice_ = context(store.fetch_all, PROJECT, f"{PROJECT}.big.Big", budget, temp_repo) + + assert slice_["pieces"][0]["source"] == "\n".join(lines[:7]) + assert slice_["used_tokens"] == budget + + +def test_a_neighbourhood_that_fits_omits_nothing(temp_repo: Path) -> None: + # Negative: a small slice is whole, with nothing counted as left out. + store = _repo( + temp_repo, + { + "app.py": "def core(x):\n return x + 1\n\n\ndef use():\n return core(1)\n", + }, + ) + + slice_ = context(store.fetch_all, PROJECT, f"{PROJECT}.app.core", 4000, temp_repo) + + assert slice_["omitted"] == [] and slice_["omitted_count"] == 0 + assert {p["why_included"] for p in slice_["pieces"]} == { + cs.CONTEXT_WHY_TARGET, + cs.CONTEXT_WHY_CALLER, + } diff --git a/codebase_rag/tests/test_context_slice.py b/codebase_rag/tests/test_context_slice.py index 7b7c71cd93..85bf54bf2f 100644 --- a/codebase_rag/tests/test_context_slice.py +++ b/codebase_rag/tests/test_context_slice.py @@ -10,7 +10,7 @@ import pytest from codebase_rag import constants as cs -from codebase_rag.context_slice import _doc_pieces, context +from codebase_rag.context_slice import _doc_pieces, _piece_tokens, context from codebase_rag.graph_updater import GraphUpdater from codebase_rag.parser_loader import load_parsers from codebase_rag.utils.token_utils import count_tokens @@ -161,12 +161,22 @@ def test_token_count_never_exceeds_the_budget( ) -> None: root, store, _updater = repo full = context(store.fetch_all, PROJECT, _qn("pkg.util.helper"), 4000, root) - for budget in (full["used_tokens"] - 1, 60, 25, 5): + # A piece's fields count against the budget too (issue #3243): the + # target with an empty source is the floor any slice of it starts at. + target = full["pieces"][0] + floor = _piece_tokens( + target["qualified_name"], + target["file"], + target["span"], + target["why_included"], + "", + ) + for budget in (full["used_tokens"] - 1, floor + 40, floor + 15, 5): slice_ = context(store.fetch_all, PROJECT, _qn("pkg.util.helper"), budget, root) assert slice_["used_tokens"] <= budget assert sum(count_tokens(p["source"]) for p in slice_["pieces"]) <= budget assert slice_["omitted"] or slice_["truncated"] - shrunk = context(store.fetch_all, PROJECT, _qn("pkg.util.helper"), 25, root) + shrunk = context(store.fetch_all, PROJECT, _qn("pkg.util.helper"), floor + 15, root) # The target keeps as many of its lines as fit and comes first. assert shrunk["truncated"] assert shrunk["pieces"][0]["why_included"] == cs.CONTEXT_WHY_TARGET diff --git a/codebase_rag/tests/test_context_slice_normalisation.py b/codebase_rag/tests/test_context_slice_normalisation.py index eb861a4f00..683dad5e05 100644 --- a/codebase_rag/tests/test_context_slice_normalisation.py +++ b/codebase_rag/tests/test_context_slice_normalisation.py @@ -112,7 +112,6 @@ def test_test_pieces_normalises_the_source_it_reads( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: from codebase_rag import context_slice as cs_mod - from codebase_rag.utils.source_extraction import extract_source_lines (tmp_path / "test_app.py").write_bytes( b"def test_run():\r\n assert run() == 6\r\n" @@ -134,26 +133,20 @@ def tests_reaching(_qn: str) -> list[dict[str, object]]: } ] - def _definition( - _fetch: object, _project: str, _qn: str, _root: Path | None - ) -> dict[str, object]: - return { - "qualified_name": "p.test_app.test_run", - "path": "test_app.py", - "start_line": 1, - "end_line": 2, - "source": extract_source_lines(tmp_path / "test_app.py", 1, 2), - } + @staticmethod + def span(_qn: str) -> tuple[int, int]: + # The reach index carries each test's span, so the source is + # read from disk without a definition query (issue #3243). + return 1, 2 monkeypatch.setattr(cs_mod, "ReachIndex", _Reach) - monkeypatch.setattr(cs_mod.graph_query, "definition", _definition) pieces = cs_mod._test_pieces(lambda *a, **k: [], "p", "p.app.run", tmp_path) assert len(pieces) == 1, repr(pieces) - assert pieces[0].source == "def test_run():\n assert run() == 6", repr( - pieces[0].source - ) + loader = pieces[0].loader + assert loader is not None + assert loader() == "def test_run():\n assert run() == 6", repr(loader()) # The indexer writes absolute_path through `cached_resolve_posix`, so the lookup diff --git a/codebase_rag/utils/source_extraction.py b/codebase_rag/utils/source_extraction.py index 8f7edfeb87..410dc6e80d 100644 --- a/codebase_rag/utils/source_extraction.py +++ b/codebase_rag/utils/source_extraction.py @@ -1,6 +1,6 @@ from __future__ import annotations -from collections.abc import Callable +from collections.abc import Callable, Sequence from pathlib import Path from loguru import logger @@ -10,6 +10,44 @@ from .source_encoding import decode_python_source +def read_source_lines(file_path: Path, encoding: str = ENCODING_UTF8) -> list[str]: + """The file's lines as the indexer read them, line endings kept.""" + raw_bytes = file_path.read_bytes() + # The indexer parsed a Python source in the encoding it declares, and + # the lines it recorded are lines of that text (issue #2445). + declared = ( + decode_python_source(raw_bytes, file_path) + if file_path.suffix in PY_EXTENSIONS + else None + ) + text = raw_bytes.decode(encoding) if declared is None else declared + return text.splitlines(keepends=True) + + +def slice_source_lines( + lines: Sequence[str], start_line: int, end_line: int, file_path: Path +) -> str | None: + """Lines `start_line`..`end_line` (1-based) of an already-read file.""" + if not lines: + return None + + if start_line > len(lines) or end_line > len(lines): + logger.warning( + ls.SOURCE_RANGE_EXCEEDS.format( + start=start_line, + end=end_line, + length=len(lines), + path=file_path, + ) + ) + end_line = min(end_line, len(lines)) + if start_line > len(lines): + return None + + extracted_lines = lines[start_line - 1 : end_line] + return "".join(extracted_lines).strip() + + def extract_source_lines( file_path: Path, start_line: int, end_line: int, encoding: str = ENCODING_UTF8 ) -> str | None: @@ -22,36 +60,9 @@ def extract_source_lines( return None try: - raw_bytes = file_path.read_bytes() - # The indexer parsed a Python source in the encoding it declares, and - # the lines it recorded are lines of that text (issue #2445). - declared = ( - decode_python_source(raw_bytes, file_path) - if file_path.suffix in PY_EXTENSIONS - else None + return slice_source_lines( + read_source_lines(file_path, encoding), start_line, end_line, file_path ) - text = raw_bytes.decode(encoding) if declared is None else declared - lines = text.splitlines(keepends=True) - - if not lines: - return None - - if start_line > len(lines) or end_line > len(lines): - logger.warning( - ls.SOURCE_RANGE_EXCEEDS.format( - start=start_line, - end=end_line, - length=len(lines), - path=file_path, - ) - ) - end_line = min(end_line, len(lines)) - if start_line > len(lines): - return None - - extracted_lines = lines[start_line - 1 : end_line] - return "".join(extracted_lines).strip() - except Exception as e: logger.warning(ls.SOURCE_EXTRACT_FAILED.format(path=file_path, error=e)) return None diff --git a/docs/architecture/context-slice.md b/docs/architecture/context-slice.md index 081e693343..e29bbf590e 100644 --- a/docs/architecture/context-slice.md +++ b/docs/architecture/context-slice.md @@ -46,11 +46,22 @@ Candidates are ordered by graph distance, then trace hotness (`dynamic_call_count` on the caller's edge, from `cgr trace ingest`), then embedding similarity for a free-text task, then name. Tokens are counted with the same encoder the query tools use, and pieces are taken in order -while they fit: the total never exceeds `budget_tokens`. The target is -special: when it alone exceeds the budget it is trimmed line by line to -fit and the slice is marked `truncated`; a budget too small for even its -first line names it in `omitted` instead. Every piece that did not fit is -listed in `omitted` with the reason it would have been included. +while they fit: the total never exceeds `budget_tokens`. A piece is +charged for what it prints, its source and the fields beside it (name, +file, span, reason), so a slice of one-line call sites cannot print many +times its budget in metadata (issue #3243). The target is special: when it +alone exceeds the budget it is trimmed to the most leading lines that fit +and the slice is marked `truncated`; a budget too small for even its first +line names it in `omitted` instead. `omitted_count` counts every +candidate that did not fit, and `omitted` names the best-ranked ten of +them with the reason each would have been included; beside the pieces, +the printed slice adds only those fields. + +A central symbol has thousands of callers and tests, nearly all of which a +budget leaves out. Their sources are read only for a candidate whose +fields still fit, each file the slice quotes is decoded once, and a test's +span comes from the reach index already built, never from a graph query +per test. ## Output @@ -59,14 +70,15 @@ listed in `omitted` with the reason it would have been included. "target": "myproj.pkg.util.helper", "resolved": "myproj.pkg.util.helper", "budget_tokens": 4000, - "used_tokens": 312, + "used_tokens": 398, "pieces": [ {"qualified_name": "myproj.pkg.util.helper", "file": "pkg/util.py", "span": [8, 13], - "why_included": "target", "source": "def helper(items: list[int]) -> Total:\n ...", "tokens": 61}, + "why_included": "target", "source": "def helper(items: list[int]) -> Total:\n ...", "tokens": 104}, {"qualified_name": "myproj.pkg.app.run", "file": "pkg/app.py", "span": [5, 5], - "why_included": "direct caller: the call line", "source": "return helper([1, 2]).value", "tokens": 9} + "why_included": "direct caller: the call line", "source": "return helper([1, 2]).value", "tokens": 52} ], "omitted": [], + "omitted_count": 0, "truncated": false } ``` From 7a6e9d09b87527ce5ce2305a6b6077afe1a5d996 Mon Sep 17 00:00:00 2001 From: vitali87 Date: Sat, 10 Oct 2026 01:44:58 +0000 Subject: [PATCH 2/6] refactor(context): size each candidate in _sized; collect omissions in _Omitted SonarCloud S3776: _fit was at 22 (15 allowed); it is now 13, with the per-candidate budget decision in _sized and the omitted sample/count in _Omitted. Behaviour is unchanged: checked against the previous _fit on 40k random candidate sets. Also splits a composite assertion (S9073). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NA8ea1a3KkMTaa9FAp5qve --- codebase_rag/context_slice.py | 87 ++++++++++++------- .../tests/test_context_budget_output.py | 3 +- 2 files changed, 56 insertions(+), 34 deletions(-) diff --git a/codebase_rag/context_slice.py b/codebase_rag/context_slice.py index 8d7f2b616a..969f682ff5 100644 --- a/codebase_rag/context_slice.py +++ b/codebase_rag/context_slice.py @@ -443,20 +443,27 @@ class _Fitted(NamedTuple): 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_count = 0 + omitted = _Omitted() used = 0 truncated = False seen: set[tuple[str, str, tuple[int, int]]] = set() - - def omit(candidate: _Candidate) -> None: - nonlocal omitted_count - omitted_count += 1 - if len(omitted) < cs.CONTEXT_OMITTED_SAMPLE: - omitted.append(f"{candidate.qualified_name} ({candidate.why})") - for candidate in sorted(candidates, key=lambda c: c.order): key = (candidate.qualified_name, candidate.why, candidate.span) if key in seen: @@ -464,32 +471,18 @@ def omit(candidate: _Candidate) -> None: seen.add(key) span = [candidate.span[0], candidate.span[1]] is_target = candidate.distance == 0 and not pieces - # 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 used + floor > budget and not is_target: - omit(candidate) + if (sized := _sized(candidate, span, budget - used, is_target)) is None: + omitted.add(candidate) continue - source = candidate.source or (candidate.loader() if candidate.loader else "") + 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 - tokens = _piece_tokens( - candidate.qualified_name, candidate.file, span, candidate.why, source - ) - if used + tokens > budget: - if not is_target: - omit(candidate) - continue - # 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, budget - used) - truncated = True - if not source: - omit(candidate) - continue used += tokens pieces.append( Piece( @@ -501,7 +494,35 @@ def omit(candidate: _Candidate) -> None: tokens=tokens, ) ) - return _Fitted(pieces, omitted, omitted_count, 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( diff --git a/codebase_rag/tests/test_context_budget_output.py b/codebase_rag/tests/test_context_budget_output.py index c62090cb5f..ff6f3dc719 100644 --- a/codebase_rag/tests/test_context_budget_output.py +++ b/codebase_rag/tests/test_context_budget_output.py @@ -237,7 +237,8 @@ def test_a_neighbourhood_that_fits_omits_nothing(temp_repo: Path) -> None: slice_ = context(store.fetch_all, PROJECT, f"{PROJECT}.app.core", 4000, temp_repo) - assert slice_["omitted"] == [] and slice_["omitted_count"] == 0 + assert slice_["omitted"] == [] + assert slice_["omitted_count"] == 0 assert {p["why_included"] for p in slice_["pieces"]} == { cs.CONTEXT_WHY_TARGET, cs.CONTEXT_WHY_CALLER, From d2d7786980cef4ac81441c5138e012f8b04924dc Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 18 Oct 2026 12:33:04 +0000 Subject: [PATCH 3/6] fix(context): charge a piece for the count it prints A piece was measured with `tokens: 0` and printed with its own count, so a count of four digits or more cost a token more than it was charged. Measure again with the count in place until it stops growing. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_013MgsnsA3jM14E4ai22xnhy --- codebase_rag/context_slice.py | 8 +++++++- codebase_rag/tests/test_context_budget_output.py | 16 ++++++++++++++++ 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/codebase_rag/context_slice.py b/codebase_rag/context_slice.py index 969f682ff5..38c0109970 100644 --- a/codebase_rag/context_slice.py +++ b/codebase_rag/context_slice.py @@ -432,7 +432,13 @@ def _piece_tokens( source=source, tokens=0, ) - return count_tokens(json.dumps(piece, indent=cs.MCP_JSON_INDENT)) + # 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. + printed = partial(json.dumps, piece, indent=cs.MCP_JSON_INDENT) + while (cost := count_tokens(printed())) > piece["tokens"]: + piece["tokens"] = cost + return piece["tokens"] class _Fitted(NamedTuple): diff --git a/codebase_rag/tests/test_context_budget_output.py b/codebase_rag/tests/test_context_budget_output.py index ff6f3dc719..8f0eb1131f 100644 --- a/codebase_rag/tests/test_context_budget_output.py +++ b/codebase_rag/tests/test_context_budget_output.py @@ -243,3 +243,19 @@ def test_a_neighbourhood_that_fits_omits_nothing(temp_repo: Path) -> None: cs.CONTEXT_WHY_TARGET, cs.CONTEXT_WHY_CALLER, } + + +def test_a_piece_is_charged_for_the_count_it_prints(temp_repo: Path) -> None: + # A piece was measured with `tokens: 0` and printed with its own count, + # and a count of four digits prints in more tokens than "0" does. + body = "class Big:\n" + "".join( + f" def m{i}(self, x):\n return x * {i}\n\n" for i in range(150) + ) + store = _repo(temp_repo, {"big.py": body}) + + slice_ = context(store.fetch_all, PROJECT, f"{PROJECT}.big.Big", 10_000, temp_repo) + + (target,) = slice_["pieces"] + assert target["tokens"] >= 1000 + printed = count_tokens(json.dumps(target, indent=cs.MCP_JSON_INDENT)) + assert printed == target["tokens"] From 9f9c64afd8665ae81f7884c1c01c3cd1413b4985 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 18 Oct 2026 12:33:04 +0000 Subject: [PATCH 4/6] fix(context): quote a file replaced under its old size and mtime afresh The decoded-lines cache was keyed by modification time and size, so an atomic replacement or edit that kept both returned the old lines. Key it by the inode and change time as well. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_013MgsnsA3jM14E4ai22xnhy --- codebase_rag/context_slice.py | 9 ++++- .../tests/test_context_budget_output.py | 38 +++++++++++++++++++ 2 files changed, 45 insertions(+), 2 deletions(-) diff --git a/codebase_rag/context_slice.py b/codebase_rag/context_slice.py index 38c0109970..3e358f7d0c 100644 --- a/codebase_rag/context_slice.py +++ b/codebase_rag/context_slice.py @@ -115,9 +115,11 @@ def _normalise(text: str) -> str: @lru_cache(maxsize=cs.CONTEXT_SOURCE_CACHE_FILES) -def _file_lines(path: Path, stamp: tuple[int, int]) -> tuple[str, ...]: +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)) @@ -132,7 +134,10 @@ def _lines(repo_root: Path | None, path: str | None, start: int, end: int) -> st return "" try: stat = candidate.stat() - lines = _file_lines(candidate, (stat.st_mtime_ns, stat.st_size)) + 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 "") diff --git a/codebase_rag/tests/test_context_budget_output.py b/codebase_rag/tests/test_context_budget_output.py index 8f0eb1131f..0b64e8ec01 100644 --- a/codebase_rag/tests/test_context_budget_output.py +++ b/codebase_rag/tests/test_context_budget_output.py @@ -11,6 +11,7 @@ from __future__ import annotations import json +import os import time from pathlib import Path @@ -259,3 +260,40 @@ def test_a_piece_is_charged_for_the_count_it_prints(temp_repo: Path) -> None: assert target["tokens"] >= 1000 printed = count_tokens(json.dumps(target, indent=cs.MCP_JSON_INDENT)) assert printed == target["tokens"] + + +def test_a_replaced_file_is_quoted_afresh(temp_repo: Path) -> None: + # An atomic replacement that keeps the old size and modification time + # is still a different file: its new lines are quoted, not the cached. + store = _repo( + temp_repo, + { + "app.py": "def core(x):\n return x + 1\n\n\ndef use():\n return core(1)\n", + }, + ) + app = temp_repo / "app.py" + before = context(store.fetch_all, PROJECT, f"{PROJECT}.app.core", 4000, temp_repo) + old = app.stat() + replacement = temp_repo / "app.py.new" + replacement.write_text( + "def core(x):\n return x + 1\n\n\ndef use():\n return core(2)\n", + encoding="utf-8", + ) + os.utime(replacement, ns=(old.st_atime_ns, old.st_mtime_ns)) + os.replace(replacement, app) + assert (app.stat().st_mtime_ns, app.stat().st_size) == ( + old.st_mtime_ns, + old.st_size, + ) + + after = context(store.fetch_all, PROJECT, f"{PROJECT}.app.core", 4000, temp_repo) + + def caller_sources(slice_: context_slice.ContextSlice) -> list[str]: + return [ + p["source"] + for p in slice_["pieces"] + if p["why_included"] == cs.CONTEXT_WHY_CALLER + ] + + assert caller_sources(before) == ["return core(1)"] + assert caller_sources(after) == ["return core(2)"] From 8e275feecedade0da5c798305e3c3969e51ee019 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 18 Oct 2026 17:19:27 +0000 Subject: [PATCH 5/6] fix(context): quote a reaching test from its first decorator The reach index span read start_line alone, the def line, so a decorated test lost its decorators where definition() had kept them (#2428). The dead-code node fetch now carries decorated_start_line and the span starts there. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_013MgsnsA3jM14E4ai22xnhy --- codebase_rag/cypher_queries.py | 1 + codebase_rag/dead_code.py | 3 + codebase_rag/graph_query.py | 10 +++- .../tests/test_context_budget_output.py | 56 +++++++++++++++++++ docs/architecture/context-slice.md | 3 +- 5 files changed, 69 insertions(+), 4 deletions(-) diff --git a/codebase_rag/cypher_queries.py b/codebase_rag/cypher_queries.py index 8a73f263a4..7c08549d39 100644 --- a/codebase_rag/cypher_queries.py +++ b/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, diff --git a/codebase_rag/dead_code.py b/codebase_rag/dead_code.py index d9ff02021b..a41c084b3a 100644 --- a/codebase_rag/dead_code.py +++ b/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). diff --git a/codebase_rag/graph_query.py b/codebase_rag/graph_query.py index 66b59230b4..8b49333f27 100644 --- a/codebase_rag/graph_query.py +++ b/codebase_rag/graph_query.py @@ -559,7 +559,9 @@ def source_root_for( return local if Path(stored).resolve() == local else None -def _decorated_start(row: ResultRow, own_start: int | None) -> int | None: +def _decorated_start( + row: ResultRow | PropertyDict, own_start: int | None +) -> int | None: """The first decorator's line when the node records one above its own start, else that start (issue #2428). @@ -1172,12 +1174,14 @@ def _walk(self, qualified_name: str) -> tuple[dict[str, int], dict[str, str]]: return depth_of, through_of def span(self, qualified_name: str) -> tuple[int, int]: - """The recorded (start, end) lines of a node, 0 where unknown.""" + """The recorded (start, end) lines of a node, 0 where unknown; the + start is the first decorator's line, as `definition()` answers.""" entry = self._by_qn.get(qualified_name) if entry is None: return 0, 0 props = entry[1] - start, end = props.get(cs.KEY_START_LINE), props.get(cs.KEY_END_LINE) + own, end = props.get(cs.KEY_START_LINE), props.get(cs.KEY_END_LINE) + start = _decorated_start(props, own if isinstance(own, int) else None) return ( start if isinstance(start, int) else 0, end if isinstance(end, int) else 0, diff --git a/codebase_rag/tests/test_context_budget_output.py b/codebase_rag/tests/test_context_budget_output.py index 0b64e8ec01..2acdee1bf6 100644 --- a/codebase_rag/tests/test_context_budget_output.py +++ b/codebase_rag/tests/test_context_budget_output.py @@ -12,6 +12,7 @@ import json import os +import re import time from pathlib import Path @@ -19,10 +20,12 @@ from codebase_rag import constants as cs from codebase_rag import context_slice +from codebase_rag import cypher_queries as cq from codebase_rag.context_slice import context from codebase_rag.graph_query import DefinitionRow, QueryFn from codebase_rag.graph_updater import GraphUpdater from codebase_rag.parser_loader import load_parsers +from codebase_rag.types_defs import PropertyParams, ResultRow from codebase_rag.utils.token_utils import count_tokens from evals.cgr_graph import _StatefulIngestor @@ -297,3 +300,56 @@ def caller_sources(slice_: context_slice.ContextSlice) -> list[str]: assert caller_sources(before) == ["return core(1)"] assert caller_sources(after) == ["return core(2)"] + + +_COLUMN = re.compile(r"\bn\.(\w+) AS (\w+)") + + +def _as_memgraph_answers(store: _StatefulIngestor) -> QueryFn: + # The stand-in answers the dead-code node query with a fixed key set; + # Memgraph answers with every column its RETURN names. + def fetch(query: str, params: PropertyParams | None) -> list[ResultRow]: + rows = store.fetch_all(query, dict(params) if params is not None else None) + if query != cq.CYPHER_DEAD_CODE_NODES: + return rows + columns = _COLUMN.findall(query) + out: list[ResultRow] = [] + for row in rows: + node = (str(row[cs.KEY_LABEL]), str(row[cs.KEY_QUALIFIED_NAME])) + props = store.nodes[node] + out.append( + {cs.KEY_LABEL: row[cs.KEY_LABEL]} + | {alias: props.get(prop) for prop, alias in columns} + ) + return out + + return fetch + + +def test_a_decorated_test_is_quoted_with_its_decorators(temp_repo: Path) -> None: + # The reach index read the def line alone, so a decorated test lost its + # decorators; it is quoted from the first one, as `definition()` quotes + # a decorated def (issue #2428). + store = _repo( + temp_repo, + { + "app.py": "def core(x):\n return x + 1\n", + "tests/__init__.py": "", + "tests/test_app.py": ( + "import pytest\n\nfrom app import core\n\n\n" + "@pytest.mark.slow\n" + "def test_core():\n assert core(1) == 2\n" + ), + }, + ) + + slice_ = context( + _as_memgraph_answers(store), PROJECT, f"{PROJECT}.app.core", 4000, temp_repo + ) + + why = cs.CONTEXT_WHY_TEST.format(depth=1, through=f"{PROJECT}.app.core") + (test,) = [p for p in slice_["pieces"] if p["why_included"] == why] + assert test["span"] == [6, 8] + assert test["source"] == ( + "@pytest.mark.slow\ndef test_core():\n assert core(1) == 2" + ) diff --git a/docs/architecture/context-slice.md b/docs/architecture/context-slice.md index e29bbf590e..a4b017b3f7 100644 --- a/docs/architecture/context-slice.md +++ b/docs/architecture/context-slice.md @@ -61,7 +61,8 @@ A central symbol has thousands of callers and tests, nearly all of which a budget leaves out. Their sources are read only for a candidate whose fields still fit, each file the slice quotes is decoded once, and a test's span comes from the reach index already built, never from a graph query -per test. +per test. That span starts at a decorated test's first decorator, as +`definition` quotes it. ## Output From d512a17100f9be927ac0de83881b5cf61edd011d Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 20 Oct 2026 21:18:00 +0000 Subject: [PATCH 6/6] fix(context): charge each piece at the indentation the slice prints it _piece_tokens measured a piece as a top-level JSON object, but the slice prints it two levels deep, every line indented by four more spaces. The gap grew with the piece count: a 10,000 budget over 300 reaching tests printed 10,478 tokens, past the budget plus its 400-token envelope. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_013MgsnsA3jM14E4ai22xnhy --- codebase_rag/constants/cli.py | 3 +++ codebase_rag/context_slice.py | 13 ++++++++++--- codebase_rag/tests/test_context_budget_output.py | 12 +++++++++--- docs/architecture/context-slice.md | 5 +++-- 4 files changed, 25 insertions(+), 8 deletions(-) diff --git a/codebase_rag/constants/cli.py b/codebase_rag/constants/cli.py index b224aa3575..a56b2cbdd5 100644 --- a/codebase_rag/constants/cli.py +++ b/codebase_rag/constants/cli.py @@ -1085,6 +1085,9 @@ class DiffMarker: # 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" diff --git a/codebase_rag/context_slice.py b/codebase_rag/context_slice.py index 3e358f7d0c..29843944e5 100644 --- a/codebase_rag/context_slice.py +++ b/codebase_rag/context_slice.py @@ -18,6 +18,7 @@ import json import re +import textwrap from collections.abc import Callable, Iterable from functools import lru_cache, partial from pathlib import Path @@ -439,9 +440,15 @@ def _piece_tokens( ) # 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. - printed = partial(json.dumps, piece, indent=cs.MCP_JSON_INDENT) - while (cost := count_tokens(printed())) > piece["tokens"]: + # 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"] diff --git a/codebase_rag/tests/test_context_budget_output.py b/codebase_rag/tests/test_context_budget_output.py index 2acdee1bf6..f835f8eac9 100644 --- a/codebase_rag/tests/test_context_budget_output.py +++ b/codebase_rag/tests/test_context_budget_output.py @@ -13,6 +13,7 @@ import json import os import re +import textwrap import time from pathlib import Path @@ -68,7 +69,7 @@ def _printed(slice_: context_slice.ContextSlice) -> int: return count_tokens(json.dumps(slice_, indent=cs.MCP_JSON_INDENT)) -@pytest.mark.parametrize("budget", [4000, 1000, 300]) +@pytest.mark.parametrize("budget", [20_000, 10_000, 4000, 1000, 300]) def test_the_printed_slice_fits_the_budget(temp_repo: Path, budget: int) -> None: store = _central(temp_repo) @@ -261,8 +262,13 @@ def test_a_piece_is_charged_for_the_count_it_prints(temp_repo: Path) -> None: (target,) = slice_["pieces"] assert target["tokens"] >= 1000 - printed = count_tokens(json.dumps(target, indent=cs.MCP_JSON_INDENT)) - assert printed == target["tokens"] + # As the slice prints it: two levels deep, inside the slice and its + # `pieces` list, every line indented by both levels. + printed = textwrap.indent( + json.dumps(target, indent=cs.MCP_JSON_INDENT), " " * 2 * cs.MCP_JSON_INDENT + ) + assert printed in json.dumps(slice_, indent=cs.MCP_JSON_INDENT) + assert count_tokens(printed) == target["tokens"] def test_a_replaced_file_is_quoted_afresh(temp_repo: Path) -> None: diff --git a/docs/architecture/context-slice.md b/docs/architecture/context-slice.md index a4b017b3f7..c115f3b051 100644 --- a/docs/architecture/context-slice.md +++ b/docs/architecture/context-slice.md @@ -48,8 +48,9 @@ embedding similarity for a free-text task, then name. Tokens are counted with the same encoder the query tools use, and pieces are taken in order while they fit: the total never exceeds `budget_tokens`. A piece is charged for what it prints, its source and the fields beside it (name, -file, span, reason), so a slice of one-line call sites cannot print many -times its budget in metadata (issue #3243). The target is special: when it +file, span, reason) at the indentation the slice prints them with, so a +slice of one-line call sites cannot print many times its budget in +metadata (issue #3243). The target is special: when it alone exceeds the budget it is trimmed to the most leading lines that fit and the slice is marked `truncated`; a budget too small for even its first line names it in `omitted` instead. `omitted_count` counts every