fix: bound knowledge-search matching so one pattern cannot stall the worker (#27471)

build_matcher compiled a caller-supplied pattern with Python's backtracking re and ran it over every line of every reachable file, with no timeout, no thread offload and no length caps. is_regex_pattern promotes any pattern containing a metacharacter, and a bare pipe counts, so no explicit regex flag is needed to reach the compiler. The search loop is synchronous inside an async handler, and UVICORN_WORKERS defaults to 1, so the cost lands on every other user of the instance. MAX_GREP_RESULTS bounds how many matches are reported, not how much work is done.

Backtracking cost is exponential in the length of the text being matched, so capping the pattern or the line does not bound it: the subject in the measurements below is 30 characters. `(x|x)*y` against a line of 30 x took 80 seconds, `(a+)+$` against 32 a took 169 seconds, and the same subject with a literal pattern took 0.6 microseconds.

Matching now runs on the regex module, which accepts a per-search timeout that re has no equivalent for. The timeout is the actual bound: regex resolves many classic catastrophic patterns instantly, but not all of them, and `(a|aa)+$` and `(?:a|a)*$` still need it. The budget covers a whole tool call rather than a single search, because a pipeline builds one matcher per segment and a per-search budget would multiply by segment count, and because a per-line timeout would allow timeout multiplied by line count. It is carried in a context variable so one command shares it without threading a parameter through every handler, and it is charged only for time spent inside search(), so database round-trips and other coroutines cannot consume it. Exhausting it raises, and both entry points already render that as an error for the model to read.

Note for anyone tracking search behaviour: re and the regex module define \w, \W and \b differently on non-ASCII text. re follows str.isalnum(), the regex module follows UTS#18, so \w no longer matches superscripts and fractions such as the ones in Nd-adjacent categories, and now does match combining marks. POSIX classes like [[:alpha:]] are interpreted rather than read as a literal set, and \p{...} compiles instead of erroring. Results on ASCII content are unchanged.

regex was already installed as a transitive dependency of nltk, tiktoken and transformers. It is now declared directly, pinned in pyproject.toml and requirements.txt to the version the lockfile already resolves.
This commit is contained in:
Classic298
2026-07-27 09:22:26 +02:00
committed by GitHub
parent 147c3b6ac8
commit 3ab2026262
3 changed files with 57 additions and 5 deletions
+55 -5
View File
@@ -7,13 +7,16 @@ for AI models to interact with knowledge bases using commands they already know.
Re-exported through builtin.py for consistent imports.
"""
import contextvars
import json
import logging
import re
import shlex
import time
from contextlib import contextmanager
from typing import Optional
import regex
from fastapi import Request
log = logging.getLogger(__name__)
@@ -26,6 +29,36 @@ DEFAULT_HEAD_LINES = 10
DEFAULT_TAIL_LINES = 10
MAX_GREP_MATCHES = 50
# Matching time allowed per tool call. Backtracking cost is exponential in the length of the
# matched text, so capping the pattern or the line does not bound it.
MATCH_BUDGET_SECONDS = 2.0
class MatchBudgetExceeded(Exception):
"""A tool call spent its whole matching budget, so the caller reports it."""
class MatchBudget:
"""Matching time remaining, counted only inside search() so awaits do not consume it."""
def __init__(self):
self.remaining = MATCH_BUDGET_SECONDS
# Scoped to the running task, so one budget covers every matcher a command builds without
# threading it through each handler.
_active_budget: contextvars.ContextVar[MatchBudget | None] = contextvars.ContextVar('kb_match_budget', default=None)
@contextmanager
def match_budget():
"""Bound the matching time of one tool call rather than of each search it runs."""
token = _active_budget.set(MatchBudget())
try:
yield
finally:
_active_budget.reset(token)
# =============================================================================
# SHARED REGEX UTILITIES — also used by builtin.py grep_knowledge_files
@@ -59,11 +92,26 @@ def build_matcher(pattern: str, case_insensitive: bool = False, use_regex: bool
if use_regex:
normalized = normalize_regex(pattern)
try:
re_flags = re.IGNORECASE if case_insensitive else 0
compiled = re.compile(normalized, re_flags)
except re.error as e:
re_flags = regex.IGNORECASE if case_insensitive else 0
compiled = regex.compile(normalized, re_flags)
except regex.error as e:
return None, f'Invalid regex: {e}'
return (lambda line: bool(compiled.search(line))), None
budget = _active_budget.get() or MatchBudget()
def matches(line: str) -> bool:
started = time.monotonic()
try:
# A negative timeout disables it, so an exhausted budget must not reach search().
if budget.remaining <= 0:
raise TimeoutError
return bool(compiled.search(line, timeout=budget.remaining))
except TimeoutError:
raise MatchBudgetExceeded(f'Search exceeded {MATCH_BUDGET_SECONDS:g}s, narrow the pattern') from None
finally:
budget.remaining -= time.monotonic() - started
return matches, None
else:
sp = pattern.lower() if case_insensitive else pattern
return (lambda line: sp in (line.lower() if case_insensitive else line)), None
@@ -1131,7 +1179,9 @@ async def kb_exec(
if not segments:
return 'Could not parse command. Run kb_exec("ls") to start.'
return await _execute_pipeline(segments, __user__, __model_knowledge__)
# One budget for the whole command: a per-search budget would multiply by segment count.
with match_budget():
return await _execute_pipeline(segments, __user__, __model_knowledge__)
except Exception as e:
log.exception(f'kb_exec error: {e}')
return f'Error: {e}'
+1
View File
@@ -13,6 +13,7 @@ authlib==1.7.2
joserfc==1.7.4
requests==2.34.2
regex==2026.5.9 # supports a per-search timeout, which `re` does not
aiohttp==3.13.5 # do not update to 3.13.3 - broken
aiodns==4.0.4 # makes aiohttp resolve DNS on the event loop instead of the threadpool
async-timeout==5.0.1
+1
View File
@@ -51,6 +51,7 @@ dependencies = [
"asgiref==3.11.1",
"tiktoken==0.13.0",
"regex==2026.5.9",
"mcp==1.27.2",
"openai==2.29.0",