search: bound Python regex CPU per request (#38)

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
 

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

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:<value>`` 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

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

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 <metric> 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.