diff --git a/README.md b/README.md index f892ac4..1365d16 100644 --- a/README.md +++ b/README.md @@ -223,10 +223,11 @@ Two different engines run in sequence, and neither is zoekt's RE2: 1. Postgres POSIX ARE (`~` / `~*`) selects which files match. An invalid pattern (any polarity) surfaces as the recoverable `regex_invalid` payload field carrying the Postgres message, not a parse error and not a failed tool call. -2. Python `re` rescans those files to produce the highlighted line matches. +2. Python `regex` (a `re`-compatible superset, still not RE2) rescans those files to produce + the highlighted line matches. The practical consequences: `^` and `$` are line anchors and `.` never crosses lines; -a Postgres-valid pattern that Python `re` rejects contributes no highlights *for that atom* +a Postgres-valid pattern that Python `regex` rejects contributes no highlights *for that atom* and sets `regex_incompatible`, so a single-atom query of that shape returns nothing while other atoms in the same query still match; and case folding can disagree on non-ASCII pairs (`ß`/`SS`, Turkish dotless `ı`). ASCII is unaffected. @@ -237,9 +238,13 @@ pattern like `/^/`, returns nothing even though the SQL predicate matched. `sym: exception: `search_code` runs a separate symbol leg, so `sym:Handler` returns definitions (carrying `symbols` and a `line`, but empty `text`) rather than falling into this hole. -One more V1 limitation worth knowing before pointing agents at it: a -catastrophic-backtracking regex on a single under-cap file runs unbounded in Python and can -stall the server. `statement_timeout` bounds the database, not the rescan. +Python-side matching is bounded by a per-request **match budget** (default 2000ms, +configurable via `CODE_SEARCH_MATCH_BUDGET_MS`) so a catastrophic-backtracking regex on a +single under-cap file cannot pin a worker: when the budget trips, scanning stops and the +result comes back with `truncated=True` and `truncation_reason="match_budget"`. This +complements `statement_timeout` (which bounds the database, not the rescan). The `regex` +module releases the GIL while matching, so a budgeted pathological scan does not block the +event loop. ## MCP tools diff --git a/app/config.py b/app/config.py index 7ee0256..a26bfbb 100644 --- a/app/config.py +++ b/app/config.py @@ -44,6 +44,11 @@ class Settings(BaseSettings): # Aggregate content bytes grep pulls/scans per request (memory bound). max_content_bytes: int = 8 * 1024 * 1024 + # Per-request wall-clock bound on Python-side pattern matching in grep (regex module's + # timeout=). Complements statement_timeout_ms (DB time) and max_content_bytes (bytes) as + # the third leg of the per-request resource triangle. + match_budget_ms: int = 2000 + # Default and hard-cap for search_code(limit): <=0 -> row_limit; > max -> max_row_limit. row_limit: int = 200 max_row_limit: int = 1000 diff --git a/app/main.py b/app/main.py index 548d79d..98cf415 100644 --- a/app/main.py +++ b/app/main.py @@ -9,8 +9,10 @@ 1. **No OBO / token forwarding.** This is a single shared-corpus service principal; there is no per-user ``X-Forwarded-Access-Token`` path. 2. **Blocking work runs off the event loop.** ``grep_search`` runs synchronous SQL *and* a - Python ``re`` rescan whose CPU is uncapped (``grep.py:45-49``); running it inline in an - async handler would stall ``/health``/``/ready`` and every concurrent request. Each tool + Python ``regex`` rescan bounded per request by a match budget (``CODE_SEARCH_MATCH_BUDGET_MS``, + default 2000ms; a trip flags ``truncated``/``truncation_reason="match_budget"``); running it + inline in an async handler would stall ``/health``/``/ready`` and every concurrent request. + The ``regex`` module releases the GIL while matching, but the SQL leg still blocks, so each tool body is dispatched to a worker thread via ``anyio.to_thread.run_sync`` under a pool-sized ``CapacityLimiter(5)`` so in-flight blocking calls never oversubscribe the 5-conn pool. 3. **The engine is a process-scoped module singleton, not lifespan-owned.** A stateful @@ -103,6 +105,8 @@ def _signals(payload: dict[str, Any]) -> dict[str, Any]: """Extract the recoverable-signal fields for the observability log line.""" return { "truncated": payload.get("truncated"), + # Distinguishes which cap/budget tripped -- byte_cap/row_cap/match_budget -- from the logs. + "truncation_reason": payload.get("truncation_reason"), "query_too_broad": payload.get("query_too_broad"), "query_parse_error": payload.get("query_parse_error"), # Without this, a Postgres-rejected regex is log-indistinguishable from a genuine @@ -247,17 +251,20 @@ async def search_code( ``commit:`` to ``query``. ``limit`` caps the number of files scanned (clamped to a server maximum). Recoverable conditions surface as fields (``query_parse_error``, ``query_too_broad``, ``truncated``, ``regex_incompatible``, ``regex_invalid``, - ``no_content_atom``, ``zero_width_only_atoms``). The middle two explain an empty result - that is NOT a true negative: ``no_content_atom`` means the query carried no affirmative - content atom to highlight -- either a filter-only query (e.g. ``lang:go`` alone) or one - that is entirely negated (e.g. ``-foo`` alone: an exclusion is never a highlight) -- and + ``no_content_atom``, ``zero_width_only_atoms``). ``truncated``/``truncation_reason`` also + cover the Python match-budget trip (``truncation_reason="match_budget"``): a pathological + pattern that exhausts the per-request CPU budget stops scanning and returns a flagged + partial result rather than pinning a worker. The middle two explain an empty result that is + NOT a true negative: ``no_content_atom`` means the query carried no affirmative content atom + to highlight -- either a filter-only query (e.g. ``lang:go`` alone) or one that is entirely + negated (e.g. ``-foo`` alone: an exclusion is never a highlight) -- and ``zero_width_only_atoms`` means every atom it carried matches zero-width (e.g. ``/^/``). A query mixing content with an exclusion (e.g. ``foo -bar``) highlights only the affirmative term; excluded terms never appear as matches. ``no_content_atom`` does not distinguish "no atom at all" from "fully negated" -- recover which one it was from the echoed ``query`` field, if it matters. ``regex_invalid`` carries the Postgres error message when a ``/regex/``, ``repo:``, ``file:``, or ``sym:`` pattern is not a valid Postgres POSIX ARE - (e.g. ``/[/``) -- distinct from ``regex_incompatible``, which means Python ``re`` (not + (e.g. ``/[/``) -- distinct from ``regex_incompatible``, which means Python ``regex`` (not Postgres) rejected an otherwise-valid pattern and only degrades highlighting. """ lc = ctx.request_context.lifespan_context diff --git a/app/requirements.txt b/app/requirements.txt index 060d7c8..c69a76b 100644 --- a/app/requirements.txt +++ b/app/requirements.txt @@ -2,7 +2,7 @@ # uv export --no-dev --no-hashes --no-emit-project -o app/requirements.txt alembic==1.18.5 # via databricks-code-search -annotated-types==0.7.0 +annotated-types==0.8.0 # via pydantic anyio==4.14.2 # via @@ -16,7 +16,7 @@ attrs==26.1.0 # via # jsonschema # referencing -certifi==2026.6.17 +certifi==2026.7.22 # via # httpcore # httpx @@ -33,11 +33,11 @@ cryptography==49.0.0 # via # google-auth # pyjwt -databricks-sdk==0.121.0 +databricks-sdk==0.122.0 # via databricks-code-search -google-auth==2.56.0 +google-auth==2.56.2 # via databricks-sdk -greenlet==3.5.3 ; platform_machine == 'AMD64' or platform_machine == 'WIN32' or platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'ppc64le' or platform_machine == 'win32' or platform_machine == 'x86_64' +greenlet==3.5.4 ; platform_machine == 'AMD64' or platform_machine == 'WIN32' or platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'ppc64le' or platform_machine == 'win32' or platform_machine == 'x86_64' # via sqlalchemy h11==0.16.0 # via @@ -110,6 +110,8 @@ referencing==0.37.0 # via # jsonschema # jsonschema-specifications +regex==2026.7.19 + # via databricks-code-search requests==2.34.2 # via databricks-sdk rpds-py==2026.6.3 @@ -120,7 +122,7 @@ sqlalchemy==2.0.51 # via # alembic # databricks-code-search -sse-starlette==3.4.5 +sse-starlette==3.4.6 # via mcp starlette==1.3.1 # via @@ -130,7 +132,7 @@ tree-sitter==0.26.0 # via # databricks-code-search # tree-sitter-language-pack -tree-sitter-language-pack==1.13.1 +tree-sitter-language-pack==1.13.3 # via databricks-code-search typing-extensions==4.16.0 # via diff --git a/app/search/AGENTS.md b/app/search/AGENTS.md index d10effe..ce2da17 100644 --- a/app/search/AGENTS.md +++ b/app/search/AGENTS.md @@ -4,12 +4,12 @@ # app/search ## Purpose -The impure serve-side execution layer sitting between the pure `app/query/` seams and the payload builders in `app/service.py`. `grep.py` runs the compiler's candidate query, then streams content one row at a time and rescans it with Python `re` for per-line zoekt-shaped highlights; `symbols.py` answers `sym:` queries with actual definitions (name/kind/line) the highlight-driven grep path cannot return; `semantic.py` is the natural-language companion — vector-ANN + BM25 legs over `chunks` fused with reciprocal-rank fusion (RRF, k=60) via the Lakebase `lakebase_ann`/`lakebase_bm25` operators; `errors.py` is the shared mapper from a raw Postgres `DBAPIError` to a typed recoverable error — `QueryTooBroadError` (statement_timeout) or `RegexInvalidError` (Postgres-invalid regex, issue #75). +The impure serve-side execution layer sitting between the pure `app/query/` seams and the payload builders in `app/service.py`. `grep.py` runs the compiler's candidate query, then streams content one row at a time and rescans it with Python `regex` for per-line zoekt-shaped highlights; `symbols.py` answers `sym:` queries with actual definitions (name/kind/line) the highlight-driven grep path cannot return; `semantic.py` is the natural-language companion — vector-ANN + BM25 legs over `chunks` fused with reciprocal-rank fusion (RRF, k=60) via the Lakebase `lakebase_ann`/`lakebase_bm25` operators; `errors.py` is the shared mapper from a raw Postgres `DBAPIError` to a typed recoverable error — `QueryTooBroadError` (statement_timeout) or `RegexInvalidError` (Postgres-invalid regex, issue #75). ## Key Files | File | Description | |------|-------------| -| `grep.py` | `grep_search(conn, query, ...) -> GrepResult`: compile → candidate ids → content fetch with `yield_per=1` (one-row server-side cursor; bare `stream_results` does NOT bound memory) → `extract_line_matches` per line. Bounds: transaction-local `statement_timeout` (cancel → `QueryTooBroadError`), aggregate `max_content_bytes` cap (trip → `truncated`/`"byte_cap"`); a Postgres-invalid regex (any polarity) raises `RegexInvalidError`. Emits UTF-8 line-local half-open `byte_ranges`, `regex_incompatible` (NOT-RE2 degradation), the shape flags `no_content_atom` / `zero_width_only_atoms` (mutually exclusive by construction; zero-width proof via guarded `re._parser.getwidth()`), and `FileCursor`-based pagination (`_Unset` sentinel: cursor kwarg supplied at all = pagination mode) | +| `grep.py` | `grep_search(conn, query, ...) -> GrepResult`: compile → candidate ids → content fetch with `yield_per=1` (one-row server-side cursor; bare `stream_results` does NOT bound memory) → `extract_line_matches` per line. Bounds: transaction-local `statement_timeout` (cancel → `QueryTooBroadError`), aggregate `max_content_bytes` cap (trip → `truncated`/`"byte_cap"`), per-request `match_budget_ms` on Python-side matching via the `regex` module's `timeout=` (trip → `truncated`/`"match_budget"`); a Postgres-invalid regex (any polarity) raises `RegexInvalidError`. Emits UTF-8 line-local half-open `byte_ranges`, `regex_incompatible` (NOT-RE2 degradation), the shape flags `no_content_atom` / `zero_width_only_atoms` (mutually exclusive by construction; zero-width proof via guarded `re._parser.getwidth()`), and `FileCursor`-based pagination (`_Unset` sentinel: cursor kwarg supplied at all = pagination mode) | | `symbols.py` | `symbol_search(conn, query, ...) -> SymbolResult`: deliberately TWO queries — the full-AST `compile_query` picks eligible file ids, then a pure projection returns matching `symbols` rows. One joined query would let SQLAlchemy auto-correlate the compiler's `sym:`/`repo:` subqueries against outer rows (silent mis-scope). No `sym:` atom → short-circuit `no_symbol_atom=True` with zero DB hits. Pure SQL: none of grep's NOT-RE2/CPU caveats apply | | `semantic.py` | `_semantic_search_payload`: flag-off short-circuit FIRST (no engine/SDK), then `app.query.semantic_filters.split_semantic_query` (pure, no DB) splits `repo:`/`file:`/`lang:`/`branch:` filters from natural-language prose — `query_parse_error` / `unsupported_filter` (`sym:`/`case:`/`commit:`/regex, remedy in `reason`) / `nothing_to_embed` (branched wording) all return BEFORE the schema probe or the embedder — then a `to_regclass('chunks')` schema probe (→ `semantic_schema_missing` payload, not a 500), embed the RESIDUAL text OUTSIDE any transaction, then the RRF query. `build_hybrid_rrf_sql(filters, branch)` / `filter_params(filters, branch)`: both derive from one shared `_normalized_filter_state` so predicates and binds can never drift; filters compile to numbered `sem_*` binds inside BOTH leg CTEs' `extra_where`; `branch` is unified with `branch:` atoms into ONE normalized, sorted branch list (no separate `:branch` bind). Two rank CTEs whose inner `ORDER BY LIMIT :topk` must stay a SINGLE expression (a second sort key silently collapses the lakebase index path to a seq scan); the query vector is bound as `(:qvec)::vector` via `format_vector_literal` (repr-based), never interpolated. The outer SELECT recomputes `c.embedding <=> (:qvec)::vector AS cosine_distance` for every fused row (Decision B1) — surfaced as `similarity = 1 - cosine_distance`, `None` for NULL embeddings. Process-scoped lazy `get_embedder` singleton mirrors `app.main.get_engine` | | `errors.py` | `QueryTooBroadError` / `RegexInvalidError` + `reraise_or_recoverable(error: DBAPIError)`: maps `psycopg.errors.QueryCanceled` (statement_timeout) to the former and `psycopg.errors.InvalidRegularExpression` (Postgres-invalid regex, issue #75) to the latter, so grep/symbols/references/semantic raise identically without importing each other; anything else re-raises unchanged | @@ -23,7 +23,7 @@ The impure serve-side execution layer sitting between the pure `app/query/` seam - **Total failure raises, partial success flags.** A statement_timeout on the candidate query raises `QueryTooBroadError` (an empty return would be a lie); a tripped byte/row cap returns `truncated=True` + `truncation_reason`. Grep's shape flags are per-leg facts — the envelope in `app/service.py` ANDs in "the symbol leg did not answer"; do not special-case `SymbolFilter` inside grep. - **`Not` (lexical negation) contributes nothing to highlighting or projection.** `grep._collect_matchers` skips `Not` subtrees (a `-foo` exclusion is SQL-only, never a green highlight) and does NOT recurse — so a purely-negated broken regex (`-/[/`) builds no pattern and never flags `regex_incompatible` (nothing was meant to highlight). `symbols._collect_symbol_patterns` skips `Not` too, so a `-sym:foo`-only query is `no_symbol_atom=True`. The compiler still lowers the exclusion in SQL regardless; only Python-side highlight/definition collection ignores it. A Postgres-invalid regex still reaches the DB either way (negated or not); it now raises a typed `RegexInvalidError` (issue #75 — `errors.py`'s `reraise_or_recoverable`), which `app/service.py` maps to the recoverable `regex_invalid` payload field rather than an uncaught fault. - **`no_content_atom`/`no_symbol_atom` mean "no affirmative atom", not "no atom".** A filter-only query (`lang:go`) and a fully-negated one (`-foo` alone) both collect zero patterns and are reported identically by these flags — deliberately, so one flag covers both shapes rather than adding a second envelope key. A caller recovers which shape it actually was from the echoed `query` field, not from a distinguishing signal. Same for `no_symbol_atom` and `-sym:foo` alone. -- **Documented V1 caveats (do not "fix" silently):** Python `re` is not Postgres POSIX ARE (uncompilable atom → skipped + `regex_incompatible`); non-ASCII case-folding divergence is unsignalled; Python regex CPU is uncapped (catastrophic backtracking holds the GIL — real fix is an RE2 binding). +- **Documented V1 caveats (do not "fix" silently):** Python `regex` (a `re`-compatible superset, still not RE2) is not Postgres POSIX ARE (uncompilable atom → skipped + `regex_incompatible`; fires less often than under `re` since `regex` accepts a superset, never more); non-ASCII case-folding divergence is unsignalled. Python regex CPU is bounded per request by `match_budget_ms` (default 2000ms / `CODE_SEARCH_MATCH_BUDGET_MS`): a catastrophic-backtracking match trips the budget → scanning stops, fully-scanned files are kept, `truncated=True`/`"match_budget"`; the `regex` module releases the GIL while matching so the event loop stays responsive. Residual risk: an attacker can still burn `match_budget_ms` × concurrent slots of CPU per volley (rate limiting is out of scope). - **Semantic determinism rules:** `row_number()` ties break on `id` and the fused `ORDER BY rrf DESC, id LIMIT` tiebreak is load-bearing, but the inner leg `ORDER BY` must NEVER gain an `, id` — see `_leg_cte`'s docstring. The ANN leg's `c.embedding IS NOT NULL` guard keeps NULL embeddings from earning RRF credit. Filter predicates (`repo:`/`file:`/`lang:`/`branch:`, from `app.query.semantic_filters`) live ONLY in each leg's inner `extra_where` — never the `ORDER BY` — so a filtered query keeps the same single-expression index path as an unfiltered one; `similarity` is recomputed in the OUTER select (never threaded through a leg CTE), so it reaches BM25-only rows the ANN leg's `:topk` window never touched. **Cross-module invariant:** `app.query.semantic_filters` derives field names by inverting the parser's `_FIELD_KINDS` and recomputes value spans with `_read_field_value` — any change to the parser's field scanning must run `tests/unit/test_semantic_filters.py` (see `app/query/AGENTS.md`). - The branch predicate (`coalesce(r.default_branch,'HEAD') = ANY(f.branches)` / `f.branches @> ARRAY[:branch]`) must stay byte-identical to the compiler/`get_file`/0003-backfill sites. - `re._parser` access stays a guarded `getattr` inside `_zero_width_only_atoms` — a module-scope import failing on a future CPython would take down the whole MCP server; `test_getwidth_private_api_canary` trips loudly instead. diff --git a/app/search/grep.py b/app/search/grep.py index 4ccfed4..94bafab 100644 --- a/app/search/grep.py +++ b/app/search/grep.py @@ -27,15 +27,17 @@ Caveats (load-bearing, documented, never silently wrong): -* **NOT RE2.** Python ``re`` is not Postgres POSIX ARE, and matching here is line-oriented: - ``^``/``$`` are line anchors, ``.`` never crosses lines, and cross-line constructs (e.g. - ``(?s)...``) do not span lines. A Postgres-valid regex that Python ``re`` rejects is - skipped (that atom contributes no highlights) and ``regex_incompatible`` is set. The SQL - predicate already selected the file; grep only degrades the *highlighting*. Case folding - can also diverge: ``re.IGNORECASE`` (Python Unicode folding) and Postgres ``lower()`` do - not agree on every non-ASCII pair (e.g. ``ß``/``SS``, Turkish dotless ``i``), so a file - the SQL predicate matched case-insensitively may yield zero Python highlights and drop - out. ASCII is unaffected. +* **NOT RE2.** Python's ``regex`` module (a ``re``-compatible superset, still not RE2) is not + Postgres POSIX ARE, and matching here is line-oriented: ``^``/``$`` are line anchors, ``.`` + never crosses lines, and cross-line constructs (e.g. ``(?s)...``) do not span lines. A + Postgres-valid regex that ``regex`` rejects is skipped (that atom contributes no highlights) + and ``regex_incompatible`` is set. Because ``regex`` accepts a strict superset of what stdlib + ``re`` accepted, this is a pure improvement: ``regex_incompatible`` fires less often than + under ``re``, never more. The SQL predicate already selected the file; grep only degrades + the *highlighting*. Case folding can also diverge: ``regex.IGNORECASE`` (Python Unicode + folding) and Postgres ``lower()`` do not agree on every non-ASCII pair (e.g. ``ß``/``SS``, + Turkish dotless ``i``), so a file the SQL predicate matched case-insensitively may yield zero + Python highlights and drop out. ASCII is unaffected. * **Highlight-driven results.** A file appears only if at least one line produces a non-empty highlight span, so two query shapes the SQL predicate does match still return no files. Both are announced by name rather than returning a silent empty result @@ -57,10 +59,22 @@ a file the SQL predicate matched case-insensitively that yields zero Python highlights drops out with patterns present, non-empty, and of non-zero width -- so it is still entirely unsignalled. Fixing that needs a new provable signal, not one of these two. -* **Uncapped Python CPU.** The byte cap bounds memory and aggregate bytes scanned but NOT - CPU/wall-clock: a catastrophic-backtracking ``re`` pattern on a single under-cap file runs - unbounded, holds the GIL, and can starve the app. ``statement_timeout`` does not cover - Python work. No guard for this ships yet; the real fix is an RE2 binding. +* **Bounded Python CPU (match budget).** The byte cap bounds memory and aggregate bytes + scanned but not CPU/wall-clock; a per-request match budget (the ``regex`` module's + ``timeout=``) covers that third leg. Its value comes from ``Settings.match_budget_ms`` / + ``CODE_SEARCH_MATCH_BUDGET_MS`` (default 2000ms), threaded through ``grep_search``. Trip + semantics: scanning stops, fully-scanned files are kept, and the result is flagged + ``truncated=True`` + ``truncation_reason="match_budget"`` (a budgeted result never + masquerades as complete). Cursor semantics have two cases -- a *pre-file* trip (the deadline + is already past when the next file's row is dequeued) leaves that file unconsumed, so + ``next_cursor`` resumes at it fresh (a later fresh budget may clear it); a *mid-file* trip + (the ``regex`` engine's ``timeout=`` fires while scanning a file) treats that file as + consumed, discards its partial matches, and ``next_cursor`` steps past it so pagination + never stalls forever on one pathological file. The ``regex`` module releases the GIL while + matching, so even a budgeted pathological match keeps the event loop responsive rather than + pinning it. Residual risk: an attacker can still burn up to ``match_budget_ms`` x + concurrent-request-slots of CPU per volley -- availability is preserved (GIL released), but + request rate limiting is out of scope here. * **Per-file memory** relies on the indexer's per-file byte cap (``MAX_FILE_BYTES``) keeping any single ``content`` bounded; ``File.size`` is nullable/unpopulated here, so a ``size`` pre-filter is intentionally NOT used. @@ -86,10 +100,12 @@ from __future__ import annotations import re +import time from collections.abc import Sequence from dataclasses import dataclass from typing import NamedTuple, assert_never +import regex from sqlalchemy import Connection, select, text, tuple_ from sqlalchemy.exc import DBAPIError @@ -117,6 +133,18 @@ DEFAULT_MAX_CONTENT_BYTES = 8 * 1024 * 1024 # Per-request DB-time bound; a cancellation surfaces as QueryTooBroadError. DEFAULT_STATEMENT_TIMEOUT_MS = 5000 +# Per-request wall-clock bound on Python-side pattern matching (regex module's timeout=); +# the third leg of the per-request resource triangle beside DB time and bytes scanned. +DEFAULT_MATCH_BUDGET_MS = 2000 + + +class MatchBudgetExceeded(Exception): + """The per-request match budget tripped mid-scan (module-private, never crosses out). + + Raised by :func:`extract_line_matches` when a supplied ``deadline`` is reached; caught + inside :func:`grep_search`, which turns it into a ``truncation_reason="match_budget"`` + partial result. It never reaches ``app/search/errors.py`` or a caller. + """ # --------------------------------------------------------------------------- contract @@ -203,7 +231,7 @@ class GrepResult: files: tuple[FileMatches, ...] # in (repo_id, path, content_sha) order truncated: bool # byte cap OR (row cap tripped AND no cursor kwarg was supplied) - truncation_reason: str | None # "byte_cap" | "row_cap" | None + truncation_reason: str | None # "byte_cap" | "row_cap" | "match_budget" | None regex_incompatible: bool # some Regex atom failed Python re.compile no_content_atom: bool # no affirmative content atom to highlight (filter-only OR fully # negated, e.g. ``lang:go`` OR ``-foo`` alone); nothing compiled away either way @@ -223,10 +251,10 @@ class GrepResult: # ----------------------------------------------------------------------- pure helpers -def _collect_matchers(node: Node, flags: int, patterns: list[re.Pattern[str]]) -> bool: +def _collect_matchers(node: Node, flags: int, patterns: list[regex.Pattern[str]]) -> bool: """Append every affirmative Substring/Regex leaf's compiled pattern to ``patterns``. - Returns True if any (affirmative) Regex leaf failed Python ``re.compile`` (NOT-RE2 + Returns True if any (affirmative) Regex leaf failed ``regex.compile`` (NOT-RE2 degradation). Filters (repo/path/lang/sym) contribute no patterns. A :class:`Not` subtree contributes nothing and is not recursed into: grep highlights only @@ -239,12 +267,12 @@ def _collect_matchers(node: Node, flags: int, patterns: list[re.Pattern[str]]) - """ match node: case Substring(value=value): - patterns.append(re.compile(re.escape(value), flags)) + patterns.append(regex.compile(regex.escape(value), flags)) return False case Regex(pattern=pattern): try: - patterns.append(re.compile(pattern, flags)) - except re.error: + patterns.append(regex.compile(pattern, flags)) + except regex.error: return True return False case Not(): @@ -267,21 +295,21 @@ def _collect_matchers(node: Node, flags: int, patterns: list[re.Pattern[str]]) - assert_never(node) -def _build_matchers(node: Node, case_sensitive: bool) -> tuple[list[re.Pattern[str]], bool]: +def _build_matchers(node: Node, case_sensitive: bool) -> tuple[list[regex.Pattern[str]], bool]: """Collect every Substring/Regex leaf (any And/Or nesting) into compiled patterns. - Filters contribute none. Substring -> ``re.compile(re.escape(value))``; Regex -> - ``re.compile(pattern)`` catching ``re.error`` (skip that atom, flag incompatible). - ``flags = re.IGNORECASE if not case_sensitive else 0``. Returns + Filters contribute none. Substring -> ``regex.compile(regex.escape(value))``; Regex -> + ``regex.compile(pattern)`` catching ``regex.error`` (skip that atom, flag incompatible). + ``flags = regex.IGNORECASE if not case_sensitive else 0``. Returns ``(patterns, regex_incompatible)``. Pure -- no DB import. """ - flags = re.IGNORECASE if not case_sensitive else 0 - patterns: list[re.Pattern[str]] = [] + flags = regex.IGNORECASE if not case_sensitive else 0 + patterns: list[regex.Pattern[str]] = [] regex_incompatible = _collect_matchers(node, flags, patterns) return patterns, regex_incompatible -def _no_content_atom(patterns: Sequence[re.Pattern[str]], regex_incompatible: bool) -> bool: +def _no_content_atom(patterns: Sequence[regex.Pattern[str]], regex_incompatible: bool) -> bool: """True when the query carries no affirmative content atom to highlight. ``patterns`` is empty for two structurally different queries, both reported identically @@ -301,7 +329,9 @@ def _no_content_atom(patterns: Sequence[re.Pattern[str]], regex_incompatible: bo return not patterns and not regex_incompatible -def _zero_width_only_atoms(patterns: Sequence[re.Pattern[str]], regex_incompatible: bool) -> bool: +def _zero_width_only_atoms( + patterns: Sequence[regex.Pattern[str]], regex_incompatible: bool +) -> bool: """True when every content atom provably matches zero-width, so nothing can highlight. ``re._parser.parse(src).getwidth()`` returns ``(min_width, max_width)``; a ``max_width`` @@ -379,7 +409,9 @@ def _char_to_byte_ranges(line: str, spans: list[tuple[int, int]]) -> tuple[tuple return tuple(ranges) -def extract_line_matches(content: str, patterns: Sequence[re.Pattern[str]]) -> list[LineMatch]: +def extract_line_matches( + content: str, patterns: Sequence[regex.Pattern[str]], *, deadline: float | None = None +) -> list[LineMatch]: """Extract per-line matches from ``content`` for the given compiled ``patterns``. Splits on ``"\\n"`` (1-based line numbers) and strips one trailing ``"\\r"`` per line @@ -388,6 +420,15 @@ def extract_line_matches(content: str, patterns: Sequence[re.Pattern[str]]) -> l any atom are merged into a sorted, non-overlapping set; merged char endpoints are converted to UTF-8 byte offsets. Only lines with >=1 span produce a :class:`LineMatch`. Pure -- no DB import. + + ``deadline`` is an absolute ``time.monotonic()`` timestamp bounding Python-side CPU. When + ``None`` (the default), matching is unbudgeted and every ``finditer`` call is issued with + no ``timeout=`` -- existing direct callers/tests pay no timeout-checking overhead and see + identical behaviour. When set, each pattern is matched with the remaining time as the + ``regex`` module's ``timeout=``; if the deadline is already past before a pattern runs, or + the ``regex`` engine trips its own ``TimeoutError`` mid-iteration on a pathological match, + this raises :class:`MatchBudgetExceeded` (the ``regex`` module releases the GIL while + matching, so a budgeted pathological scan does not starve the event loop). """ if not patterns: return [] @@ -396,9 +437,21 @@ def extract_line_matches(content: str, patterns: Sequence[re.Pattern[str]]) -> l line = raw_line[:-1] if raw_line.endswith("\r") else raw_line spans: list[tuple[int, int]] = [] for pattern in patterns: - for m in pattern.finditer(line): - if m.end() > m.start(): # drop zero-width matches - spans.append((m.start(), m.end())) + if deadline is None: + # Unbudgeted path: no timeout= kwarg, no per-call clock read. + finditer = pattern.finditer(line) + else: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise MatchBudgetExceeded() + finditer = pattern.finditer(line, timeout=remaining) + try: + for m in finditer: + if m.end() > m.start(): # drop zero-width matches + spans.append((m.start(), m.end())) + except TimeoutError: + # regex module raises builtin TimeoutError mid-iteration when timeout= trips. + raise MatchBudgetExceeded() from None if not spans: continue byte_ranges = _char_to_byte_ranges(line, _merge_spans(spans)) @@ -416,6 +469,7 @@ def grep_search( row_limit: int = DEFAULT_ROW_LIMIT, max_content_bytes: int = DEFAULT_MAX_CONTENT_BYTES, statement_timeout_ms: int = DEFAULT_STATEMENT_TIMEOUT_MS, + match_budget_ms: int = DEFAULT_MATCH_BUDGET_MS, cursor: FileCursor | None | _Unset = _UNSET, ) -> GrepResult: """Run a zoekt-style ``query`` and return file-grouped, per-line matches. @@ -428,10 +482,12 @@ def grep_search( A per-request ``statement_timeout`` bounds DB time (a cancellation raises :class:`QueryTooBroadError`); ``max_content_bytes`` bounds app memory / aggregate bytes - scanned (a cap sets ``truncated`` + ``truncation_reason``). A Postgres-invalid regex - pattern (any polarity) raises :class:`~app.search.errors.RegexInvalidError`. See the - module docstring for the NOT-RE2 and uncapped-Python-CPU caveats. Raises - ``QueryParseError`` on a malformed query (propagated from :func:`parse`). + scanned (a cap sets ``truncated`` + ``truncation_reason``); ``match_budget_ms`` bounds + Python-side pattern-matching wall-clock (a trip sets ``truncated`` + + ``truncation_reason="match_budget"``). A Postgres-invalid regex pattern (any polarity) + raises :class:`~app.search.errors.RegexInvalidError`. See the module docstring for the + NOT-RE2 and match-budget caveats. Raises ``QueryParseError`` on a malformed query + (propagated from :func:`parse`). ``cursor`` activates pagination mode when supplied at all, including ``None`` (page 1): a plain row-cap fill then reports ``truncated=False`` plus a non-null @@ -459,6 +515,7 @@ def grep_search( files: list[FileMatches] = [] byte_capped = False + budget_tripped = False last_candidate: FileCursor | None = None with conn.begin(): @@ -506,11 +563,32 @@ def grep_search( .order_by(File.repo_id, File.path, File.content_sha) .execution_options(yield_per=1) # one-row server-side cursor; NOT bare stream_results ) + # Computed AFTER the candidate query returns so candidate-selection DB time never eats + # the match budget. Per-row content-fetch latency (`yield_per=1`, below) DOES fall + # inside this deadline window by design -- it is interleaved with the Python-side + # matching the budget bounds, not excluded from it (see the pre-file check's comment + # for the consequence on the very first row of a page). + deadline = time.monotonic() + match_budget_ms / 1000.0 running = 0 result = conn.execute(content_stmt) try: for row in result: content = row.content or "" + # Pre-file budget check: mirrors the byte-cap break's position (before + # `last_candidate` advances for this row). A trip here leaves this file + # unconsumed. If a PRIOR row was already consumed this call, `next_cursor` + # (below) points at it and a resumed page re-fetches this file fresh. If this + # trips on the very FIRST row of a page -- reachable via slow content-fetch + # latency, not just an undersized budget, since fetch time counts against the + # deadline too -- `last_candidate` is still None; falling back to `next_cursor + # = resume` (unchanged) below re-issues this same page with a fresh budget + # instead of a `next_cursor=None` dead end that would silently drop the rest of + # the corpus. Only page 1 with no incoming `resume` and no row consumed yet has + # no cursor to fall back to; that residual case yields `next_cursor=None` + + # `truncated=True`, safely retried by re-issuing the identical query. + if time.monotonic() >= deadline: + budget_tripped = True + break # Char count is a valid lower bound on UTF-8 byte count, so this never # under-counts the cap; checked BEFORE .encode()/processing so the cap is a # real bound (overshoot <= one file) and avoids a transient copy of a huge file. @@ -532,7 +610,14 @@ def grep_search( # the resume key tracks candidates CONSUMED, not candidates EMITTED. See # FileCursor's docstring. last_candidate = FileCursor(row.repo_id, row.path, row.content_sha) - line_matches = extract_line_matches(content, patterns) + # Mid-file budget trip: `last_candidate` is ALREADY advanced past this row + # (above), so this file is CONSUMED and its partial line_matches are discarded + # entirely -- pagination steps past it rather than re-scanning it forever. + try: + line_matches = extract_line_matches(content, patterns, deadline=deadline) + except MatchBudgetExceeded: + budget_tripped = True + break if line_matches: files.append( FileMatches( @@ -549,10 +634,22 @@ def grep_search( finally: result.close() - truncated = byte_capped or (row_capped and not pagination_mode) + truncated = byte_capped or budget_tripped or (row_capped and not pagination_mode) legacy_row_capped = row_capped and not pagination_mode - reason = "byte_cap" if byte_capped else ("row_cap" if legacy_row_capped else None) - next_cursor = last_candidate if (row_capped or byte_capped) else None + reason = ( + "byte_cap" + if byte_capped + else ("match_budget" if budget_tripped else ("row_cap" if legacy_row_capped else None)) + ) + next_cursor: FileCursor | None + if budget_tripped and last_candidate is None and resume is not None: + # Pre-file trip on the very first row of a resumed page: nothing new was consumed this + # call, so falling back to the unchanged incoming `resume` cursor makes this page + # retryable with a fresh budget rather than reporting `next_cursor=None` (exhausted) + # while `truncated=True` -- see the pre-file check's comment above. + next_cursor = resume + else: + next_cursor = last_candidate if (row_capped or byte_capped or budget_tripped) else None return GrepResult( files=tuple(files), truncated=truncated, diff --git a/app/service.py b/app/service.py index 3c1bbed..992e15f 100644 --- a/app/service.py +++ b/app/service.py @@ -517,6 +517,7 @@ def search_code_payload( "row_limit": limit, "max_content_bytes": cfg.max_content_bytes, "statement_timeout_ms": cfg.statement_timeout_ms, + "match_budget_ms": cfg.match_budget_ms, } if pagination_mode: grep_kwargs["cursor"] = decoded_cursor diff --git a/pyproject.toml b/pyproject.toml index d0b7b7c..71cd407 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,6 +21,9 @@ dependencies = [ "tree-sitter-language-pack", "pydantic-settings", "pyyaml>=6.0", + # >=2021.11.10: floor for finditer(..., timeout=...) support, which app/search/grep.py's + # per-request match budget depends on. + "regex>=2021.11.10", ] # webui/ is a SEPARATE deploy artifact from the MCP app: it imports app.* from a @@ -45,6 +48,7 @@ dev = [ "httpx", "asgi-lifespan", "types-PyYAML", + "types-regex", ] [tool.setuptools.packages.find] diff --git a/tests/integration/test_grep.py b/tests/integration/test_grep.py index 038bce5..c7fd0e5 100644 --- a/tests/integration/test_grep.py +++ b/tests/integration/test_grep.py @@ -14,6 +14,7 @@ from __future__ import annotations +import time import uuid from collections.abc import Iterator from typing import NamedTuple @@ -689,3 +690,143 @@ def test_negated_content_predicate_does_not_resurrect_null_content_row(seeded: S # The two real go files both contain "foo" -> excluded by -foo, not by the null check -- # so candidate_paths is empty for an entirely different (and also correct) reason. assert candidate_paths == set() + + +# ------------------------------------------------------------------- 13. match budget (#38) +# +# The pathological regex below was reverified empirically against the installed `regex` +# version: `(?:(?:a{1,10}){1,10}){1,10}b` over a long run of `a`s with no trailing `b` +# catastrophically backtracks in Python. Postgres's NFA engine does NOT backtrack, so the +# file is selected cheaply as a candidate via the short `aaab` line (which the SQL predicate +# matches); the CPU blow-up only happens in the Python line-by-line rescan, which the match +# budget bounds. A benign OR arm supplies a second file the SQL predicate selects independently +# so a fully-scanned file survives the trip. + +_CATASTROPHIC = r"(?:(?:a{1,10}){1,10}){1,10}b" + + +@pytest.mark.integration +def test_match_budget_trip_is_recoverable_and_partial() -> None: + # File A (benign) sorts before file B (pathological) in (repo_id, path) order, so A is + # FULLY scanned before the budget trips inside B. A's matches survive; B is discarded. + schema = _unique(SCHEMA_PREFIX) + engine = create_db_engine() + conn = engine.connect() + try: + conn.execute(text(f"CREATE SCHEMA {schema}")) + conn.execute(text(f"SET search_path TO {schema}, public")) + conn.execute(text("CREATE EXTENSION IF NOT EXISTS pg_trgm")) + conn.commit() + Base.metadata.create_all(bind=conn) + conn.commit() + repo_id = _insert_repo(conn, "solo/repo") + _insert_file(conn, repo_id, "a_benign.txt", lang=None, content="alpha match here\n") + # Line 1: a long run of `a`s with no `b` (the catastrophic input). Line 2 `aaab` is what + # the SQL predicate matches to make this file a candidate at all. + _insert_file(conn, repo_id, "z_patho.txt", lang=None, content="a" * 55 + "\naaab\n") + conn.commit() + + query = f"alpha OR /{_CATASTROPHIC}/" + started = time.monotonic() + result = grep_search(conn, query, match_budget_ms=100) + elapsed = time.monotonic() - started + + # One-sided timing only: the budgeted call returns instead of hanging on the backtrack. + assert elapsed < 5 + assert result.truncated is True + assert result.truncation_reason == "match_budget" + # The fully-scanned benign file A survives; the pathological B is discarded. + assert _paths(result) == ["a_benign.txt"] + finally: + conn.rollback() + conn.execute(text(f"DROP SCHEMA IF EXISTS {schema} CASCADE")) + conn.commit() + conn.close() + engine.dispose() + + +@pytest.mark.integration +def test_match_budget_pagination_steps_past_pathological_file() -> None: + # Pagination mode: page 1 dies INSIDE the pathological file B (which sorts before benign C). + # Because the mid-file trip treats B as consumed, next_cursor points PAST B; a page 2 with a + # fresh budget then returns C -- pagination never stalls forever on the pathological file. + schema = _unique(SCHEMA_PREFIX) + engine = create_db_engine() + conn = engine.connect() + try: + conn.execute(text(f"CREATE SCHEMA {schema}")) + conn.execute(text(f"SET search_path TO {schema}, public")) + conn.execute(text("CREATE EXTENSION IF NOT EXISTS pg_trgm")) + conn.commit() + Base.metadata.create_all(bind=conn) + conn.commit() + repo_id = _insert_repo(conn, "solo/repo") + patho_content = "a" * 55 + "\naaab\n" + _insert_file(conn, repo_id, "b_patho.txt", lang=None, content=patho_content) + _insert_file(conn, repo_id, "c_benign.txt", lang=None, content="hit here\n") + conn.commit() + + query = f"hit OR /{_CATASTROPHIC}/" + page1 = grep_search(conn, query, match_budget_ms=100, cursor=None) + assert page1.truncated is True + assert page1.truncation_reason == "match_budget" + # Mid-file trip consumed B: the cursor resumes PAST it, not at it. + assert page1.next_cursor == FileCursor(repo_id, "b_patho.txt", content_sha(patho_content)) + + # Page 2 with a FRESH normal budget resumes after B and returns the benign C. + started = time.monotonic() + page2 = grep_search(conn, query, match_budget_ms=2000, cursor=page1.next_cursor) + assert time.monotonic() - started < 5 + assert _paths(page2) == ["c_benign.txt"] + assert page2.truncated is False + assert page2.truncation_reason is None + finally: + conn.rollback() + conn.execute(text(f"DROP SCHEMA IF EXISTS {schema} CASCADE")) + conn.commit() + conn.close() + engine.dispose() + + +@pytest.mark.integration +def test_match_budget_pre_file_trip_on_first_row_of_resumed_page_falls_back_to_resume_cursor() -> ( + None +): + # Regression test: a pre-file trip on the very FIRST content row of a resumed page must not + # strand the cursor. `last_candidate` is None (nothing consumed yet this call), so without + # the fallback, `next_cursor` would come out None while `truncated=True` -- indistinguishable + # from "exhausted" to a caller, silently dropping the rest of the corpus. A negative + # `match_budget_ms` deterministically puts the deadline in the past BEFORE the content query + # even runs, so this trips on row 0 with no dependency on DB timing (fully one-sided/non-flaky, + # unlike relying on a slow first-row fetch). + schema = _unique(SCHEMA_PREFIX) + engine = create_db_engine() + conn = engine.connect() + try: + conn.execute(text(f"CREATE SCHEMA {schema}")) + conn.execute(text(f"SET search_path TO {schema}, public")) + conn.execute(text("CREATE EXTENSION IF NOT EXISTS pg_trgm")) + conn.commit() + Base.metadata.create_all(bind=conn) + conn.commit() + repo_id = _insert_repo(conn, "solo/repo") + _insert_file(conn, repo_id, "m_benign.txt", lang=None, content="alpha match here\n") + conn.commit() + + # Sorts before the seeded file's (repo_id, path, content_sha) key, so the file remains + # a candidate on resume -- mimics a caller resuming after some earlier, unrelated file. + resume_cursor = FileCursor(repo_id, "", "") + result = grep_search(conn, "alpha", match_budget_ms=-1000, cursor=resume_cursor) + + assert result.truncated is True + assert result.truncation_reason == "match_budget" + assert result.files == () + # Falls back to the unchanged incoming cursor, not None -- the page is retryable with a + # fresh budget instead of looking exhausted. + assert result.next_cursor == resume_cursor + finally: + conn.rollback() + conn.execute(text(f"DROP SCHEMA IF EXISTS {schema} CASCADE")) + conn.commit() + conn.close() + engine.dispose() diff --git a/tests/unit/test_grep.py b/tests/unit/test_grep.py index 53103f3..63313b6 100644 --- a/tests/unit/test_grep.py +++ b/tests/unit/test_grep.py @@ -9,13 +9,16 @@ import ast import inspect import re +import time import pytest +import regex from app.query.parser import parse, resolve_case from app.search import grep as grep_module from app.search.grep import ( LineMatch, + MatchBudgetExceeded, _build_matchers, _no_content_atom, _zero_width_only_atoms, @@ -422,3 +425,54 @@ def test_flags_are_mutually_exclusive(query: str) -> None: # needs `bool(patterns)`. The envelope only ever clears flags, so this holds there too. no_content, zero_width = _flags(query) assert not (no_content and zero_width) + + +# ------------------------------------------------------------------ match budget (#38) +# +# extract_line_matches gains an optional `deadline` (a time.monotonic() absolute timestamp). +# None (the default) is the unbudgeted path existing callers rely on; a set deadline bounds +# Python-side CPU via the `regex` module's timeout=, raising MatchBudgetExceeded on a trip. +# +# The pathological exemplar below was reverified empirically against the installed `regex` +# version (2026.7.19): `(?:(?:a{1,10}){1,10}){1,10}b` over a run of `a`s with no `b` trips +# the timeout in well under a second at a small budget. A naive `(a+)+b` gets optimized away +# by the engine and does NOT trip -- do not substitute it. + +_CATASTROPHIC_PATTERN = r"(?:(?:a{1,10}){1,10}){1,10}b" + + +@pytest.mark.unit +def test_catastrophic_regex_finditer_is_interrupted_by_deadline() -> None: + pattern = regex.compile(_CATASTROPHIC_PATTERN) + content = "a" * 50 # no trailing "b": forces the engine into catastrophic backtracking + deadline = time.monotonic() + 0.05 + started = time.monotonic() + with pytest.raises(MatchBudgetExceeded): + extract_line_matches(content, [pattern], deadline=deadline) + # One-sided only: the trip must be bounded well under an unbudgeted run (which would take + # many seconds / effectively hang). No lower bound -- timing races must never make this flaky. + assert time.monotonic() - started < 5 + + +@pytest.mark.unit +def test_deadline_in_past_raises_before_scanning() -> None: + # A deadline already in the past raises immediately, before any matching happens. + pattern = regex.compile("foo") + deadline = time.monotonic() - 1 + with pytest.raises(MatchBudgetExceeded): + extract_line_matches("foo bar foo", [pattern], deadline=deadline) + + +@pytest.mark.unit +def test_deadline_none_preserves_unbudgeted_behavior() -> None: + # Passing deadline=None (and omitting it entirely) behaves exactly as the pre-existing + # unbudgeted path: same matches, same byte ranges as test_substring_single_match_*. + patterns, _ = _patterns("foo") + content = "first line\nsecond foo here" + explicit_none = extract_line_matches(content, patterns, deadline=None) + omitted = extract_line_matches(content, patterns) + assert explicit_none == omitted + (m,) = explicit_none + assert m.line_number == 2 + assert m.line_text == "second foo here" + assert m.byte_ranges == ((7, 10),) diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index dbe3b45..e396d6f 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -1278,6 +1278,30 @@ def test_signals_log_includes_both_flags() -> None: assert signals["zero_width_only_atoms"] is False +@pytest.mark.unit +def test_match_budget_ms_defaults_to_2000_and_is_env_overridable( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # Default matches DEFAULT_MATCH_BUDGET_MS in app/search/grep.py; the CODE_SEARCH_ prefix + # gives an automatic env override, exactly like statement_timeout_ms / max_content_bytes. + monkeypatch.delenv("CODE_SEARCH_MATCH_BUDGET_MS", raising=False) + assert Settings(lakebase_endpoint=None).match_budget_ms == 2000 + + monkeypatch.setenv("CODE_SEARCH_MATCH_BUDGET_MS", "500") + assert Settings(lakebase_endpoint=None).match_budget_ms == 500 + + +@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 + # log line, not just the bare `truncated` bool. + signals = main._signals({"truncated": True, "truncation_reason": "match_budget"}) + assert signals["truncated"] is True + assert signals["truncation_reason"] == "match_budget" + # None-safe on payloads that never carry it (e.g. list_repos). + assert main._signals({})["truncation_reason"] is None + + @pytest.mark.observability def test_signals_log_includes_regex_invalid() -> None: # Without this, a Postgres-rejected regex is log-indistinguishable from a genuine diff --git a/tests/unit/test_service.py b/tests/unit/test_service.py index f8a929f..5411bc8 100644 --- a/tests/unit/test_service.py +++ b/tests/unit/test_service.py @@ -214,6 +214,24 @@ def test_page_one_cursor_none_includes_next_cursor_key(monkeypatch: pytest.Monke assert payload["next_cursor"] is None # exhausted -- grep.GrepResult.next_cursor was None +@pytest.mark.unit +def test_match_budget_truncation_reason_passes_through_envelope( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # A grep match-budget trip surfaces as truncated=True + truncation_reason="match_budget" + # in the envelope unchanged -- the service layer never rewrites grep's truncation reason. + monkeypatch.setattr( + service, + "grep_search", + lambda *a, **k: _grep(truncated=True, truncation_reason="match_budget"), + ) + monkeypatch.setattr(service, "symbol_search", lambda *a, **k: _no_sym()) + + payload = service.search_code_payload(_FakeEngine([]), _cfg(), "foo", 50) + assert payload["truncated"] is True + assert payload["truncation_reason"] == "match_budget" + + @pytest.mark.unit def test_pagination_mode_encodes_grep_next_cursor(monkeypatch: pytest.MonkeyPatch) -> None: file_cursor = FileCursor(repo_id=7, path="src/handler.go", content_sha="deadbeef")