From 059f59e4dc88f564c67c7589aa5e3a5cfb94eea5 Mon Sep 17 00:00:00 2001 From: Tanner Wendland Date: Fri, 24 Jul 2026 12:46:14 -0700 Subject: [PATCH 1/6] config: add mcp_max_response_bytes byte-budget setting (Step 1) Byte ceiling on serialized MCP tool responses, default 100_000 (~25k tokens), following the existing CODE_SEARCH_ env-prefixed tunable pattern. Foundation for the MCP response-shaping work; not yet wired into app/main.py. Co-authored-by: Isaac --- app/config.py | 6 ++++++ tests/unit/test_main.py | 11 +++++++++++ 2 files changed, 17 insertions(+) diff --git a/app/config.py b/app/config.py index a26bfbb..e9900f8 100644 --- a/app/config.py +++ b/app/config.py @@ -53,6 +53,12 @@ class Settings(BaseSettings): row_limit: int = 200 max_row_limit: int = 1000 + # Byte ceiling on every serialized MCP tool response (~4 bytes/token heuristic: + # 100_000 bytes ≈ 25k tokens). A per-request `max_bytes` tool param clamps this DOWN, + # never up (app/main.py's `_effective_budget`). No tokenizer dependency -- the budget is + # enforced against the exact `json.dumps` wire string, not an estimate. + mcp_max_response_bytes: int = 100_000 + # Gates the semantic_search code path. Default True: the target Lakebase project's # managed shared_preload_libraries including lakebase_vector,lakebase_text is a stated # project assumption (see docs/runbooks/semantic-enablement.md). Opt out with diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index e396d6f..8212cd4 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -1291,6 +1291,17 @@ def test_match_budget_ms_defaults_to_2000_and_is_env_overridable( assert Settings(lakebase_endpoint=None).match_budget_ms == 500 +@pytest.mark.unit +def test_mcp_max_response_bytes_defaults_to_100000_and_is_env_overridable( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("CODE_SEARCH_MCP_MAX_RESPONSE_BYTES", raising=False) + assert Settings(lakebase_endpoint=None).mcp_max_response_bytes == 100_000 + + monkeypatch.setenv("CODE_SEARCH_MCP_MAX_RESPONSE_BYTES", "5000") + assert Settings(lakebase_endpoint=None).mcp_max_response_bytes == 5000 + + @pytest.mark.observability def test_signals_log_includes_truncation_reason() -> None: # Which cap/budget tripped -- byte_cap/row_cap/match_budget -- must be recoverable from the From 5fed5ee5d2ee20dd17aa69d60f027004dd0aca21 Mon Sep 17 00:00:00 2001 From: Tanner Wendland Date: Fri, 24 Jul 2026 12:51:46 -0700 Subject: [PATCH 2/6] main: MCP response shaping -- projection, byte budget, dispatch telemetry (Steps 2-6) Implements the byte-budgeted MCP response pipeline entirely in app/main.py per the consensus plan (.omc/plans/ralplan-token-limiter.md): - project_for_mcp() drops webui-only fields (duration_ns, content_sha, byte_ranges) from the MCP wire payload; service.py/webui builders are untouched. - _effective_budget/_fit_list/_shape_response implement the serialize -> measure -> tail-trim -> re-verify pipeline against the exact json.dumps wire string, never an estimate. - Per-tool truncators: static tail-trim table for list_repos/ find_references/list_imports/semantic_search; a search_code truncator that synthesizes next_cursor via a bounded repo name->id SELECT, degrading to a flagged handle-less truncation on a no-row/fault; a dedicated get_file line-paging shaper (start_line/next_start_line, split("\n") congruent with grep.py:436). - _dispatch gains max_bytes threading, response_bytes_pre/response_bytes telemetry, and duration_ns/cursor_invalid in _signals(). - search_code always runs in pagination mode (cursor param, CursorError caught into a structured cursor_invalid payload); get_file gains start_line; all six tools gain max_bytes. Migrated the two search_code wrapper tests whose fakes needed the new cursor kwarg; added response_bytes_pre/response_bytes and duration_ns-before-projection assertions to the dispatch log tests. Per-tool truncator/shaping unit tests land in a follow-up commit (tests/unit/test_mcp_shaping.py). Co-authored-by: Isaac --- app/main.py | 599 ++++++++++++++++++++++++++++++++++++++-- tests/unit/test_main.py | 26 +- 2 files changed, 600 insertions(+), 25 deletions(-) diff --git a/app/main.py b/app/main.py index 98cf415..d7d5e13 100644 --- a/app/main.py +++ b/app/main.py @@ -38,20 +38,23 @@ import logging import threading import time -from collections.abc import AsyncIterator, Callable +from collections.abc import AsyncIterator, Callable, Sequence from contextlib import asynccontextmanager from typing import Any import anyio from mcp.server.fastmcp import Context, FastMCP +from sqlalchemy import select from sqlalchemy.engine import Engine from starlette.applications import Starlette from starlette.requests import Request from starlette.responses import JSONResponse from app import service -from app.config import get_settings +from app.config import Settings, get_settings from app.db.client import create_db_engine +from app.db.models import Repo +from app.search.grep import FileCursor from app.search.semantic import _semantic_search_payload logger = logging.getLogger("app.tools") @@ -133,33 +136,475 @@ def _signals(payload: dict[str, Any]) -> dict[str, Any]: "unsupported_direction": payload.get("unsupported_direction"), "missing_repo": payload.get("missing_repo"), "missing_target": payload.get("missing_target"), + # search_code's duration_ns is read HERE, before MCP projection drops the field from + # the wire payload (see project_for_mcp below) -- so duration observability survives + # the byte-budget work even though the field itself never reaches an MCP caller. + "duration_ns": payload.get("duration_ns"), + # Set only by the search_code tool wrapper's structured CursorError catch (a garbled + # `cursor` string) -- None-safe on every other payload shape. + "cursor_invalid": payload.get("cursor_invalid"), } -async def _dispatch(name: str, build: Callable[[], dict[str, Any]]) -> str: - """Run a tool's blocking ``build`` off-loop, log its outcome, and serialize to ``str``. +# ---------------------------------------------------------------- MCP response shaping +# +# Everything in this section runs ONLY on the MCP path -- app/service.py's payload builders +# and webui/ never see it (Round-1 constraint: one builder, byte-identical for both +# consumers). The pipeline per call: capture pre-projection signals -> drop webui-only +# fields (project_for_mcp) -> serialize once -> if the wire string exceeds the effective +# byte budget, per-tool tail-trim with measured re-serialization -> flag +# `truncated`/`truncation_reason="token_budget"` (extending the repo's existing +# byte_cap/row_cap/match_budget reason enum) -> attach a resume handle where plumbing exists +# (`next_cursor` for search_code, `next_start_line` for get_file). Truncation is always a +# payload fact, never an exception -- see the module docstring's recoverable-conditions +# contract; the same "never hard-error" guarantee extends to the byte budget. +# +# The budget is enforced against the EXACT wire string: `json.dumps` with its DEFAULT +# separators (this module's existing convention), never an estimate -- so a compact-separator +# switch later would only loosen the effective budget, never violate it. + +# Floor so a hostile/typo'd `max_bytes=1` request param cannot force a truncator into a +# degenerate (or infinite) trim loop; still far below the 100_000 default. +_MIN_MAX_BYTES = 1024 + + +def _effective_budget(max_bytes: int | None, cfg: Settings) -> int: + """Resolve the byte budget for one call: `max_bytes` clamps the env default DOWN, never up. + + `None` or a non-positive value means "no per-request override" -> the env-configured + ceiling. Otherwise clamped to `[_MIN_MAX_BYTES, cfg.mcp_max_response_bytes]` (AC6). + """ + if max_bytes is None or max_bytes <= 0: + return cfg.mcp_max_response_bytes + return max(_MIN_MAX_BYTES, min(max_bytes, cfg.mcp_max_response_bytes)) + + +def project_for_mcp(tool: str, payload: dict[str, Any]) -> None: + """Mutate ``payload`` in place, dropping webui-only fields from the MCP wire response. + + Table-driven (a future drop is a one-line addition) and None-safe on skeletal/minimal + payloads (every access is a defensive ``.get`` that no-ops when the key/list is absent) -- + the wrapper tests' fakes return bare dicts like ``{"query": q}``. Today only + ``search_code`` carries any of the three dropped fields (verified against every other + builder's shape): ``duration_ns`` (rendered nowhere; ``_signals()`` already read it before + this runs), per-file ``content_sha`` (near-zero LLM value; the cursor keeps its own copy + inside the opaque token), and per-match ``byte_ranges`` (webui highlighting only, derivable + from ``text``). ``permalink_branch``/``commit``/ranking metadata/``rrf_score``/``similarity`` + are deliberately KEPT (spec constraint) -- this function only ever removes the three named + fields, nothing else. + """ + payload.pop("duration_ns", None) + if tool != "search_code": + return + for file_entry in payload.get("files") or []: + file_entry.pop("content_sha", None) + for match in file_entry.get("matches") or []: + match.pop("byte_ranges", None) - The single choke-point every tool routes through: it times the call, logs the signal set - and limiter/pool saturation on success, and on an UNEXPECTED fault logs the full traceback - (``logger.exception``) then re-raises — the fault is never swallowed. Recoverable - conditions are turned into payload fields inside ``build``, so only genuine faults land - here. + +def _fit_list(items: Sequence[Any], budget_for_items: int) -> int: + """Return the largest prefix length of ``items`` whose combined serialized size fits. + + Prefix-sum over ``len(json.dumps(item)) + 2`` (a separator allowance for the ``", "`` + joining each item in the enclosing JSON array) against ``budget_for_items``. Exact against + default-separator ``json.dumps``, not an estimate. + """ + used = 0 + for index, item in enumerate(items): + cost = len(json.dumps(item)) + 2 + if used + cost > budget_for_items: + return index + used += cost + return len(items) + + +def _make_tail_trim_truncator( + list_key: str, recompute: Callable[[dict[str, Any]], None] +) -> Callable[[dict[str, Any], int], None]: + """Build a truncator that tail-trims ``payload[list_key]`` to fit ``budget``. + + Shared by the four tools with no resume-handle plumbing (``list_repos``, + ``find_references``, ``list_imports``, ``semantic_search``): drop the tail of the + dominant list (relevance-ranked for ``semantic_search``, so the tail is the least + relevant), recompute the list's derived count/summary fields, and flag + ``truncated``/``truncation_reason``. None-safe: an absent or empty list still flags + (lossy, no resume handle for these four) rather than raising. + """ + + def _truncate(payload: dict[str, Any], budget: int) -> None: + payload.setdefault("truncated", False) + payload.setdefault("truncation_reason", None) + items: list[Any] = payload.get(list_key) or [] + if not items: + payload["truncated"] = True + if payload.get("truncation_reason") is None: + payload["truncation_reason"] = "token_budget" + return + + envelope_overhead = len(json.dumps(payload)) - sum( + len(json.dumps(item)) + 2 for item in items + ) + budget_for_items = max(0, budget - envelope_overhead) + keep = _fit_list(items, budget_for_items) + payload[list_key] = items[:keep] + recompute(payload) + payload["truncated"] = True + if payload.get("truncation_reason") is None: + payload["truncation_reason"] = "token_budget" + + # Safety net: the per-item "+2" separator allowance is a close but not + # byte-for-byte-guaranteed estimate once `recompute` changes derived scalars (e.g. a + # count's digit width). Re-verify against the real wire string and keep halving until + # it fits or the list is empty -- bounded (log2(len(items)) iterations), never infinite. + while keep > 0 and len(json.dumps(payload)) > budget: + keep //= 2 + payload[list_key] = items[:keep] + recompute(payload) + + return _truncate + + +def _recompute_list_repos(payload: dict[str, Any]) -> None: + payload["count"] = len(payload.get("repos") or []) + + +def _recompute_reference_sites(payload: dict[str, Any]) -> None: + """Recompute ``site_count``/``resolution_summary`` from the surviving (trimmed) sites -- + both currently derive from the returned list, so recomputing keeps them internally + consistent after a tail-trim.""" + sites: list[dict[str, Any]] = payload.get("sites") or [] + payload["site_count"] = len(sites) + summary = {"unique": 0, "ambiguous": 0, "unresolved": 0} + for site in sites: + resolution = site.get("resolution") + if resolution in summary: + summary[resolution] += 1 + payload["resolution_summary"] = summary + + +def _recompute_semantic_results(payload: dict[str, Any]) -> None: + payload["count"] = len(payload.get("results") or []) + + +def _resolve_repo_id(engine: Engine, cfg: Settings, name: str) -> int | None: + """One bounded ``repos.name -> id`` SELECT, fired only when byte-budget truncation of + ``search_code`` needs to synthesize a resume cursor (D3) -- the built payload resolves + ``repo_id`` to a name and drops the id, so the MCP layer must resolve it back. + + A no-row result (reachable: the builder itself falls back to ``str(repo_id)`` when a repo + id has no name, ``service.py:764``, so the payload's ``repo`` string may match no + ``repos.name``) or any unexpected fault returns ``None`` -- the caller then leaves the + response flagged ``truncated``/``"token_budget"`` with no cursor (an honest lossy + truncation), never raising through ``_dispatch``. + """ + try: + with engine.connect() as conn: + with conn.begin(): + conn.exec_driver_sql( + f"SET LOCAL statement_timeout = {int(cfg.statement_timeout_ms)}" + ) + return conn.execute(select(Repo.id).where(Repo.name == name)).scalar_one_or_none() + except Exception: + logger.warning("search_code cursor resolve failed for repo=%r", name, exc_info=True) + return None + + +def _make_search_code_truncator( + engine: Engine, cfg: Settings +) -> Callable[[dict[str, Any], int, list[tuple[str, int]] | None], None]: + """Build the ``search_code`` truncator: a closure over ``engine``/``cfg`` so an actual + byte-budget trim can synthesize ``next_cursor`` (D3). + + Pure tail-trim of ``files`` in the payload's existing ``(repo_id, path, content_sha)`` sort + order (never reordered to protect symbol-bearing files) -- the kept files are always a + maximal contiguous prefix of that order, so the last kept file is a valid resume point: + every candidate after it (kept or dropped) sorts strictly later, exactly the invariant + ``search_code``'s cursor-seek predicate resumes on. ``snapshot`` is + ``[(content_sha, span_count), ...]`` captured by ``_shape_response`` BEFORE + ``project_for_mcp`` strips ``content_sha``/``byte_ranges`` from ``payload["files"]`` -- + trimming here is positional against that snapshot, not content-aware. + """ + + def _truncate( + payload: dict[str, Any], budget: int, snapshot: list[tuple[str, int]] | None + ) -> None: + files: list[dict[str, Any]] = payload.get("files") or [] + snapshot = snapshot or [] + if not files: + payload["truncated"] = True + if payload.get("truncation_reason") is None: + payload["truncation_reason"] = "token_budget" + payload["next_cursor"] = None + return + + def _apply(keep: int) -> None: + kept = files[:keep] + payload["files"] = kept + payload["file_count"] = len(kept) + payload["match_count"] = sum(count for _sha, count in snapshot[:keep]) + + envelope_overhead = len(json.dumps(payload)) - sum(len(json.dumps(f)) + 2 for f in files) + budget_for_items = max(0, budget - envelope_overhead) + keep = _fit_list(files, budget_for_items) + _apply(keep) + payload["truncated"] = True + if payload.get("truncation_reason") is None: + payload["truncation_reason"] = "token_budget" + + while keep > 0 and len(json.dumps(payload)) > budget: + keep //= 2 + _apply(keep) + + next_cursor: str | None = None + if keep > 0: + last_file = payload["files"][keep - 1] + last_sha = snapshot[keep - 1][0] if keep - 1 < len(snapshot) else None + repo_name = last_file.get("repo") + path = last_file.get("file") + if repo_name is not None and path is not None and last_sha: + repo_id = _resolve_repo_id(engine, cfg, repo_name) + if repo_id is not None: + next_cursor = service.encode_cursor( + FileCursor(repo_id=repo_id, path=path, content_sha=last_sha) + ) + payload["next_cursor"] = next_cursor + + return _truncate + + +# Static table for the four tools with no resume-handle plumbing. search_code and get_file +# are NOT here: search_code's truncator must close over engine/cfg (the cursor-synthesis +# SELECT), and get_file's shaping is not a "trim a list" operation at all (see +# _shape_get_file_response) -- both are wired per-request by their own tool wrapper instead. +_TRUNCATORS: dict[str, Callable[[dict[str, Any], int], None]] = { + "list_repos": _make_tail_trim_truncator("repos", _recompute_list_repos), + "find_references": _make_tail_trim_truncator("sites", _recompute_reference_sites), + "list_imports": _make_tail_trim_truncator("sites", _recompute_reference_sites), + "semantic_search": _make_tail_trim_truncator("results", _recompute_semantic_results), +} + + +def _shape_response( + tool: str, + payload: dict[str, Any], + budget: int, + truncator: Callable[..., None] | None, +) -> tuple[str, dict[str, Any]]: + """Project, measure, and (if needed) truncate ``payload`` for the MCP wire. + + Returns ``(body, log_fields)`` where ``log_fields`` carries ``pre_bytes``/``post_bytes`` + (pre- and post-projection serialized sizes, for ``_dispatch``'s telemetry line) and the + pre-projection ``signals`` dict (so ``duration_ns`` observability survives projection + dropping the field). Never raises: an irreducible over-budget envelope (a giant echoed + scalar with nothing left to trim) is returned flagged rather than erroring (Principle 4). + """ + pre_bytes = len(json.dumps(payload)) + signals = _signals(payload) + + # search_code only: snapshot (content_sha, span_count) per file BEFORE projection strips + # content_sha/byte_ranges -- the truncator needs both to synthesize a resume cursor and to + # recompute match_count over the surviving tail (symbol matches count 1, grep matches count + # len(byte_ranges), exactly mirroring service.py's own match_count arithmetic). + snapshot: list[tuple[str, int]] | None = None + if tool == "search_code": + snapshot = [ + ( + entry.get("content_sha") or "", + sum( + len(match.get("byte_ranges") or ()) or (1 if match.get("symbols") else 0) + for match in (entry.get("matches") or []) + ), + ) + for entry in (payload.get("files") or []) + ] + + project_for_mcp(tool, payload) + body = json.dumps(payload) + if len(body) <= budget or truncator is None: + return body, {"pre_bytes": pre_bytes, "post_bytes": len(body), "signals": signals} + + if tool == "search_code": + truncator(payload, budget, snapshot) + else: + truncator(payload, budget) + + body = json.dumps(payload) + return body, {"pre_bytes": pre_bytes, "post_bytes": len(body), "signals": signals} + + +def _line_cost(line: str) -> int: + """The exact JSON-escaped body-byte cost of ``line`` (excluding the two quote chars + ``json.dumps`` adds around any standalone string) -- JSON string escaping has no + cross-character interaction, so this is exact when lines are later joined, not approximate. + """ + return len(json.dumps(line)) - 2 + + +def _fit_lines(lines: Sequence[str], budget_for_content: int) -> int: + """Largest prefix of ``lines`` whose ``"\\n"``-rejoined content fits ``budget_for_content`` + JSON-string-body bytes. Each line after the first adds 2 bytes for the re-appended + ``"\\n"`` (which ``json.dumps`` encodes as the two characters ``\\n``) that ``split("\\n")`` + stripped and a per-line ``json.dumps`` would not otherwise capture. + """ + used = 0 + for index, line in enumerate(lines): + cost = _line_cost(line) + (2 if index > 0 else 0) + if used + cost > budget_for_content: + return index + used += cost + return len(lines) + + +def _shape_get_file_response( + payload: dict[str, Any], start_line: int, budget: int +) -> tuple[str, dict[str, Any]]: + """Slice ``payload["content"]`` to a ``start_line``-anchored page fitting ``budget``. + + Per D2: splits with ``content.split("\\n")`` -- the SAME rule ``grep.py:436`` uses for + ``search_code`` line numbers, so ``get_file`` pages stay congruent with search match lines + on every input, including form feeds and ``U+2028``/``U+2029`` (which ``str.splitlines()`` + would wrongly treat as line breaks for this purpose). Reassembly re-appends ``"\\n"`` to + every segment except the last, which is byte-exact for CRLF (the ``\\r`` stays attached to + its own segment) and no-trailing-newline files alike: pages joined with ``"\\n"`` across + page boundaries reconstruct the exact original content. + + A miss (``found: false``) never truncates: ``start_line`` echoes, ``next_start_line`` is + ``null``, ``truncated`` is ``False``. On a hit, at least one line is always returned even if + it alone exceeds ``budget`` (Principle 4's progress guarantee outranks strict enforcement + for that documented, degenerate edge -- a single line larger than the whole budget); + ``next_start_line`` still advances past it. + """ + pre_bytes = len(json.dumps(payload)) + signals = _signals(payload) + project_for_mcp("get_file", payload) + start_line = max(1, start_line) + + if not payload.get("found"): + payload["start_line"] = start_line + payload["next_start_line"] = None + payload["truncated"] = False + payload["truncation_reason"] = None + body = json.dumps(payload) + return body, {"pre_bytes": pre_bytes, "post_bytes": len(body), "signals": signals} + + content = payload.get("content") or "" + lines = content.split("\n") # grep.py:436 congruence -- NOT str.splitlines() + total_lines = len(lines) + start_idx = min(start_line - 1, total_lines) + page_lines = lines[start_idx:] + + payload["start_line"] = start_line + payload["content"] = "\n".join(page_lines) + payload["truncated"] = False + payload["truncation_reason"] = None + payload["next_start_line"] = None + body = json.dumps(payload) + if len(body) <= budget: + return body, {"pre_bytes": pre_bytes, "post_bytes": len(body), "signals": signals} + + # envelope_overhead = everything except the content string's own escaped-body bytes + # (excluding its two outer quotes) -- the exact per-line cost unit _fit_lines uses. + content_cost = len(json.dumps(payload["content"])) - 2 + envelope_overhead = len(body) - content_cost + budget_for_content = max(0, budget - envelope_overhead) + keep = _fit_lines(page_lines, budget_for_content) + if keep == 0 and page_lines: + keep = 1 # progress guarantee: always return >= 1 line, even oversized (documented edge) + + def _apply(k: int) -> None: + kept_lines = page_lines[:k] + payload["content"] = "\n".join(kept_lines) + next_idx = start_idx + k + payload["next_start_line"] = (next_idx + 1) if next_idx < total_lines else None + + _apply(keep) + payload["truncated"] = True + payload["truncation_reason"] = "token_budget" + body = json.dumps(payload) + while keep > 1 and len(body) > budget: + keep -= 1 + _apply(keep) + body = json.dumps(payload) + return body, {"pre_bytes": pre_bytes, "post_bytes": len(body), "signals": signals} + + +def _cursor_invalid_payload(query: str, error: Exception) -> dict[str, Any]: + """A garbled/tampered/version-mismatched ``cursor`` string: a structured, remedy-bearing + rejection -- the repo's recoverable-payload idiom (mirrors ``unsupported_filter``) -- never + an uncaught :class:`~app.service.CursorError` through ``_dispatch``. Carries the full + pinned empty envelope so a caller's existing key access never KeyErrors. + """ + return { + "query": query, + "file_count": 0, + "match_count": 0, + "duration_ns": 0, + "files": [], + "truncated": False, + "truncation_reason": None, + "regex_incompatible": False, + "regex_invalid": None, + "query_too_broad": False, + "query_parse_error": None, + "no_content_atom": False, + "zero_width_only_atoms": False, + "next_cursor": None, + "cursor_invalid": True, + "reason": f"invalid cursor: {error}", + } + + +async def _dispatch( + name: str, + build: Callable[[], dict[str, Any]], + *, + max_bytes: int | None = None, + shape: Callable[[dict[str, Any], int], tuple[str, dict[str, Any]]] | None = None, +) -> str: + """Run a tool's blocking ``build`` off-loop, shape the result to the byte budget, log the + outcome, and return the serialized wire string. + + The single choke-point every tool routes through: it times the call, logs the signal set, + the pre-/post-shaping response byte sizes, and limiter/pool saturation on success, and on + an UNEXPECTED fault logs the full traceback (``logger.exception``) then re-raises — the + fault is never swallowed. Recoverable conditions are turned into payload fields inside + ``build``, so only genuine faults land here. + + ``shape`` defaults to :func:`_shape_response` bound to ``name``'s static truncator (the + four no-resume-handle tools); ``search_code``/``get_file`` pass their own ``shape`` closure + (over ``engine``/``cfg``/``start_line``) since their shaping needs request-scoped state a + static per-name lookup cannot carry. Both ``build`` and ``shape`` run INSIDE the same + worker-thread call (``_run_blocking``), so the multi-MB ``json.dumps`` calls and any + truncation-time DB round trip never touch the event loop. """ t0 = time.monotonic() + cfg = get_settings() + budget = _effective_budget(max_bytes, cfg) + shape_fn = shape or ( + lambda payload, b: _shape_response(name, payload, b, _TRUNCATORS.get(name)) + ) + + def _run() -> tuple[str, dict[str, Any]]: + payload = build() + return shape_fn(payload, budget) + try: - payload = await _run_blocking(build) + body, log_fields = await _run_blocking(_run) except Exception: logger.exception("tool=%s failed", name) raise logger.info( - "tool=%s duration_ms=%.1f signals=%s limiter_borrowed=%d/%d", + "tool=%s duration_ms=%.1f response_bytes_pre=%d response_bytes=%d signals=%s " + "limiter_borrowed=%d/%d", name, (time.monotonic() - t0) * 1e3, - _signals(payload), + log_fields["pre_bytes"], + log_fields["post_bytes"], + log_fields["signals"], _DB_LIMITER.borrowed_tokens, _DB_LIMITER.total_tokens, ) - return json.dumps(payload) + return body # ------------------------------------------------------------------------ payload builders @@ -228,6 +673,8 @@ async def search_code( limit: int = 200, branch: str | None = None, commit: str | None = None, + cursor: str | None = None, + max_bytes: int | None = None, ) -> str: """Search the indexed corpus with a zoekt-style query; returns file-grouped line matches. @@ -266,6 +713,26 @@ async def search_code( ``/regex/``, ``repo:``, ``file:``, or ``sym:`` pattern is not a valid Postgres POSIX ARE (e.g. ``/[/``) -- distinct from ``regex_incompatible``, which means Python ``regex`` (not Postgres) rejected an otherwise-valid pattern and only degrades highlighting. + + This tool always runs in pagination mode: page 1 omits ``cursor`` (or passes ``null``), + and every response carries ``next_cursor`` (``str | null``) -- resume a traversal by + passing the previous response's ``next_cursor`` back as ``cursor``. **Behavior change**: + because of this, a plain row-cap fill now reports ``truncated: false`` + a non-null + ``next_cursor`` instead of the old ``truncated: true``/``truncation_reason: "row_cap"`` + (there is a next page, not an error) -- ``truncated: true``/``"token_budget"`` is the + byte-budget signal below, and a match-budget trip still reports + ``truncation_reason="match_budget"``. A garbled/tampered/unrecognized ``cursor`` string + never raises: it comes back as a structured ``cursor_invalid: true`` payload with a remedy + ``reason`` and the normal empty envelope. + + ``max_bytes`` caps the serialized response size in bytes (default + ``CODE_SEARCH_MCP_MAX_RESPONSE_BYTES``, ~4 bytes/token; a request value only clamps the + server ceiling DOWN, never up). An over-budget response is truncated to fit -- tail-trimmed + in the payload's existing file order -- and flagged ``truncated: true``/ + ``truncation_reason: "token_budget"`` with a synthesized ``next_cursor`` so a caller can + keep paging through the CONTENT matches losslessly; a ``sym:`` query's page-1-only symbol + definitions that land in a truncated tail are lost from that traversal (flagged, not + silently dropped) since the symbol leg never re-runs on a continuation page. """ lc = ctx.request_context.lifespan_context engine, cfg = lc["engine"], lc["config"] @@ -274,11 +741,28 @@ async def search_code( query = _append_branch_atom(query, branch) if commit: query = _append_commit_atom(query, commit) - return await _dispatch("search_code", lambda: _search_code_payload(engine, cfg, query, limit)) + + def _build() -> dict[str, Any]: + try: + return _search_code_payload(engine, cfg, query, limit, cursor=cursor) + except service.CursorError as error: + return _cursor_invalid_payload(query, error) + + truncator = _make_search_code_truncator(engine, cfg) + return await _dispatch( + "search_code", + _build, + max_bytes=max_bytes, + shape=lambda payload, budget: _shape_response("search_code", payload, budget, truncator), + ) async def semantic_search( - query: str, ctx: Context, limit: int = 50, branch: str | None = None + query: str, + ctx: Context, + limit: int = 50, + branch: str | None = None, + max_bytes: int | None = None, ) -> str: """Semantic + BM25 hybrid search: rank indexed chunks by relevance to a free-text query. @@ -322,36 +806,88 @@ async def semantic_search( rank score), and ``similarity`` (raw cosine similarity against the query embedding, defined as ``1 - cosine_distance``; ``null`` for chunks with no embedding) -- ``rrf_score`` alone is not comparable across queries, ``similarity`` is. + + ``max_bytes`` caps the serialized response size in bytes (default + ``CODE_SEARCH_MCP_MAX_RESPONSE_BYTES``, ~4 bytes/token; a request value only clamps the + server ceiling DOWN, never up). An over-budget response tail-trims the (already + relevance-ranked) ``results`` list -- dropping the least relevant first -- and flags + ``truncated: true``/``truncation_reason: "token_budget"``; there is no resume handle for + this tool, so a truncated response is lossy (re-run with a smaller ``limit`` or a narrower + query to see what was cut). """ lc = ctx.request_context.lifespan_context engine, cfg = lc["engine"], lc["config"] limit = _clamp_limit(limit, cfg) return await _dispatch( - "semantic_search", lambda: _semantic_search_payload(engine, cfg, query, limit, branch) + "semantic_search", + lambda: _semantic_search_payload(engine, cfg, query, limit, branch), + max_bytes=max_bytes, ) -async def list_repos(ctx: Context) -> str: - """List every indexed repository with its branches and per-branch last-indexed metadata.""" +async def list_repos(ctx: Context, max_bytes: int | None = None) -> str: + """List every indexed repository with its branches and per-branch last-indexed metadata. + + ``max_bytes`` caps the serialized response size in bytes (default + ``CODE_SEARCH_MCP_MAX_RESPONSE_BYTES``, ~4 bytes/token; a request value only clamps the + server ceiling DOWN, never up). An over-budget response tail-trims ``repos`` and flags + ``truncated: true``/``truncation_reason: "token_budget"``; there is no resume handle for + this tool, so a truncated response is lossy. + """ lc = ctx.request_context.lifespan_context - return await _dispatch("list_repos", lambda: _list_repos_payload(lc["engine"], lc["config"])) + return await _dispatch( + "list_repos", + lambda: _list_repos_payload(lc["engine"], lc["config"]), + max_bytes=max_bytes, + ) -async def get_file(repo: str, path: str, ctx: Context, branch: str | None = None) -> str: - """Return the full content of a file by repository name and path (miss -> ``found:false``). +async def get_file( + repo: str, + path: str, + ctx: Context, + branch: str | None = None, + start_line: int = 1, + max_bytes: int | None = None, +) -> str: + """Return a file's content by repository name and path (miss -> ``found:false``). ``branch`` scopes the lookup to the content version indexed on that branch (one path may have several); omitted, it resolves to the repo's default branch. The resolved branch is echoed back in the payload. + + ``start_line`` (1-based; values ``< 1`` clamp to 1) pages through large files: the response + always carries ``start_line`` (echo) and ``next_start_line`` (the next page's ``start_line`` + when the file continues past what was returned, ``null`` when the tail fits). Content is + split on ``"\\n"`` -- the SAME rule ``search_code``'s match ``line`` numbers use, so a + ``get_file`` page's line numbers stay congruent with search results. Paging through every + page from ``start_line=1`` and rejoining each page's ``content`` with ``"\\n"`` reconstructs + the file byte-exactly (CRLF and no-trailing-newline files included). ``max_bytes`` caps the + serialized response size in bytes (default ``CODE_SEARCH_MCP_MAX_RESPONSE_BYTES``, ~4 + bytes/token; a request value only clamps the server ceiling DOWN, never up); an over-budget + page is cut to the largest whole-line prefix that fits and flagged ``truncated: true``/ + ``truncation_reason: "token_budget"``. Edge case: a single line whose JSON-encoded size + alone exceeds ``max_bytes`` (e.g. a minified one-line file) is still returned alone -- + flagged, with the response exceeding the budget for that one call -- because always making + forward progress outranks strict enforcement for that degenerate case; ``next_start_line`` + still advances past it. """ lc = ctx.request_context.lifespan_context + start_line = max(1, start_line) return await _dispatch( - "get_file", lambda: _get_file_payload(lc["engine"], lc["config"], repo, path, branch) + "get_file", + lambda: _get_file_payload(lc["engine"], lc["config"], repo, path, branch), + max_bytes=max_bytes, + shape=lambda payload, budget: _shape_get_file_response(payload, start_line, budget), ) async def find_references( - symbol: str, ctx: Context, limit: int = 200, branch: str | None = None + symbol: str, + ctx: Context, + limit: int = 200, + branch: str | None = None, + max_bytes: int | None = None, ) -> str: """Find candidate call sites of ``symbol`` corpus-wide, each with its ranked definitions. @@ -380,6 +916,13 @@ async def find_references( Composition -- "what tests cover symbol X": call ``find_references(X)`` and client-side filter ``sites`` by your test-path convention (e.g. ``file`` starts with ``"tests/"``); each surviving site's ``enclosing_symbol`` names the covering test. No separate tool is needed. + + ``max_bytes`` caps the serialized response size in bytes (default + ``CODE_SEARCH_MCP_MAX_RESPONSE_BYTES``, ~4 bytes/token; a request value only clamps the + server ceiling DOWN, never up). An over-budget response tail-trims ``sites`` (recomputing + ``site_count``/``resolution_summary`` from the survivors) and flags ``truncated: true``/ + ``truncation_reason: "token_budget"``; there is no resume handle for this tool, so a + truncated response is lossy (re-run with a smaller ``limit`` to see what was cut). """ lc = ctx.request_context.lifespan_context engine, cfg = lc["engine"], lc["config"] @@ -387,6 +930,7 @@ async def find_references( return await _dispatch( "find_references", lambda: _find_references_payload(engine, cfg, symbol, limit, branch), + max_bytes=max_bytes, ) @@ -397,6 +941,7 @@ async def list_imports( direction: str = "imports", branch: str | None = None, limit: int = 200, + max_bytes: int | None = None, ) -> str: """Enumerate ``import`` edge sites in one of two directions (candidate-set semantics). @@ -424,6 +969,13 @@ async def list_imports( ``file``, ``line``, ``edge_kind`` = ``"import"``, ``target_name``, ``enclosing_symbol`` | ``null`` for module scope, ``resolution``, ``candidate_count``, ``candidates_truncated``, ranked ``candidates``). ``limit`` caps the sites scanned (clamped to a server maximum). + + ``max_bytes`` caps the serialized response size in bytes (default + ``CODE_SEARCH_MCP_MAX_RESPONSE_BYTES``, ~4 bytes/token; a request value only clamps the + server ceiling DOWN, never up). An over-budget response tail-trims ``sites`` (recomputing + ``site_count``/``resolution_summary`` from the survivors) and flags ``truncated: true``/ + ``truncation_reason: "token_budget"``; there is no resume handle for this tool, so a + truncated response is lossy (re-run with a smaller ``limit`` to see what was cut). """ lc = ctx.request_context.lifespan_context engine, cfg = lc["engine"], lc["config"] @@ -433,6 +985,7 @@ async def list_imports( lambda: _list_imports_payload( engine, cfg, repo, limit, branch, target=target, direction=direction ), + max_bytes=max_bytes, ) diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 8212cd4..172db40 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -992,8 +992,11 @@ async def test_search_code_tool_appends_branch_atom_to_query( ) -> None: captured: dict[str, Any] = {} - def _fake_payload(engine: Any, cfg: Settings, query: str, limit: int) -> dict[str, Any]: + def _fake_payload( + engine: Any, cfg: Settings, query: str, limit: int, cursor: str | None = None + ) -> dict[str, Any]: captured["query"] = query + captured["cursor"] = cursor return {"query": query} monkeypatch.setattr(main, "_search_code_payload", _fake_payload) @@ -1002,6 +1005,8 @@ def _fake_payload(engine: Any, cfg: Settings, query: str, limit: int) -> dict[st await main.search_code("foo", ctx, branch="release/1.0") # type: ignore[arg-type] assert captured["query"] == 'foo branch:"release/1.0"' + # search_code always runs in pagination mode: page 1 passes cursor=None explicitly. + assert captured["cursor"] is None @pytest.mark.unit @@ -1011,7 +1016,9 @@ async def test_search_code_tool_leaves_query_untouched_without_branch( ) -> None: captured: dict[str, Any] = {} - def _fake_payload(engine: Any, cfg: Settings, query: str, limit: int) -> dict[str, Any]: + def _fake_payload( + engine: Any, cfg: Settings, query: str, limit: int, cursor: str | None = None + ) -> dict[str, Any]: captured["query"] = query return {"query": query} @@ -1267,6 +1274,21 @@ def _build() -> dict[str, Any]: assert "tool=search_code" in line assert "query_too_broad" in line assert "limiter_borrowed=" in line # pool/limiter saturation signal is wired + # AC5 / Step 4: pre- and post-shaping response byte sizes ship in the same log line. + assert "response_bytes_pre=" in line + assert "response_bytes=" in line + + +@pytest.mark.observability +@pytest.mark.asyncio +async def test_signals_log_includes_duration_ns_before_projection() -> None: + # duration_ns is read into _signals() BEFORE project_for_mcp drops it from the wire + # payload -- so log-line observability survives even though MCP callers never see the + # field itself. + payload = {"duration_ns": 123456, "files": []} + body, log_fields = main._shape_response("search_code", payload, 100_000, None) + assert log_fields["signals"]["duration_ns"] == 123456 + assert '"duration_ns"' not in body @pytest.mark.observability From e30e1191fe95da6241da4f5421bb03754f092199 Mon Sep 17 00:00:00 2001 From: Tanner Wendland Date: Fri, 24 Jul 2026 12:59:25 -0700 Subject: [PATCH 3/6] test: MCP response-shaping unit tests (Step 8, fate 3) New tests/unit/test_mcp_shaping.py: pure-function pins for _effective_budget's clamp matrix, _fit_list/_fit_lines exactness, project_for_mcp's drop-list (and None-safety on skeletal payloads), each static truncator's flag/recompute behavior, the pinned post-projection MCP search_code envelope shape, and the cursor_invalid structured payload. Fake-engine integration tests cover the search_code byte-budget truncation path end to end: cursor-traversal losslessness for content matches, the no-row and lookup-fault cursor-synthesis degrade paths (never raising), and the documented mixed grep+symbol carve-out (a page-1-only symbol match in a truncated tail is lost and flagged, while content matches for that same file remain recoverable). get_file gets a parametrized multi-page reassembly property test (plain/CRLF/no-trailing-newline/form-feed/unicode-line-separator fixtures), a line-numbering-congruence pin against split("\n") vs splitlines(), and the single-oversized-line documented edge. test_all_tools_respect_budget parametrizes all six tools (incl. a 200-site x 32-candidate find_references worst case) over the env default and a small max_bytes (AC2/AC6). test_lane3_fixture_reduction pins projection alone at >=25% serialized-byte reduction on a synthetic 200-file/6-match/6-range search_code fixture (AC4). Co-authored-by: Isaac --- tests/unit/test_mcp_shaping.py | 1006 ++++++++++++++++++++++++++++++++ 1 file changed, 1006 insertions(+) create mode 100644 tests/unit/test_mcp_shaping.py diff --git a/tests/unit/test_mcp_shaping.py b/tests/unit/test_mcp_shaping.py new file mode 100644 index 0000000..279b7e8 --- /dev/null +++ b/tests/unit/test_mcp_shaping.py @@ -0,0 +1,1006 @@ +"""Unit tests for the MCP response-shaping pipeline (byte budget, projection, truncation, +resume handles) added in app/main.py -- the "MCP response shaping" section between +``_signals`` and ``_dispatch``. + +No DB, no SDK: pure-function tests for ``project_for_mcp``/``_effective_budget``/``_fit_list``/ +``_fit_lines``/the per-tool truncators/``_shape_response``/``_shape_get_file_response``, plus a +few fake-engine integration tests for the ``search_code`` cursor-synthesis path (which needs a +name->id SELECT) and the full async tool wrappers. Complements (never duplicates) +``tests/unit/test_main.py``'s zoekt-parity pins: those assert the UNCHANGED service-layer +payload shape; these assert the MCP-only shaping layered on top. +""" + +from __future__ import annotations + +import copy +import json +from collections.abc import Callable +from typing import Any + +import pytest + +from app import main, service +from app.config import Settings +from app.search.grep import FileCursor, FileMatches, GrepResult, LineMatch +from app.search.symbols import SymbolMatch, SymbolResult +from tests.unit.test_main import ( + _cfg, + _FakeEngine, + _FakeLifespanContext, + _FakeResult, + _no_sym, + _Row, +) + +# --------------------------------------------------------------------------- _effective_budget + + +@pytest.mark.unit +@pytest.mark.parametrize( + "max_bytes,expected", + [ + (None, 100_000), + (0, 100_000), + (-5, 100_000), + (5_000, 5_000), + (100_000, 100_000), + (500_000, 100_000), # over the env max clamps DOWN + (10, 1024), # under the floor clamps UP to _MIN_MAX_BYTES + ], +) +def test_effective_budget_clamp_matrix(max_bytes: int | None, expected: int) -> None: + cfg = Settings(lakebase_endpoint=None, mcp_max_response_bytes=100_000) + assert main._effective_budget(max_bytes, cfg) == expected + + +# --------------------------------------------------------------------------------- _fit_list + + +@pytest.mark.unit +def test_fit_list_exactness() -> None: + items = [{"a": 1}, {"a": 2}, {"a": 3}] + costs = [len(json.dumps(item)) + 2 for item in items] + + assert main._fit_list(items, costs[0] + costs[1]) == 2 + assert main._fit_list(items, costs[0] + costs[1] - 1) == 1 + assert main._fit_list(items, 0) == 0 + assert main._fit_list(items, sum(costs)) == 3 + assert main._fit_list([], 100) == 0 + + +# ------------------------------------------------------------------------------- _fit_lines + + +@pytest.mark.unit +def test_fit_lines_newline_cost_term() -> None: + # cost("abc") = len(json.dumps("abc")) - 2 = 3 (the escaped body, excluding the two quotes + # json.dumps adds); each line after the first adds 2 more for the re-appended "\n". + assert main._line_cost("abc") == 3 + lines = ["abc", "de"] + assert main._fit_lines(lines, 3) == 1 # only the first line fits + assert main._fit_lines(lines, 6) == 1 # 3 + (2 + 2) = 7 > 6, still only the first + assert main._fit_lines(lines, 7) == 2 # exactly fits both + assert main._fit_lines([], 100) == 0 + + +# --------------------------------------------------------------------------- project_for_mcp + + +@pytest.mark.unit +def test_projection_drops_exactly_the_three_fields() -> None: + payload = { + "query": "foo", + "duration_ns": 123, + "files": [ + { + "repo": "acme/widgets", + "file": "f.py", + "content_sha": "deadbeef", + "permalink_branch": "main", + "matches": [{"line": 1, "text": "foo", "byte_ranges": [[0, 3]]}], + } + ], + } + main.project_for_mcp("search_code", payload) + assert "duration_ns" not in payload + assert "content_sha" not in payload["files"][0] + assert "byte_ranges" not in payload["files"][0]["matches"][0] + # Everything else survives untouched. + assert payload["files"][0]["repo"] == "acme/widgets" + assert payload["files"][0]["permalink_branch"] == "main" + assert payload["files"][0]["matches"][0]["text"] == "foo" + + +@pytest.mark.unit +def test_projection_is_noop_for_non_search_code_tools() -> None: + for tool in ("list_repos", "find_references", "list_imports", "semantic_search", "get_file"): + payload = {"repos": [{"name": "r"}], "count": 1} + before = copy.deepcopy(payload) + main.project_for_mcp(tool, payload) + assert payload == before + + +@pytest.mark.unit +def test_projection_is_none_safe_on_minimal_payloads() -> None: + # The wrapper tests' fakes return skeletal dicts like {"query": q}; projection must no-op + # rather than KeyError/TypeError. + payload: dict[str, Any] = {"query": "foo"} + main.project_for_mcp("search_code", payload) + assert payload == {"query": "foo"} + + payload2: dict[str, Any] = {"files": [{"repo": "r"}]} # no matches key at all + main.project_for_mcp("search_code", payload2) + assert payload2 == {"files": [{"repo": "r"}]} + + +# ---------------------------------------------------------------- static tail-trim truncators + + +@pytest.mark.unit +def test_list_repos_truncator_trims_tail_and_recomputes_count() -> None: + payload = { + "repos": [{"name": f"org/repo-{i}", "branches": ["main"]} for i in range(50)], + "count": 50, + } + main._TRUNCATORS["list_repos"](payload, 500) + assert payload["truncated"] is True + assert payload["truncation_reason"] == "token_budget" + assert payload["count"] == len(payload["repos"]) + assert 0 < payload["count"] < 50 + assert len(json.dumps(payload)) <= 500 + + +@pytest.mark.unit +def test_list_repos_truncator_none_safe_on_minimal_payload() -> None: + payload: dict[str, Any] = {"query": "x"} + main._TRUNCATORS["list_repos"](payload, 10) + assert payload["truncated"] is True + assert payload["truncation_reason"] == "token_budget" + + +def _reference_site(index: int, *, resolution: str = "unique") -> dict[str, Any]: + return { + "repo": "org/repo", + "file": f"src/site_{index}.py", + "line": index, + "edge_kind": "call", + "target_name": "Handler", + "enclosing_symbol": {"name": f"func_{index}", "kind": "function"}, + "resolution": resolution, + "candidate_count": 1, + "candidates_truncated": False, + "candidates": [ + { + "repo": "org/repo", + "file": "src/handler.py", + "line": 1, + "name": "Handler", + "kind": "function", + "same_repo": True, + "same_file": False, + "kind_match": True, + } + ], + } + + +@pytest.mark.unit +@pytest.mark.parametrize("tool", ["find_references", "list_imports"]) +def test_reference_sites_truncator_recomputes_summary(tool: str) -> None: + sites = [_reference_site(i) for i in range(300)] + payload = { + "sites": sites, + "site_count": len(sites), + "resolution_summary": {"unique": len(sites), "ambiguous": 0, "unresolved": 0}, + "truncated": False, + "truncation_reason": None, + } + main._TRUNCATORS[tool](payload, 2000) + assert payload["truncated"] is True + assert payload["truncation_reason"] == "token_budget" + assert payload["site_count"] == len(payload["sites"]) + assert payload["resolution_summary"]["unique"] == len(payload["sites"]) + assert payload["resolution_summary"]["ambiguous"] == 0 + assert len(json.dumps(payload)) <= 2000 + + +@pytest.mark.unit +def test_reference_sites_truncator_none_safe_on_minimal_payload() -> None: + payload: dict[str, Any] = {"query": "Handler"} + main._TRUNCATORS["find_references"](payload, 10) + assert payload["truncated"] is True + assert payload["truncation_reason"] == "token_budget" + + +@pytest.mark.unit +def test_semantic_results_truncator_drops_least_relevant_tail() -> None: + results = [ + { + "repo": "org/repo", + "file": f"src/f{i}.py", + "chunk_index": i, + "content": "x" * 200, + "start_line": 1, + "end_line": 10, + "rrf_score": 1.0 / (i + 1), + "similarity": 0.9, + } + for i in range(100) + ] + payload = {"query": "auth flow", "semantic_enabled": True, "results": results, "count": 100} + main._TRUNCATORS["semantic_search"](payload, 3000) + assert payload["truncated"] is True + assert payload["truncation_reason"] == "token_budget" + assert payload["count"] == len(payload["results"]) + assert payload["count"] < 100 + # Tail-trim: the surviving prefix is exactly the most-relevant leading results. + assert [r["chunk_index"] for r in payload["results"]] == list(range(payload["count"])) + assert len(json.dumps(payload)) <= 3000 + + +@pytest.mark.unit +def test_semantic_results_truncator_none_safe_on_minimal_payload() -> None: + payload: dict[str, Any] = {"query": "x"} + main._TRUNCATORS["semantic_search"](payload, 10) + assert payload["truncated"] is True + assert payload["truncation_reason"] == "token_budget" + + +# -------------------------------------------------------------------------- _shape_response + + +@pytest.mark.unit +def test_shape_response_under_budget_is_a_noop_pass_through() -> None: + payload = {"repos": [{"name": "r"}], "count": 1} + body, log_fields = main._shape_response( + "list_repos", payload, 100_000, main._TRUNCATORS["list_repos"] + ) + out = json.loads(body) + assert out == {"repos": [{"name": "r"}], "count": 1} + assert log_fields["pre_bytes"] == log_fields["post_bytes"] + + +@pytest.mark.unit +def test_shape_response_returns_flagged_irreducible_envelope_without_raising() -> None: + # No truncator (e.g. an unregistered tool name): never crash, just pass through even if + # over budget -- Principle 4, never hard-error. + payload = {"query": "x" * 10_000} + body, _ = main._shape_response("mystery_tool", payload, 100, None) + assert json.loads(body)["query"] == "x" * 10_000 + + +@pytest.mark.unit +def test_signals_log_includes_duration_ns_before_projection() -> None: + payload = {"duration_ns": 123456, "files": []} + body, log_fields = main._shape_response("search_code", payload, 100_000, None) + assert log_fields["signals"]["duration_ns"] == 123456 + assert '"duration_ns"' not in body + + +# ------------------------------------------------------------------------------ cursor_invalid + + +@pytest.mark.unit +def test_cursor_invalid_payload_shape() -> None: + error = service.CursorError("malformed pagination cursor: 'garbled'") + payload = main._cursor_invalid_payload("foo", error) + assert payload["cursor_invalid"] is True + assert "malformed pagination cursor" in payload["reason"] + assert payload["files"] == [] + assert payload["next_cursor"] is None + assert payload["truncated"] is False + + +@pytest.mark.unit +@pytest.mark.asyncio +async def test_search_code_tool_returns_cursor_invalid_payload_never_raises( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def _raise(*_a: object, **_k: object) -> dict[str, Any]: + raise service.CursorError("malformed pagination cursor: 'garbled'") + + monkeypatch.setattr(main, "_search_code_payload", _raise) + ctx = _FakeLifespanContext(_FakeEngine([]), _cfg()) + + out = await main.search_code("foo", ctx, cursor="garbled") # type: ignore[arg-type] + + payload = json.loads(out) + assert payload["cursor_invalid"] is True + assert payload["files"] == [] + + +@pytest.mark.unit +@pytest.mark.asyncio +async def test_search_code_tool_preserves_row_cap_pagination_signal_under_budget( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # Under budget, MCP shaping must not interfere with the builder's own pagination signal: + # a plain grep row-cap fill in pagination mode is truncated=False + a real next_cursor. + def _fake_payload( + engine: Any, cfg: Settings, query: str, limit: int, cursor: str | None = None + ) -> dict[str, Any]: + return { + "query": query, + "files": [], + "file_count": 0, + "match_count": 0, + "truncated": False, + "truncation_reason": None, + "next_cursor": "opaquetoken", + } + + monkeypatch.setattr(main, "_search_code_payload", _fake_payload) + ctx = _FakeLifespanContext(_FakeEngine([]), _cfg()) + + out = await main.search_code("foo", ctx) # type: ignore[arg-type] + + payload = json.loads(out) + assert payload["truncated"] is False + assert payload["next_cursor"] == "opaquetoken" + + +# ------------------------------------------------------------- pinned MCP envelope (post-proj) + + +@pytest.mark.unit +def test_mcp_search_code_envelope_shape_is_pinned_minus_dropped_plus_cursor() -> None: + payload = { + "query": "foo", + "file_count": 1, + "match_count": 1, + "duration_ns": 999, + "files": [ + { + "repo": "r", + "file": "f.py", + "language": "python", + "branches": ["main"], + "matches": [{"line": 1, "text": "foo", "byte_ranges": [[0, 3]]}], + "content_sha": "sha", + "permalink_branch": None, + } + ], + "truncated": False, + "truncation_reason": None, + "regex_incompatible": False, + "regex_invalid": None, + "query_too_broad": False, + "query_parse_error": None, + "no_content_atom": False, + "zero_width_only_atoms": False, + "next_cursor": None, + } + body, _ = main._shape_response("search_code", payload, 100_000, None) + out = json.loads(body) + + assert "duration_ns" not in out + assert set(out) == { + "query", + "file_count", + "match_count", + "files", + "truncated", + "truncation_reason", + "regex_incompatible", + "regex_invalid", + "query_too_broad", + "query_parse_error", + "no_content_atom", + "zero_width_only_atoms", + "next_cursor", + } + (file_entry,) = out["files"] + assert "content_sha" not in file_entry + assert set(file_entry) == { + "repo", + "file", + "language", + "branches", + "matches", + "permalink_branch", + } + (match,) = file_entry["matches"] + assert "byte_ranges" not in match + assert set(match) == {"line", "text"} + + +# -------------------------------------------------------------------------------- get_file + + +def _get_file_payload_fixture(content: str, *, found: bool = True) -> dict[str, Any]: + return { + "repo": "acme/widgets", + "path": "big.txt", + "branch": "main", + "content": content if found else None, + "found": found, + "commit": "abc1234" if found else None, + } + + +@pytest.mark.unit +def test_get_file_miss_never_truncates() -> None: + payload = _get_file_payload_fixture("irrelevant", found=False) + body, _ = main._shape_get_file_response(payload, 1, 100_000) + out = json.loads(body) + assert out["truncated"] is False + assert out["truncation_reason"] is None + assert out["next_start_line"] is None + assert out["start_line"] == 1 + assert out["found"] is False + + +@pytest.mark.unit +@pytest.mark.parametrize( + "content", + [ + "\n".join(f"line {i} " + "x" * 50 for i in range(2000)), + "\n".join(f"line {i}\r" for i in range(2000)), # CRLF-shaped: trailing \r per line + "\n".join(f"line {i}" for i in range(500)) + "\nlast line no trailing newline", + "\n".join(f"line {i} \x0c tail" for i in range(500)), # form feed + "\n".join(f"line {i} 

 tail" for i in range(500)), # U+2028 / U+2029 + ], + ids=["plain", "crlf", "no-trailing-newline", "form-feed", "unicode-seps"], +) +def test_get_file_traversal_reassembles_exact_content(content: str) -> None: + budget = 2000 # small enough to force many pages + start_line = 1 + pages: list[str] = [] + guard = 0 + while True: + guard += 1 + assert guard < 5000, "traversal did not terminate" + payload = _get_file_payload_fixture(content) + body, _ = main._shape_get_file_response(payload, start_line, budget) + assert len(body) <= budget + page = json.loads(body) + pages.append(page["content"]) + if page["next_start_line"] is None: + break + assert page["next_start_line"] > start_line + start_line = page["next_start_line"] + assert "\n".join(pages) == content + + +@pytest.mark.unit +def test_get_file_line_numbering_matches_split_not_splitlines() -> None: + # grep.py:436 numbers lines via content.split("\n"); str.splitlines() would (wrongly, for + # this purpose) also break on \x0c/
/
, diverging from search_code's line numbers. + content = "a\x0cb\nc
d\ne" + assert len(content.splitlines()) > len(content.split("\n")) + + payload = _get_file_payload_fixture(content) + body, _ = main._shape_get_file_response(payload, 1, 100_000) + out = json.loads(body) + assert out["content"] == content + assert out["next_start_line"] is None + assert out["truncated"] is False + + +@pytest.mark.unit +def test_get_file_single_oversized_line_returned_flagged_over_budget() -> None: + huge_line = "x" * 5000 + content = f"short\n{huge_line}\nshort2" + payload = _get_file_payload_fixture(content) + + body, _ = main._shape_get_file_response(payload, 2, 1000) # page starting at the huge line + + out = json.loads(body) + assert out["truncated"] is True + assert out["truncation_reason"] == "token_budget" + assert out["content"] == huge_line # progress guarantee: >= 1 line always returned + assert len(body) > 1000 # documented edge: exceeds the budget for this one call + assert out["next_start_line"] == 3 + + +@pytest.mark.unit +def test_get_file_start_line_clamps_below_one() -> None: + content = "a\nb\nc" + payload = _get_file_payload_fixture(content) + body, _ = main._shape_get_file_response(payload, 0, 100_000) + out = json.loads(body) + assert out["start_line"] == 1 + assert out["content"] == content + + +@pytest.mark.unit +def test_get_file_start_line_past_eof_returns_empty_untruncated() -> None: + content = "a\nb\nc" + payload = _get_file_payload_fixture(content) + body, _ = main._shape_get_file_response(payload, 100, 100_000) + out = json.loads(body) + assert out["content"] == "" + assert out["truncated"] is False + assert out["next_start_line"] is None + + +# ------------------------------------------------------ search_code cursor synthesis (D3) + + +class _ScalarOnly: + def __init__(self, value: int | None) -> None: + self._value = value + + def scalar_one_or_none(self) -> int | None: + return self._value + + +class _FakeSearchConn: + """Understands exactly the two SELECT shapes search_code_payload/_resolve_repo_id issue: + ``select(Repo.id, Repo.name)`` (the repo-name map) and ``select(Repo.id).where(Repo.name == + ...)`` (the truncation-time cursor-synthesis lookup) -- dispatched by column name rather + than a real dialect/DB, using SQLAlchemy's own compiled-bind-param introspection.""" + + def __init__(self, name_by_id: dict[int, str], *, raise_on_id_lookup: bool = False) -> None: + self._name_by_id = name_by_id + self._id_by_name = {name: repo_id for repo_id, name in name_by_id.items()} + self._raise_on_id_lookup = raise_on_id_lookup + self.driver_sql: list[str] = [] + + def __enter__(self) -> _FakeSearchConn: + return self + + def __exit__(self, *exc: object) -> None: + return None + + def begin(self) -> _FakeSearchConn: + return self + + def exec_driver_sql(self, sql: str) -> None: + self.driver_sql.append(sql) + + def execute(self, stmt: Any) -> Any: + cols = list(stmt.selected_columns.keys()) + if cols == ["id", "name"]: + rows = [_Row(id=repo_id, name=name) for repo_id, name in self._name_by_id.items()] + return _FakeResult(rows) + if cols == ["id"]: + if self._raise_on_id_lookup: + raise RuntimeError("simulated lookup fault") + params = stmt.compile().params + name = next(iter(params.values())) + return _ScalarOnly(self._id_by_name.get(name)) + raise AssertionError(f"unexpected query shape: {cols}") + + +class _FakeSearchEngine: + def __init__(self, name_by_id: dict[int, str], *, raise_on_id_lookup: bool = False) -> None: + self._conn = _FakeSearchConn(name_by_id, raise_on_id_lookup=raise_on_id_lookup) + + def connect(self) -> _FakeSearchConn: + return self._conn + + +def _make_grep_stub(all_files: list[FileMatches]) -> Callable[..., GrepResult]: + """A cursor-aware service.grep_search stand-in: returns every candidate strictly after + ``cursor`` in (repo_id, path, content_sha) order, UNCAPPED (no row_cap) -- so byte-budget + truncation at the MCP layer is the only truncation source these tests exercise. + """ + ordered = sorted(all_files, key=lambda f: (f.repo_id, f.path, f.content_sha)) + + def _grep( + conn: Any, query: str, *, cursor: FileCursor | None = None, **_kwargs: Any + ) -> GrepResult: + if cursor is None: + remaining = ordered + else: + remaining = [ + f + for f in ordered + if (f.repo_id, f.path, f.content_sha) + > (cursor.repo_id, cursor.path, cursor.content_sha) + ] + return GrepResult( + files=tuple(remaining), + truncated=False, + truncation_reason=None, + regex_incompatible=False, + no_content_atom=False, + zero_width_only_atoms=False, + next_cursor=None, + ) + + return _grep + + +@pytest.mark.unit +@pytest.mark.asyncio +async def test_search_code_cursor_traversal_equals_uncapped( + monkeypatch: pytest.MonkeyPatch, +) -> None: + all_files = [ + FileMatches( + repo_id=7, + path=f"file{i}.py", + lang="python", + content_sha=f"sha{i}", + branches=("main",), + line_matches=(LineMatch(1, f"needle in file {i}", ((0, 6),)),), + ) + for i in range(10) + ] + monkeypatch.setattr(service, "grep_search", _make_grep_stub(all_files)) + monkeypatch.setattr(service, "symbol_search", lambda *a, **k: _no_sym()) + engine = _FakeSearchEngine({7: "acme/widgets"}) + ctx = _FakeLifespanContext(engine, _cfg()) + + # Page 1: a small max_bytes forces MCP-level (pure tail-trim) truncation. + out1 = await main.search_code("needle", ctx, max_bytes=1200) # type: ignore[arg-type] + page1 = json.loads(out1) + assert page1["truncated"] is True + assert page1["truncation_reason"] == "token_budget" + page1_paths = {f["file"] for f in page1["files"]} + assert page1_paths, "page 1 must keep at least one file (progress guarantee)" + assert "file9.py" not in page1_paths # lexicographically last -> tail-dropped + cursor = page1["next_cursor"] + assert cursor is not None + + # Page 2: a generous budget lets the (already smaller) remainder fit in one page. + out2 = await main.search_code( # type: ignore[arg-type] + "needle", ctx, cursor=cursor, max_bytes=1_000_000 + ) + page2 = json.loads(out2) + page2_paths = {f["file"] for f in page2["files"]} + + assert page1_paths | page2_paths == {f"file{i}.py" for i in range(10)} + assert not (page1_paths & page2_paths) # no double-counting across pages + + +@pytest.mark.unit +@pytest.mark.asyncio +async def test_search_code_cursor_no_row_degrades_without_cursor_or_raise( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # The builder's own str(repo_id) fallback (service.py:764) means the payload's `repo` + # string may match no `repos.name` row -- the truncator must degrade to a flagged, + # handle-less truncation rather than raising. + all_files = [ + FileMatches( + repo_id=7, + path=f"file{i}.py", + lang="python", + content_sha=f"sha{i}", + branches=("main",), + line_matches=(LineMatch(1, f"needle in file {i}", ((0, 6),)),), + ) + for i in range(10) + ] + monkeypatch.setattr(service, "grep_search", _make_grep_stub(all_files)) + monkeypatch.setattr(service, "symbol_search", lambda *a, **k: _no_sym()) + # No name in the map at all -> _repo_name_map falls back to str(repo_id) for every file, + # so _resolve_repo_id's name lookup can never find a row. + engine = _FakeSearchEngine({}) + ctx = _FakeLifespanContext(engine, _cfg()) + + out = await main.search_code("needle", ctx, max_bytes=1200) # type: ignore[arg-type] + payload = json.loads(out) + + assert payload["truncated"] is True + assert payload["truncation_reason"] == "token_budget" + assert payload["next_cursor"] is None # degraded: no cursor synthesized, never raised + + +@pytest.mark.unit +@pytest.mark.asyncio +async def test_search_code_cursor_resolve_fault_degrades_without_raising( + monkeypatch: pytest.MonkeyPatch, +) -> None: + all_files = [ + FileMatches( + repo_id=7, + path=f"file{i}.py", + lang="python", + content_sha=f"sha{i}", + branches=("main",), + line_matches=(LineMatch(1, f"needle in file {i}", ((0, 6),)),), + ) + for i in range(10) + ] + monkeypatch.setattr(service, "grep_search", _make_grep_stub(all_files)) + monkeypatch.setattr(service, "symbol_search", lambda *a, **k: _no_sym()) + engine = _FakeSearchEngine({7: "acme/widgets"}, raise_on_id_lookup=True) + ctx = _FakeLifespanContext(engine, _cfg()) + + out = await main.search_code("needle", ctx, max_bytes=1200) # type: ignore[arg-type] + payload = json.loads(out) + + assert payload["truncated"] is True + assert payload["next_cursor"] is None # the id lookup faulted -- degrade, never raise + + +@pytest.mark.unit +@pytest.mark.asyncio +async def test_mixed_grep_symbol_truncation_traversal(monkeypatch: pytest.MonkeyPatch) -> None: + # AC3's documented carve-out: the symbol leg folds in page-1-only. When byte-budget + # truncation drops the tail file carrying a symbol match, that symbol is lost from the + # traversal (never reappears on a continuation page), while CONTENT matches for every file + # -- including the one that lost its symbol -- are still recoverable losslessly via cursor. + all_files = [ + FileMatches( + repo_id=7, + path=f"file{i}.py", + lang="python", + content_sha=f"sha{i}", + branches=("main",), + line_matches=(LineMatch(1, f"needle in file {i}", ((0, 6),)),), + ) + for i in range(10) + ] + monkeypatch.setattr(service, "grep_search", _make_grep_stub(all_files)) + monkeypatch.setattr( + service, + "symbol_search", + lambda *a, **k: SymbolResult( + symbols=( + SymbolMatch( + repo_id=7, + path="file9.py", # lexicographically last -> guaranteed tail-dropped + lang="python", + content_sha="sha9", + branches=("main",), + name="Handler", + kind="function", + start_line=1, + ), + ), + truncated=False, + truncation_reason=None, + no_symbol_atom=False, + ), + ) + engine = _FakeSearchEngine({7: "acme/widgets"}) + ctx = _FakeLifespanContext(engine, _cfg()) + + out1 = await main.search_code( # type: ignore[arg-type] + "needle sym:Handler", ctx, max_bytes=1500 + ) + page1 = json.loads(out1) + assert page1["truncated"] is True + page1_paths = {f["file"] for f in page1["files"]} + assert "file9.py" not in page1_paths + + def _has_symbols(payload: dict[str, Any]) -> bool: + return any( + "symbols" in match for file_entry in payload["files"] for match in file_entry["matches"] + ) + + assert _has_symbols(page1) is False # dropped before it could even be considered + + cursor = page1["next_cursor"] + assert cursor is not None + out2 = await main.search_code( # type: ignore[arg-type] + "needle sym:Handler", ctx, cursor=cursor, max_bytes=1_000_000 + ) + page2 = json.loads(out2) + page2_paths = {f["file"] for f in page2["files"]} + + # Content matches: fully recoverable across the two pages, including file9's own content. + assert page1_paths | page2_paths == {f"file{i}.py" for i in range(10)} + # The symbol match itself never reappears -- continuation pages skip the symbol leg + # entirely (page-1-only folding), so it is genuinely and permanently lost. + assert _has_symbols(page2) is False + + +# --------------------------------------------------------------------- AC2: budget respected + + +def _worst_case_payload(tool: str) -> dict[str, Any]: + if tool == "list_repos": + return { + "repos": [ + { + "name": f"org/{'x' * 40}-repo-{i}", + "branches": ["main", "develop", "release/1.0"], + "index_time": "2026-07-18T00:00:00+00:00", + "default_branch": "main", + "last_indexed_commit": "a" * 40, + "branch_details": [ + { + "branch": "main", + "last_indexed_commit": "a" * 40, + "index_time": "2026-07-18T00:00:00+00:00", + } + ], + } + for i in range(500) + ], + "count": 500, + } + if tool in ("find_references", "list_imports"): + sites = [] + for i in range(200): + candidates = [ + { + "repo": f"org/repo-{j}", + "file": f"src/module_{j}.py", + "line": j, + "name": "Handler", + "kind": "function", + "same_repo": j % 2 == 0, + "same_file": False, + "kind_match": True, + } + for j in range(32) + ] + sites.append( + { + "repo": "org/repo", + "file": f"src/site_{i}.py", + "line": i, + "edge_kind": "call" if tool == "find_references" else "import", + "target_name": "Handler", + "enclosing_symbol": {"name": f"func_{i}", "kind": "function"}, + "resolution": "ambiguous", + "candidate_count": 32, + "candidates_truncated": False, + "candidates": candidates, + } + ) + base: dict[str, Any] = { + "query": "Handler", + "sites": sites, + "site_count": 200, + "resolution_summary": {"unique": 0, "ambiguous": 200, "unresolved": 0}, + "truncated": False, + "truncation_reason": None, + "query_too_broad": False, + } + if tool == "find_references": + base.update({"kind": "references", "symbol": "Handler", "branch": None}) + else: + base.update( + { + "kind": "imports", + "direction": "imports", + "repo": "org/repo", + "repo_known": True, + "target": None, + "branch": None, + } + ) + return base + if tool == "semantic_search": + return { + "query": "how does auth work", + "semantic_enabled": True, + "results": [ + { + "repo": "org/repo", + "file": f"src/file_{i}.py", + "chunk_index": i, + "content": "def handler():\n pass\n" * 20, + "start_line": 1, + "end_line": 40, + "rrf_score": 0.01, + "similarity": 0.9, + } + for i in range(200) + ], + "count": 200, + } + if tool == "search_code": + files = [ + { + "repo": "org/repo", + "file": f"src/file_{i}.py", + "language": "python", + "branches": ["main"], + "matches": [ + {"line": j, "text": "match text " * 5, "byte_ranges": [[0, 5], [10, 15]]} + for j in range(10) + ], + "content_sha": "a" * 40, + "permalink_branch": "main", + } + for i in range(200) + ] + return { + "query": "handler", + "file_count": 200, + "match_count": 2000, + "duration_ns": 123456, + "files": files, + "truncated": False, + "truncation_reason": None, + "regex_incompatible": False, + "regex_invalid": None, + "query_too_broad": False, + "query_parse_error": None, + "no_content_atom": False, + "zero_width_only_atoms": False, + "next_cursor": None, + } + if tool == "get_file": + content = "\n".join(f"line {i} " + "x" * 60 for i in range(20_000)) + return { + "repo": "org/repo", + "path": "big.py", + "branch": "main", + "content": content, + "found": True, + "commit": "a" * 40, + } + raise ValueError(tool) + + +@pytest.mark.unit +@pytest.mark.parametrize("max_bytes", [None, 5000]) +@pytest.mark.parametrize( + "tool", + ["list_repos", "find_references", "list_imports", "semantic_search", "search_code", "get_file"], +) +def test_all_tools_respect_budget(tool: str, max_bytes: int | None) -> None: + cfg = Settings(lakebase_endpoint=None, mcp_max_response_bytes=100_000) + budget = main._effective_budget(max_bytes, cfg) + payload = _worst_case_payload(tool) + + if tool == "get_file": + body, _ = main._shape_get_file_response(payload, 1, budget) + elif tool == "search_code": + truncator = main._make_search_code_truncator(_FakeSearchEngine({}), cfg) + body, _ = main._shape_response(tool, payload, budget, truncator) + else: + body, _ = main._shape_response(tool, payload, budget, main._TRUNCATORS[tool]) + + assert len(body) <= budget + + +# ------------------------------------------------------------------- AC4: Lane-3 fixture + + +def _lane3_search_code_fixture( + n_files: int = 200, n_matches: int = 6, n_ranges: int = 6 +) -> dict[str, Any]: + """A synthetic search_code payload sized like the trace's measured shape (dense matches, + each carrying content_sha/byte_ranges/duration_ns) -- ported as a deterministic fixture + rather than depending on the scratchpad measure_tokens.py script.""" + files = [] + for i in range(n_files): + matches = [ + { + "line": j, + "text": "x" * 24, + "byte_ranges": [[k, k + 3] for k in range(n_ranges)], + } + for j in range(n_matches) + ] + files.append( + { + "repo": "org/repo", + "file": f"src/file_{i}.py", + "language": "python", + "branches": ["main"], + "matches": matches, + "content_sha": "a" * 40, + "permalink_branch": "main", + } + ) + return { + "query": "handler", + "file_count": n_files, + "match_count": n_files * n_matches, + "duration_ns": 123456, + "files": files, + "truncated": False, + "truncation_reason": None, + "regex_incompatible": False, + "regex_invalid": None, + "query_too_broad": False, + "query_parse_error": None, + "no_content_atom": False, + "zero_width_only_atoms": False, + "next_cursor": None, + } + + +@pytest.mark.unit +def test_lane3_fixture_reduction_at_least_25pct() -> None: + payload = _lane3_search_code_fixture() + original_size = len(json.dumps(payload)) + + projected = copy.deepcopy(payload) + main.project_for_mcp("search_code", projected) + projected_size = len(json.dumps(projected)) + + reduction = 1 - (projected_size / original_size) + assert reduction >= 0.25, f"projection-only reduction was {reduction:.1%}, need >= 25%" From 65447bad977e725302ed1aabd0bce28ec930cd96 Mon Sep 17 00:00:00 2001 From: Tanner Wendland Date: Fri, 24 Jul 2026 13:01:17 -0700 Subject: [PATCH 4/6] docs: document MCP response-size limits and the search_code pagination change Update the MCP tools table (cursor/start_line/max_bytes params) and add a "Response size limits" section covering CODE_SEARCH_MCP_MAX_RESPONSE_BYTES (default 100000, ~4 bytes/token heuristic), per-request max_bytes clamping down, the truncated/truncation_reason="token_budget" signal, and the search_code/get_file resume handles (next_cursor/next_start_line). Release-notes the search_code behavior change: because the MCP tool now always runs in pagination mode, a plain row-cap fill reports truncated=false + next_cursor instead of the old truncated=true/truncation_reason="row_cap". Co-authored-by: Isaac --- README.md | 69 +++++++++++++++++++++++++++++++++++++++++++++---------- 1 file changed, 57 insertions(+), 12 deletions(-) diff --git a/README.md b/README.md index 482877f..6b8d7ad 100644 --- a/README.md +++ b/README.md @@ -250,16 +250,59 @@ event loop. | Tool | Parameters | Returns | |---|---|---| -| `search_code` | `query`, `limit=200`, `branch=None`, `commit=None` | file-grouped line matches with byte ranges | -| `semantic_search` | `query`, `limit=50`, `branch=None` | ranked chunks with `rrf_score` | -| `list_repos` | — | indexed repos with per-branch last-indexed metadata | -| `get_file` | `repo`, `path`, `branch=None` | full file content, or `found: false` | -| `find_references` | `symbol`, `limit=200`, `branch=None` | ranked candidate reference (call) sites with enclosing symbols | -| `list_imports` | `repo=None`, `target=None`, `direction=imports`, `branch=None`, `limit=200` | import edge sites; `repo` required for `imports`, `target` required for `imported_by` | +| `search_code` | `query`, `limit=200`, `branch=None`, `commit=None`, `cursor=None`, `max_bytes=None` | file-grouped line matches with byte ranges | +| `semantic_search` | `query`, `limit=50`, `branch=None`, `max_bytes=None` | ranked chunks with `rrf_score` | +| `list_repos` | `max_bytes=None` | indexed repos with per-branch last-indexed metadata | +| `get_file` | `repo`, `path`, `branch=None`, `start_line=1`, `max_bytes=None` | a page of file content, or `found: false` | +| `find_references` | `symbol`, `limit=200`, `branch=None`, `max_bytes=None` | ranked candidate reference (call) sites with enclosing symbols | +| `list_imports` | `repo=None`, `target=None`, `direction=imports`, `branch=None`, `limit=200`, `max_bytes=None` | import edge sites; `repo` required for `imports`, `target` required for `imported_by` | Every tool returns a JSON string. `limit` is clamped server-side: a non-positive value falls back to 200, and anything above 1000 is capped there. +### Response size limits + +Every MCP tool response is capped at a byte-denominated budget — `CODE_SEARCH_MCP_MAX_RESPONSE_BYTES`, +default `100000` bytes (~25k tokens at a ~4 bytes/token heuristic; no tokenizer dependency — +the budget is enforced against the exact serialized JSON string sent over the wire). Every +tool also accepts a per-request `max_bytes`, which only clamps the server ceiling **down**, +never up. This is enforced entirely on the MCP surface (`app/main.py`); the web UI's REST API +and the underlying payload builders are unaffected. + +An over-budget response is never a failed tool call: it is truncated to fit, tail-trimmed in +each tool's dominant list (or, for `get_file`, cut to the largest whole-line prefix that fits), +and flagged `truncated: true` / `truncation_reason: "token_budget"` (extending the existing +`byte_cap`/`row_cap`/`match_budget` reasons). `search_code` and `get_file` carry resume +handles so a byte-budget trim is still traversable: + +- **`search_code`** now accepts `cursor` and always runs in pagination mode — page 1 omits + `cursor` (or passes `null`), and every response carries `next_cursor` (`str | null`); pass + the previous response's `next_cursor` back as `cursor` to resume. When byte-budget + truncation drops tail files, `next_cursor` is synthesized to resume exactly after the last + file kept, so a full traversal recovers every **content** match losslessly. A `sym:` query's + definitions fold in on page 1 only, so a symbol match that lands in a byte-budget-truncated + tail is lost from that traversal (flagged `truncated`, not silently dropped) — a continuation + page never re-runs the symbol leg. A garbled/tampered `cursor` string never raises: it comes + back as `{"cursor_invalid": true, "reason": "..."}` with the normal empty envelope. +- **`get_file`** now accepts `start_line` (1-based; values below 1 clamp to 1) and every + response carries `next_start_line` (the next page's `start_line`, or `null` when the file's + tail already fit). Content is split on `"\n"` — the same rule `search_code`'s match line + numbers use — so paging from `start_line=1` and rejoining each page's `content` with `"\n"` + reconstructs the file byte-exactly (CRLF and no-trailing-newline files included). A single + line whose JSON-encoded size alone exceeds the budget is still returned alone, flagged, with + that one response exceeding the budget — always making forward progress outranks strict + enforcement for that degenerate (e.g. minified one-line file) case. + +`list_repos`, `find_references`, `list_imports`, and `semantic_search` have no cursor +plumbing, so their truncation (tail-trimmed `repos`/`sites`/`results`) is lossy: re-run with a +smaller `limit` or a narrower query to see what was cut. + +**Behavior change:** because `search_code` now always runs in pagination mode, a plain grep +row-cap fill reports `truncated: false` plus a non-null `next_cursor` (there is a next page, +not an error) instead of the previous `truncated: true` / `truncation_reason: "row_cap"`. +`truncated: true` / `truncation_reason: "token_budget"` is the new byte-budget signal, and a +match-budget trip still reports `truncation_reason: "match_budget"` as before. + `branch` behaves differently per tool because `search_code` takes zoekt grammar and `semantic_search` takes natural language: on `search_code` it is sugar for appending `branch:""` to the query string (quoted, so `/`, `.`, and spaces need no escaping @@ -272,13 +315,15 @@ repo with no `default_branch` recorded). Recoverable conditions come back as payload fields — `query_parse_error`, `query_too_broad`, `truncated`, `regex_incompatible`, `regex_invalid`, -`no_content_atom`, `zero_width_only_atoms`, `commit_not_indexed` — rather than errors, so -an agent can react without a failed tool call. `regex_invalid` is distinct from -`regex_incompatible`: the latter means Python `regex` (not Postgres) rejected an otherwise-valid -pattern and only degrades highlighting; `regex_invalid` means Postgres rejected the pattern -outright and the query did not run. Pagination rides the same envelope as +`no_content_atom`, `zero_width_only_atoms`, `commit_not_indexed`, `cursor_invalid` — rather +than errors, so an agent can react without a failed tool call. `regex_invalid` is distinct +from `regex_incompatible`: the latter means Python `regex` (not Postgres) rejected an +otherwise-valid pattern and only degrades highlighting; `regex_invalid` means Postgres +rejected the pattern outright and the query did not run. Pagination rides the same envelope as `next_cursor`, and the semantic tool adds its own status fields (`semantic_enabled`, -`semantic_schema_missing`). +`semantic_schema_missing`). See [Response size limits](#response-size-limits) for the +byte-budget truncation signal (`truncation_reason: "token_budget"`) and the resume handles +(`next_cursor`/`next_start_line`) that ship with it. `semantic_search` is natural-language hybrid search (vector ANN + BM25 fused by reciprocal rank). It is **on by default** — the `chunks` schema rides the core migration chain and From 204a11c6c3cb9b476d9a39ec2cb13650673e7878 Mon Sep 17 00:00:00 2001 From: Tanner Wendland Date: Fri, 24 Jul 2026 13:02:52 -0700 Subject: [PATCH 5/6] test: pin max_bytes threading for all six MCP tools (AC6) test_all_six_tools_thread_max_bytes_to_dispatch monkeypatches main._dispatch to capture the max_bytes kwarg each tool wrapper passes through, closing the per-tool threading gap the plan called out alongside test_effective_budget_clamp_matrix. Co-authored-by: Isaac --- tests/unit/test_main.py | 34 ++++++++++++++++++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 172db40..d9e82ee 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -1167,6 +1167,40 @@ def _fake_payload( assert captured["repo"] is None +# ------------------------------------------------------------- max_bytes threading (AC6) + + +@pytest.mark.unit +@pytest.mark.asyncio +async def test_all_six_tools_thread_max_bytes_to_dispatch(monkeypatch: pytest.MonkeyPatch) -> None: + captured: dict[str, int | None] = {} + + async def _fake_dispatch( + name: str, build: Any, *, max_bytes: int | None = None, **_kw: Any + ) -> str: + captured[name] = max_bytes + return "{}" + + monkeypatch.setattr(main, "_dispatch", _fake_dispatch) + ctx = _FakeLifespanContext(_FakeEngine([]), _cfg()) + + await main.search_code("foo", ctx, max_bytes=111) # type: ignore[arg-type] + await main.semantic_search("foo", ctx, max_bytes=222) # type: ignore[arg-type] + await main.list_repos(ctx, max_bytes=333) # type: ignore[arg-type] + await main.get_file("acme/widgets", "f.py", ctx, max_bytes=444) # type: ignore[arg-type] + await main.find_references("Handler", ctx, max_bytes=555) # type: ignore[arg-type] + await main.list_imports(ctx, repo="acme/widgets", max_bytes=666) # type: ignore[arg-type] + + assert captured == { + "search_code": 111, + "semantic_search": 222, + "list_repos": 333, + "get_file": 444, + "find_references": 555, + "list_imports": 666, + } + + # ------------------------------------------------- search_code: divergent content_sha merge From a0d014a6e40f506e0232cb307ffba3ff4160600a Mon Sep 17 00:00:00 2001 From: Tanner Wendland Date: Fri, 24 Jul 2026 13:14:17 -0700 Subject: [PATCH 6/6] fix: search_code truncator loses progress guarantee on an oversized file Blocking defect from code review: when the first file's own serialized size alone exceeded the byte budget, _fit_list returned keep=0, and the `if keep > 0` guard around cursor synthesis meant the response came back as `files: [], next_cursor: null, truncated: true` -- a silent, unrecoverable dead end under default config (confirmed: one file with 2000 matches serializes past 300KB > the 100_000 default budget). Fix mirrors get_file's existing single-oversized-line edge: floor `keep` at 1 whenever the untruncated payload had at least one file (that one response may still exceed the budget -- the same documented trade-off get_file already makes), and always attempt cursor synthesis from whatever file was kept. The safety re-serialize/shrink loop is now floored at keep=1 rather than allowed to reach 0, so the guarantee can never be undone there either. Also, per review: - Removed dangling "Principle 4"/"AC6"/"D2"/"D3" plan-doc references in app/main.py comments/docstrings, replaced with the actual constraint spelled out in words (app/main.py, README.md). - Deduped test_signals_log_includes_duration_ns_before_projection, which existed verbatim in both test_main.py and test_mcp_shaping.py; kept the copy in test_mcp_shaping.py (its charter is the shaping pipeline). - Added test_search_code_single_oversized_file_progress_guarantee: a single file whose own serialized size exceeds max_bytes still comes back alone, flagged, with a next_cursor that lets a caller traverse past it to the remaining files -- exercising the exact scenario the prior test only asserted in a comment without ever constructing. Co-authored-by: Isaac --- README.md | 6 +++- app/main.py | 64 ++++++++++++++++++++++------------ tests/unit/test_main.py | 15 ++------ tests/unit/test_mcp_shaping.py | 58 ++++++++++++++++++++++++++++++ 4 files changed, 107 insertions(+), 36 deletions(-) diff --git a/README.md b/README.md index 6b8d7ad..e553810 100644 --- a/README.md +++ b/README.md @@ -279,7 +279,11 @@ handles so a byte-budget trim is still traversable: `cursor` (or passes `null`), and every response carries `next_cursor` (`str | null`); pass the previous response's `next_cursor` back as `cursor` to resume. When byte-budget truncation drops tail files, `next_cursor` is synthesized to resume exactly after the last - file kept, so a full traversal recovers every **content** match losslessly. A `sym:` query's + file kept, so a full traversal recovers every **content** match losslessly. At least one + file is always kept and a `next_cursor` always synthesized as long as there was at least one + file to begin with — a single file whose own serialized size alone exceeds `max_bytes` (e.g. + one file with thousands of matches) is still returned alone, with that one response exceeding + the budget, rather than coming back as an empty, unresumable dead end. A `sym:` query's definitions fold in on page 1 only, so a symbol match that lands in a byte-budget-truncated tail is lost from that traversal (flagged `truncated`, not silently dropped) — a continuation page never re-runs the symbol leg. A garbled/tampered `cursor` string never raises: it comes diff --git a/app/main.py b/app/main.py index d7d5e13..eec9bc2 100644 --- a/app/main.py +++ b/app/main.py @@ -172,7 +172,9 @@ def _effective_budget(max_bytes: int | None, cfg: Settings) -> int: """Resolve the byte budget for one call: `max_bytes` clamps the env default DOWN, never up. `None` or a non-positive value means "no per-request override" -> the env-configured - ceiling. Otherwise clamped to `[_MIN_MAX_BYTES, cfg.mcp_max_response_bytes]` (AC6). + ceiling. Otherwise clamped to `[_MIN_MAX_BYTES, cfg.mcp_max_response_bytes]` -- a request + can shrink the ceiling but never raise it above the server-configured maximum, and never + below the floor. """ if max_bytes is None or max_bytes <= 0: return cfg.mcp_max_response_bytes @@ -288,7 +290,7 @@ def _recompute_semantic_results(payload: dict[str, Any]) -> None: def _resolve_repo_id(engine: Engine, cfg: Settings, name: str) -> int | None: """One bounded ``repos.name -> id`` SELECT, fired only when byte-budget truncation of - ``search_code`` needs to synthesize a resume cursor (D3) -- the built payload resolves + ``search_code`` needs to synthesize a resume cursor -- the built payload resolves ``repo_id`` to a name and drops the id, so the MCP layer must resolve it back. A no-row result (reachable: the builder itself falls back to ``str(repo_id)`` when a repo @@ -313,7 +315,9 @@ def _make_search_code_truncator( engine: Engine, cfg: Settings ) -> Callable[[dict[str, Any], int, list[tuple[str, int]] | None], None]: """Build the ``search_code`` truncator: a closure over ``engine``/``cfg`` so an actual - byte-budget trim can synthesize ``next_cursor`` (D3). + byte-budget trim can synthesize ``next_cursor`` by resolving the last kept file's repo + name back to an id (the built payload only carries the name, not the id ``grep``'s cursor + needs). Pure tail-trim of ``files`` in the payload's existing ``(repo_id, path, content_sha)`` sort order (never reordered to protect symbol-bearing files) -- the kept files are always a @@ -323,6 +327,13 @@ def _make_search_code_truncator( ``[(content_sha, span_count), ...]`` captured by ``_shape_response`` BEFORE ``project_for_mcp`` strips ``content_sha``/``byte_ranges`` from ``payload["files"]`` -- trimming here is positional against that snapshot, not content-aware. + + Progress guarantee: whenever the untruncated payload had at least one file, at least one + file is ALWAYS kept, even if that single file's own serialized size alone exceeds + ``budget`` -- mirroring ``_shape_get_file_response``'s single-oversized-line edge. Without + this floor, a single huge file (e.g. one file with thousands of matches) would fit zero + items, yielding ``files: []`` with no ``next_cursor`` to resume from: an unrecoverable dead + end that contradicts the "always advance" contract this tool promises callers. """ def _truncate( @@ -346,27 +357,31 @@ def _apply(keep: int) -> None: envelope_overhead = len(json.dumps(payload)) - sum(len(json.dumps(f)) + 2 for f in files) budget_for_items = max(0, budget - envelope_overhead) keep = _fit_list(files, budget_for_items) + if keep == 0: + keep = 1 # progress guarantee -- see docstring; this page may exceed budget _apply(keep) payload["truncated"] = True if payload.get("truncation_reason") is None: payload["truncation_reason"] = "token_budget" - while keep > 0 and len(json.dumps(payload)) > budget: - keep //= 2 + # Safety net (mirrors _make_tail_trim_truncator's): re-verify against the real wire + # string and shrink further if the "+2" separator estimate under-counted. Floored at 1, + # never 0, so the progress guarantee above can never be undone here. + while keep > 1 and len(json.dumps(payload)) > budget: + keep = max(1, keep // 2) _apply(keep) + last_file = payload["files"][keep - 1] + last_sha = snapshot[keep - 1][0] if keep - 1 < len(snapshot) else None + repo_name = last_file.get("repo") + path = last_file.get("file") next_cursor: str | None = None - if keep > 0: - last_file = payload["files"][keep - 1] - last_sha = snapshot[keep - 1][0] if keep - 1 < len(snapshot) else None - repo_name = last_file.get("repo") - path = last_file.get("file") - if repo_name is not None and path is not None and last_sha: - repo_id = _resolve_repo_id(engine, cfg, repo_name) - if repo_id is not None: - next_cursor = service.encode_cursor( - FileCursor(repo_id=repo_id, path=path, content_sha=last_sha) - ) + if repo_name is not None and path is not None and last_sha: + repo_id = _resolve_repo_id(engine, cfg, repo_name) + if repo_id is not None: + next_cursor = service.encode_cursor( + FileCursor(repo_id=repo_id, path=path, content_sha=last_sha) + ) payload["next_cursor"] = next_cursor return _truncate @@ -396,7 +411,8 @@ def _shape_response( (pre- and post-projection serialized sizes, for ``_dispatch``'s telemetry line) and the pre-projection ``signals`` dict (so ``duration_ns`` observability survives projection dropping the field). Never raises: an irreducible over-budget envelope (a giant echoed - scalar with nothing left to trim) is returned flagged rather than erroring (Principle 4). + scalar with nothing left to trim) is returned flagged rather than raising -- forward + progress for the caller always wins over strict budget enforcement. """ pre_bytes = len(json.dumps(payload)) signals = _signals(payload) @@ -460,7 +476,7 @@ def _shape_get_file_response( ) -> tuple[str, dict[str, Any]]: """Slice ``payload["content"]`` to a ``start_line``-anchored page fitting ``budget``. - Per D2: splits with ``content.split("\\n")`` -- the SAME rule ``grep.py:436`` uses for + Splits with ``content.split("\\n")`` -- the SAME rule ``grep.py:436`` uses for ``search_code`` line numbers, so ``get_file`` pages stay congruent with search match lines on every input, including form feeds and ``U+2028``/``U+2029`` (which ``str.splitlines()`` would wrongly treat as line breaks for this purpose). Reassembly re-appends ``"\\n"`` to @@ -470,9 +486,9 @@ def _shape_get_file_response( A miss (``found: false``) never truncates: ``start_line`` echoes, ``next_start_line`` is ``null``, ``truncated`` is ``False``. On a hit, at least one line is always returned even if - it alone exceeds ``budget`` (Principle 4's progress guarantee outranks strict enforcement - for that documented, degenerate edge -- a single line larger than the whole budget); - ``next_start_line`` still advances past it. + it alone exceeds ``budget`` -- always making forward progress for the caller outranks strict + budget enforcement for that documented, degenerate edge (a single line larger than the whole + budget); ``next_start_line`` still advances past it. """ pre_bytes = len(json.dumps(payload)) signals = _signals(payload) @@ -732,7 +748,11 @@ async def search_code( ``truncation_reason: "token_budget"`` with a synthesized ``next_cursor`` so a caller can keep paging through the CONTENT matches losslessly; a ``sym:`` query's page-1-only symbol definitions that land in a truncated tail are lost from that traversal (flagged, not - silently dropped) since the symbol leg never re-runs on a continuation page. + silently dropped) since the symbol leg never re-runs on a continuation page. At least one + file is always kept and a ``next_cursor`` always synthesized when there was at least one + file to begin with -- a single file whose own serialized size alone exceeds ``max_bytes`` + (e.g. one file with thousands of matches) is still returned alone, with that one response + exceeding the budget, rather than coming back as an empty, unresumable dead end. """ lc = ctx.request_context.lifespan_context engine, cfg = lc["engine"], lc["config"] diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index d9e82ee..93f5a92 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -1308,23 +1308,12 @@ def _build() -> dict[str, Any]: assert "tool=search_code" in line assert "query_too_broad" in line assert "limiter_borrowed=" in line # pool/limiter saturation signal is wired - # AC5 / Step 4: pre- and post-shaping response byte sizes ship in the same log line. + # Pre- and post-shaping response byte sizes ship in the same log line (see + # tests/unit/test_mcp_shaping.py for the duration_ns-survives-projection pin). assert "response_bytes_pre=" in line assert "response_bytes=" in line -@pytest.mark.observability -@pytest.mark.asyncio -async def test_signals_log_includes_duration_ns_before_projection() -> None: - # duration_ns is read into _signals() BEFORE project_for_mcp drops it from the wire - # payload -- so log-line observability survives even though MCP callers never see the - # field itself. - payload = {"duration_ns": 123456, "files": []} - body, log_fields = main._shape_response("search_code", payload, 100_000, None) - assert log_fields["signals"]["duration_ns"] == 123456 - assert '"duration_ns"' not in body - - @pytest.mark.observability def test_signals_log_includes_both_flags() -> None: # Read straight off the payload dict, so a filter-only query is diagnosable from the logs diff --git a/tests/unit/test_mcp_shaping.py b/tests/unit/test_mcp_shaping.py index 279b7e8..ddc5169 100644 --- a/tests/unit/test_mcp_shaping.py +++ b/tests/unit/test_mcp_shaping.py @@ -646,6 +646,64 @@ async def test_search_code_cursor_traversal_equals_uncapped( assert not (page1_paths & page2_paths) # no double-counting across pages +@pytest.mark.unit +@pytest.mark.asyncio +async def test_search_code_single_oversized_file_progress_guarantee( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # Regression for the blocking defect found in review: a single file whose OWN serialized + # size already exceeds the budget must still come back as exactly that one file (never an + # empty, unresumable `files: []` with `next_cursor: null`), flagged truncated, with a + # next_cursor that lets a caller advance past it -- mirroring get_file's single-oversized- + # line edge. Under the OLD code, _fit_list returned keep=0 here, `next_cursor` stayed None, + # and the response was a silent, permanent dead end. + huge_file = FileMatches( + repo_id=7, + path="file0_huge.py", + lang="python", + content_sha="sha-huge", + branches=("main",), + line_matches=tuple( + LineMatch(i, f"needle match number {i} " + "x" * 40, ((0, 6),)) for i in range(1, 501) + ), + ) + small_files = [ + FileMatches( + repo_id=7, + path=f"file{i}_small.py", + lang="python", + content_sha=f"sha-small-{i}", + branches=("main",), + line_matches=(LineMatch(1, f"needle in small file {i}", ((0, 6),)),), + ) + for i in range(1, 3) + ] + all_files = [huge_file, *small_files] + monkeypatch.setattr(service, "grep_search", _make_grep_stub(all_files)) + monkeypatch.setattr(service, "symbol_search", lambda *a, **k: _no_sym()) + engine = _FakeSearchEngine({7: "acme/widgets"}) + ctx = _FakeLifespanContext(engine, _cfg()) + + # A budget far smaller than the huge file's own serialized size (the file's ~500 matches + # alone serialize to well over 10x this) but the file still must come back, not be dropped. + out1 = await main.search_code("needle", ctx, max_bytes=2000) # type: ignore[arg-type] + page1 = json.loads(out1) + + assert page1["truncated"] is True + assert page1["truncation_reason"] == "token_budget" + assert [f["file"] for f in page1["files"]] == ["file0_huge.py"] # exactly 1 file kept + assert len(out1) > 2000 # the documented edge: this one response exceeds the budget + cursor = page1["next_cursor"] + assert cursor is not None, "progress guarantee: a resume cursor must still be synthesized" + + # Traversal past the oversized file must work: the remaining (small) files come back. + out2 = await main.search_code( # type: ignore[arg-type] + "needle", ctx, cursor=cursor, max_bytes=100_000 + ) + page2 = json.loads(out2) + assert {f["file"] for f in page2["files"]} == {"file1_small.py", "file2_small.py"} + + @pytest.mark.unit @pytest.mark.asyncio async def test_search_code_cursor_no_row_degrades_without_cursor_or_raise(