From 4e9e8ca207eefd4ba20525ea4c9e31f08b4247cb Mon Sep 17 00:00:00 2001 From: Patrick Buckley Date: Mon, 4 May 2026 02:24:44 +0000 Subject: [PATCH] Bound search tool output against pathological inputs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replaces the per-line truncation with a fully bounded pipeline so the search tool can no longer overflow the LLM context — or OOM the parent — on minified bundles, multi-GB JSONL records, or huge result sets. Backend: - Prefer ripgrep when on PATH; grep is the fallback. Detection is cached via functools.cache. - ripgrep flags do most of the bounding natively: --max-columns 1024 + --max-columns-preview, --max-filesize 10M, --max-count 100, --no-config, --no-messages, plus negative globs for the same noisy directories grep has been excluding. - ripgrep added to the Dockerfile. Streaming subprocess (_search_capture): - subprocess.Popen with a streaming, byte-capped stdout read (4 MB). Defends against single-line files (training data, minified bundles) that would have OOM'd the previous subprocess.run capture. - threading.Timer watchdog enforces tool_timeout even when the pipe read is blocked in the kernel — proc.wait(timeout=…) alone was insufficient because the read sat ahead of it. - Stderr drained in a daemon thread to avoid pipe-deadlock when the child writes to stderr while we're still reading stdout. Cap on captured stderr keeps a hostile child from growing the buffer. Tier-based formatter (_format_search_results): - Tier 1: full path:line:content output, stream-emitted with a running-cost short-circuit so we never materialize past the budget. - Tier 2: K samples per file with overflow notes; K is computed analytically from budget / file_count / avg-line-length so we hit the right ladder rung in a single pass. - Tier 3: per-file counts only, also budget-bounded with a tail line reporting the omitted files. Sorted by descending count. - Total output budget (32 KB) is well under tool_truncation, so the head+tail _truncate_output strategy never silently drops middle files in a search result. Argument injection fix: - The ripgrep arg list was missing the `--` separator that the grep branch already had. With auto_approve on the search tool, that was exploitable: path='--pre=COMMAND' would have made ripgrep run the script as a per-file preprocessor and surface its stdout. Added `--` and a regression test. State-machine cleanup in _exec_search: - rc < 0 (signal-killed by something other than us) now surfaces a dedicated 'killed by signal N' message instead of being parsed as success. - capped + zero parsed records (e.g. one multi-MB line with no \n) now returns a dedicated byte-cap message instead of the malformed- output message that previously masked the real cause. - _report_tool_result descriptions now match the returned payload (no more 'no matches' tag on a 'malformed' payload). Defence-in-depth on env scrub: - RIPGREP_CONFIG_PATH, GIT_CONFIG, GIT_CONFIG_GLOBAL, GIT_CONFIG_SYSTEM added to _EXPLICIT_SCRUB. We pass --no-config on the rg CLI today, but if a future caller forgets the flag, an attacker who can set one of these env vars could plant a config containing --pre=… and recreate the same RCE shape. Tests: - TestSearchLineTruncation rewritten to mock _search_capture instead of subprocess.run (the previous tests passed ChatSession kwargs that no longer satisfy the constructor). - TestSearchBackendSelection covers rg/grep detection and arg construction, including the --pre flag-injection regression. - TestSearchOutputBudget exercises Tier 1/2/3 directly. - TestSearchCaptureStreaming spawns real Python subprocess writers to exercise the byte-cap trim, mega-line-no-newline edge case, the watchdog timeout when the child writes nothing, and the stderr drain under load. - test_env_scrub picks up the new tool-config keys. --- Dockerfile | 7 +- tests/test_env_scrub.py | 10 + tests/test_session.py | 396 +++++++++++++++++++++++++++++++++ turnstone/core/env.py | 9 + turnstone/core/session.py | 448 +++++++++++++++++++++++++++++++++----- 5 files changed, 814 insertions(+), 56 deletions(-) diff --git a/Dockerfile b/Dockerfile index cb7ed75a..443cd227 100644 --- a/Dockerfile +++ b/Dockerfile @@ -13,9 +13,12 @@ COPY --from=ghcr.io/astral-sh/uv:0.11.8 /uv /usr/local/bin/uv # Remove the slim image's man page exclusion so man-db has actual content RUN rm -f /etc/dpkg/dpkg.cfg.d/docker -# System dependencies: psycopg (libpq5), developer tooling for agent workflows +# System dependencies: psycopg (libpq5), developer tooling for agent workflows. +# ripgrep is the preferred backend for the search tool — natively bounds +# per-line, per-file, and per-filesize so pathological inputs (minified +# bundles, training-data JSONL with multi-MB single records) can't OOM us. RUN apt-get update && apt-get upgrade -y && apt-get install -y --no-install-recommends \ - libpq5 git curl jq man-db manpages procps file \ + libpq5 git curl jq man-db manpages procps file ripgrep \ && rm -rf /var/lib/apt/lists/* # Node.js LTS (for npx-based MCP servers like @modelcontextprotocol/server-github) diff --git a/tests/test_env_scrub.py b/tests/test_env_scrub.py index 6500376d..047de4af 100644 --- a/tests/test_env_scrub.py +++ b/tests/test_env_scrub.py @@ -15,6 +15,16 @@ class TestIsSecret: assert _is_secret("TURNSTONE_JWT_SECRET") is True assert _is_secret("AWS_SECRET_ACCESS_KEY") is True + def test_tool_config_paths_scrubbed(self): + """Tool-config env vars whose target files load executable + directives must be scrubbed even though they don't match a + secret-suffix pattern. Defence-in-depth alongside on-CLI + ``--no-config`` for ripgrep and friends.""" + assert _is_secret("RIPGREP_CONFIG_PATH") is True + assert _is_secret("GIT_CONFIG") is True + assert _is_secret("GIT_CONFIG_GLOBAL") is True + assert _is_secret("GIT_CONFIG_SYSTEM") is True + def test_suffix_matching(self): assert _is_secret("MY_CUSTOM_API_KEY") is True assert _is_secret("DB_PASSWORD") is True diff --git a/tests/test_session.py b/tests/test_session.py index 996c379c..1c8d38bb 100644 --- a/tests/test_session.py +++ b/tests/test_session.py @@ -3,8 +3,11 @@ import base64 import contextlib import json +import subprocess from unittest.mock import MagicMock, patch +import pytest + from turnstone.core.session import _IMAGE_EXTENSIONS, _IMAGE_SIZE_CAP, ChatSession @@ -83,6 +86,25 @@ def _make_session( return ChatSession(**defaults) +def _run_exec_search(session, capture_return): + """Patch ``_search_capture`` to ``capture_return`` and run ``_exec_search``. + + Returns the formatted output string. The fixed call args + (``call_id``/``pattern``/``path``) are deliberately uniform across the + line-truncation tests — only the captured stdout/rc/stderr/capped tuple + varies between cases. + """ + with patch.object(session, "_search_capture", return_value=capture_return): + _, output = session._exec_search( + { + "call_id": "test_call", + "pattern": "test_pattern", + "path": "/workspace/turnstone", + } + ) + return output + + class TestChatSessionConstruction: def test_system_messages_created(self, tmp_db): session = _make_session() @@ -2891,3 +2913,377 @@ class TestSessionUIBaseToolReminderHook: "tool_call_id": "call_abc123", } ] + + +class TestSearchLineTruncation: + """Tests for search tool line truncation to prevent context overflow.""" + + def test_search_truncates_long_lines_preserves_path(self): + """Long lines are truncated but path:line: prefix is preserved for file counting.""" + from turnstone.core.session import ( + _MAX_SEARCH_LINE_LENGTH, + _SEARCH_LINE_MARGIN, + _SEARCH_TRUNCATION_SUFFIX, + ) + + # path:line:content where content is way over the cap+margin + long_content = "x" * 5000 + stdout = f"turnstone/core/session.py:100:{long_content}\n".encode() + + output = _run_exec_search(_make_session(), (stdout, 0, b"", False)) + + assert _SEARCH_TRUNCATION_SUFFIX in output + assert "turnstone/core/session.py" in output + # The *content portion* (after the 2nd colon) is what's bounded by + # the per-line cap; the path prefix is unbounded. + max_content_len = ( + _MAX_SEARCH_LINE_LENGTH + len(_SEARCH_TRUNCATION_SUFFIX) + _SEARCH_LINE_MARGIN + ) + for line in output.splitlines(): + if "matches across" in line or not line.strip(): + continue + parts = line.split(":", 2) + if len(parts) == 3: + assert len(parts[2]) <= max_content_len + + def test_search_file_counting_with_truncated_lines(self): + """File counting works correctly even with truncated lines.""" + stdout = ( + "turnstone/core/session.py:100:" + "x" * 5000 + "\n" + "turnstone/core/auth.py:50:normal line\n" + "turnstone/core/session.py:200:" + "y" * 3000 + "\n" + ).encode() + + output = _run_exec_search(_make_session(), (stdout, 0, b"", False)) + + assert "3 matches across 2 files" in output + assert "turnstone/core/session.py" in output + assert "turnstone/core/auth.py" in output + + def test_search_drops_lines_without_colon(self): + """Lines without any colon are dropped at the parsing step.""" + from turnstone.core.session import _SEARCH_ALL_TRUNCATED_MSG + + # No colon anywhere — parsed records list is empty. + stdout = ("turnstone/core/session.py" + "x" * 5000 + "\n").encode() + + output = _run_exec_search(_make_session(), (stdout, 0, b"", False)) + + assert output == _SEARCH_ALL_TRUNCATED_MSG + + def test_search_handles_single_colon_lines(self): + """Lines with one colon and a non-numeric line-number portion are dropped.""" + from turnstone.core.session import _SEARCH_ALL_TRUNCATED_MSG + + # path:100xxxxx... — partition's lineno chunk has trailing junk, .isdigit() fails + stdout = ("turnstone/core/session.py:100" + "x" * 5000 + "\n").encode() + + output = _run_exec_search(_make_session(), (stdout, 0, b"", False)) + + assert output == _SEARCH_ALL_TRUNCATED_MSG + + def test_search_no_truncation_for_short_lines(self): + """Short lines pass through unchanged.""" + stdout = b"turnstone/core/session.py:100:short line\n" + + output = _run_exec_search(_make_session(), (stdout, 0, b"", False)) + + assert "...[truncated" not in output + assert "short line" in output + + def test_search_no_matches(self): + """rc==1 (no matches) returns the friendly no-matches sentinel.""" + output = _run_exec_search(_make_session(), (b"", 1, b"", False)) + assert output == "(no matches)" + + def test_search_error_propagates_stderr(self): + """rc>1 surfaces stderr text, not a generic message, when stderr is non-empty.""" + output = _run_exec_search( + _make_session(), + (b"", 2, b"grep: foo: No such file or directory\n", False), + ) + assert "No such file or directory" in output + + def test_search_capped_flag_in_output(self): + """When raw stdout is byte-capped, results note the partial output.""" + stdout = b"a/b.py:1:line1\na/b.py:2:line2\n" + output = _run_exec_search(_make_session(), (stdout, 0, b"", True)) + assert "byte cap" in output or "capped" in output + + +class TestSearchBackendSelection: + """Tests for backend detection (rg vs grep) and arg construction.""" + + def test_detect_uses_rg_when_on_path(self): + from turnstone.core.session import _detect_search_backend + + # Reset cache so the patch takes effect. + _detect_search_backend.cache_clear() + try: + with patch("turnstone.core.session.shutil.which", return_value="/usr/bin/rg"): + assert _detect_search_backend() == "rg" + finally: + _detect_search_backend.cache_clear() + + def test_detect_falls_back_to_grep(self): + from turnstone.core.session import _detect_search_backend + + _detect_search_backend.cache_clear() + try: + with patch("turnstone.core.session.shutil.which", return_value=None): + assert _detect_search_backend() == "grep" + finally: + _detect_search_backend.cache_clear() + + def test_detect_caches_result(self): + from turnstone.core.session import _detect_search_backend + + _detect_search_backend.cache_clear() + try: + with patch( + "turnstone.core.session.shutil.which", return_value="/usr/bin/rg" + ) as mock_which: + _detect_search_backend() + _detect_search_backend() + _detect_search_backend() + assert mock_which.call_count == 1 + finally: + _detect_search_backend.cache_clear() + + def test_rg_args_include_size_and_column_caps(self): + from turnstone.core.session import ( + _MAX_SEARCH_LINE_LENGTH, + _SEARCH_MAX_FILESIZE, + _build_search_args, + ) + + args = _build_search_args("foo", "/some/path", "rg") + assert args[0] == "rg" + # Per-line cap with preview marker (the load-bearing flag pair) + assert "--max-columns" in args + assert str(_MAX_SEARCH_LINE_LENGTH) in args + assert "--max-columns-preview" in args + # Per-file size guard against multi-MB JSONL records + assert "--max-filesize" in args + assert _SEARCH_MAX_FILESIZE in args + # Per-file match cap + assert "--max-count" in args + # ``-e `` form so patterns starting with ``-`` are safe; + # ``--`` separator before the path so paths starting with ``-`` + # (e.g. ``--pre=/tmp/x``) cannot be parsed as ripgrep flags. + assert "-e" in args + e_idx = args.index("-e") + assert args[e_idx + 1] == "foo" + assert "--" in args + sep = args.index("--") + assert args[sep + 1] == "/some/path" + assert args[-1] == "/some/path" + + def test_rg_args_protect_path_from_flag_injection(self): + """A ``path`` starting with ``-`` cannot inject ripgrep flags. + + Regression test for an RCE vector: without the ``--`` separator, + ``path="--pre=/tmp/x.sh"`` would have made ripgrep execute the + script as a per-file preprocessor and surface its stdout as + search results. + """ + from turnstone.core.session import _build_search_args + + args = _build_search_args("foo", "--pre=/tmp/evil.sh", "rg") + assert "--" in args + sep = args.index("--") + assert args[sep + 1] == "--pre=/tmp/evil.sh" + # And the malicious path is the last token, not interspersed with flags. + assert args[-1] == "--pre=/tmp/evil.sh" + + def test_grep_args_include_excludes_and_separator(self): + from turnstone.core.session import _build_search_args + + args = _build_search_args("foo", "/some/path", "grep") + assert args[0] == "grep" + assert "-rn" in args + assert "-I" in args + assert "-E" in args + # Excludes for noisy build dirs + assert any(a == "--exclude-dir=node_modules" for a in args) + assert any(a == "--exclude-dir=.git" for a in args) + # ``--`` separator is what protects pattern-as-flag in grep + assert "--" in args + sep = args.index("--") + assert args[sep + 1] == "foo" + assert args[sep + 2] == "/some/path" + + +class TestSearchOutputBudget: + """Tests for tier-based degradation when output exceeds the budget.""" + + def test_tier1_fits_full_output(self): + from turnstone.core.session import _format_search_results + + records = [ + ("foo.py", "1", "small match"), + ("bar.py", "2", "another match"), + ("foo.py", "3", "third match"), + ] + out = _format_search_results(records, capped=False) + assert "foo.py:1:small match" in out + assert "bar.py:2:another match" in out + assert "foo.py:3:third match" in out + assert "3 matches across 2 files" in out + + def test_tier2_samples_when_over_budget(self): + """Many matches per file → degrade to K samples per file with overflow notes.""" + from turnstone.core.session import _SEARCH_OUTPUT_BUDGET, _format_search_results + + # 3 files × 200 matches/file × ~80 chars/line ≈ 48 KB → over the 32 KB budget + records = [] + line = "x" * 60 + for f in ("a.py", "b.py", "c.py"): + for i in range(200): + records.append((f, str(i), line)) + out = _format_search_results(records, capped=False) + # Should have collapsed to per-file samples + overflow note + assert "showing first" in out + assert "more in a.py" in out + assert "more in b.py" in out + assert "more in c.py" in out + assert len(out) <= _SEARCH_OUTPUT_BUDGET + 512 # small slack for header + + def test_tier3_counts_only_when_too_many_files(self): + """Thousands of files × matches → degrade to per-file counts.""" + from turnstone.core.session import _SEARCH_OUTPUT_BUDGET, _format_search_results + + records = [] + # 2000 files × 50 matches × 80 chars = 8 MB; well past budget even at 1/file + line = "x" * 60 + for f_idx in range(2000): + for i in range(50): + records.append((f"path/to/file_{f_idx:04}.py", str(i), line)) + out = _format_search_results(records, capped=False) + assert "Counts only" in out + assert "path/to/file_0000.py: 50 matches" in out + assert len(out) <= _SEARCH_OUTPUT_BUDGET + 512 + + def test_tier1_preserves_file_order(self): + """Tier 1 emits files in insertion order (so first-seen file appears first).""" + from turnstone.core.session import _format_search_results + + records = [ + ("z.py", "1", "first"), + ("a.py", "2", "second"), + ("z.py", "3", "third"), + ] + out = _format_search_results(records, capped=False) + z_idx = out.index("z.py:1:") + a_idx = out.index("a.py:2:") + assert z_idx < a_idx, "first-seen file (z.py) should appear before later-seen (a.py)" + + def test_capped_flag_propagates_to_summary(self): + from turnstone.core.session import _format_search_results + + records = [("foo.py", "1", "match")] + out = _format_search_results(records, capped=True) + assert "byte cap" in out or "capped" in out + + +class TestSearchCaptureStreaming: + """Direct tests for ``_search_capture`` — the streaming subprocess + layer that backs ``_exec_search``. These tests do NOT mock subprocess; + they spawn small ``python -c`` writers so the byte-cap, last-newline + trim, and timeout paths actually execute in real OS processes. + """ + + def test_byte_cap_trims_to_last_newline(self): + """Writer emits >cap bytes of well-formed lines; capture caps and + trims to the last newline so the parser never sees a partial + trailing line.""" + import sys + + from turnstone.core.session import _SEARCH_RAW_BYTE_CAP + + session = _make_session() + # Each line is "p:1:" + 1023 'x' chars + '\n' = 1028 bytes; emit + # enough lines to comfortably exceed the 4 MB cap. + line_count = (_SEARCH_RAW_BYTE_CAP // 1028) + 100 + writer = ( + "import sys\n" + f"line = 'p:1:' + ('x' * 1023) + '\\n'\n" + f"sys.stdout.buffer.write(line.encode() * {line_count})\n" + ) + stdout, rc, stderr, capped = session._search_capture([sys.executable, "-c", writer]) + assert capped is True + assert len(stdout) <= _SEARCH_RAW_BYTE_CAP + # Trim was applied — every parsed line is well-formed (no partial + # trailing line). The buffer is sliced at the last newline, which + # discards the (possibly partial) bytes after it. + lines = stdout.splitlines() + assert lines, "expected at least one complete line" + for raw in lines: + assert raw.startswith(b"p:1:") + assert len(raw) == 1027 # "p:1:" + 1023 x's, no trailing \n + + def test_byte_cap_mega_line_no_newline(self): + """A single multi-MB line with no newline is the worst-case input + (think a JSONL training record on one line). The cap fires and + ``last_nl == -1`` skips the trim — _exec_search distinguishes + this from 'all malformed' via the dedicated byte-cap message.""" + import sys + + from turnstone.core.session import _SEARCH_RAW_BYTE_CAP + + session = _make_session() + # 5 MB of bytes, no newlines anywhere. + writer = "import sys\nsys.stdout.buffer.write(b'a' * (5 * 1024 * 1024))\n" + stdout, rc, stderr, capped = session._search_capture([sys.executable, "-c", writer]) + assert capped is True + assert len(stdout) == _SEARCH_RAW_BYTE_CAP + assert b"\n" not in stdout + + def test_timeout_raises_even_when_child_writes_nothing(self): + """Watchdog enforces tool_timeout regardless of whether the + child has written anything to stdout — ``proc.stdout.read`` is a + blocking pipe read that wouldn't otherwise honour the timeout. + Regression test for bug-1. + """ + import sys + + session = _make_session(tool_timeout=1) + # Sleep silently — never writes to stdout — so the read blocks. + sleeper = "import time; time.sleep(30)\n" + with pytest.raises(subprocess.TimeoutExpired): + session._search_capture([sys.executable, "-c", sleeper]) + + def test_clean_exit_returns_full_output_uncapped(self): + """A child that writes a small amount and exits cleanly returns + ``capped=False`` and the full output verbatim.""" + import sys + + session = _make_session() + writer = "import sys; sys.stdout.write('a.py:1:hello\\n')\n" + stdout, rc, stderr, capped = session._search_capture([sys.executable, "-c", writer]) + assert capped is False + assert rc == 0 + assert stdout == b"a.py:1:hello\n" + + def test_stderr_drained_without_deadlock(self): + """If a child writes stderr in parallel with stdout, the drain + thread must keep the pipe flowing so the child doesn't block on + a full stderr buffer while we're reading stdout.""" + import sys + + session = _make_session() + # Write more to stderr than the OS pipe buffer (~64KB) while + # also writing stdout. Without the drain thread, the child + # blocks on stderr.write and we deadlock waiting for stdout EOF. + writer = ( + "import sys\n" + "sys.stderr.buffer.write(b'e' * (200 * 1024))\n" + "sys.stdout.buffer.write(b'a.py:1:done\\n')\n" + ) + stdout, rc, stderr, capped = session._search_capture([sys.executable, "-c", writer]) + assert rc == 0 + assert stdout == b"a.py:1:done\n" + # stderr was drained; the captured prefix is bounded by the cap. + from turnstone.core.session import _SEARCH_STDERR_CAP + + assert len(stderr) <= _SEARCH_STDERR_CAP diff --git a/turnstone/core/env.py b/turnstone/core/env.py index b87e5301..f112576c 100644 --- a/turnstone/core/env.py +++ b/turnstone/core/env.py @@ -75,6 +75,15 @@ _EXPLICIT_SCRUB: frozenset[str] = frozenset( "GOOGLE_APPLICATION_CREDENTIALS", "DATABASE_URL", # conventional name (Heroku, Railway, etc.) — kept for defence-in-depth "TURNSTONE_DB_URL", + # Tool-config env vars whose target files can directly load + # executable directives (preprocessor commands, pagers, etc.). + # Defence-in-depth alongside the on-CLI ``--no-config`` we pass + # to ripgrep — if a future caller forgets that flag, an attacker + # who can set one of these can plant a config that runs commands. + "RIPGREP_CONFIG_PATH", + "GIT_CONFIG", + "GIT_CONFIG_GLOBAL", + "GIT_CONFIG_SYSTEM", } ) diff --git a/turnstone/core/session.py b/turnstone/core/session.py index e44a4060..3ef9dc41 100644 --- a/turnstone/core/session.py +++ b/turnstone/core/session.py @@ -15,6 +15,7 @@ import contextlib import copy import dataclasses import difflib +import functools import hashlib import json import mimetypes @@ -195,6 +196,261 @@ def _encode_image_data_uri(raw: bytes, mime: str) -> str: # Upper bound on total skill content injected into system messages _MAX_SKILL_CONTENT: int = 32768 +# Maximum length (characters) for a single line in search results. +# Lines longer than this are truncated to prevent context overflow from +# pathological files (minified blobs, base64 data, etc.). +_MAX_SEARCH_LINE_LENGTH: int = 1024 +# Margin over the per-line cap before re-truncating, so backend-supplied +# preview markers (e.g. ripgrep's ``[... omitted end of long line]``) pass +# through cleanly without redundant " ...[truncated]" stacking. +_SEARCH_LINE_MARGIN: int = 128 +_SEARCH_TRUNCATION_SUFFIX: str = f"...[truncated, line length > {_MAX_SEARCH_LINE_LENGTH}]" +_SEARCH_ALL_TRUNCATED_MSG: str = ( + "(all matches returned were malformed -- re-check your search query, " + "if the issue persists there may be a problem with the search backend or the filesystem)" +) +# Total search-output budget (chars). Chosen well under ``tool_truncation`` +# (typically 256 KB+) so the head+tail ``_truncate_output`` strategy never +# kicks in for search results — that strategy silently drops middle files +# alphabetically, which is exactly the wrong shape for a grep result. +_SEARCH_OUTPUT_BUDGET: int = 32_768 +# Hard cap on raw bytes read from the search subprocess. Defends against +# pathological single-line files (multi-GB JSONL training records, etc.) +# that would otherwise OOM the parent process via ``subprocess.run``. +_SEARCH_RAW_BYTE_CAP: int = 4 * 1024 * 1024 +# Files larger than this are skipped entirely (ripgrep only — grep has no +# native equivalent and falls back to the byte cap above). +_SEARCH_MAX_FILESIZE: str = "10M" +# Per-file sample-count ladder for Tier 2 degradation. Each step is tried in +# order; the first K whose total emission fits the budget wins. The full +# 5/3/1 curve documents the degradation: prefer 5 samples per file, fall to +# 3, then a single representative sample before giving up to Tier 3. +_SEARCH_TIER2_SAMPLE_LADDER: tuple[int, ...] = (5, 3, 1) +# Bytes reserved at the end of the Tier 3 body for the +# "(plus N more files with M matches between them)" tail line, so we don't +# blow the budget when the count list itself is enormous. +_SEARCH_TIER3_TAIL_RESERVE: int = 80 +# Stderr-drain knobs for ``_search_capture``: bound the captured stderr so +# a hostile child can't grow the buffer indefinitely, and drain in +# moderate-sized chunks so the OS pipe buffer doesn't deadlock the child. +_SEARCH_STDERR_CAP: int = 64 * 1024 +_SEARCH_DRAIN_CHUNK: int = 8192 +# How long we wait for the stderr drain thread to finish after the child +# exits. The thread reads from a closed pipe at that point; a small +# timeout keeps shutdown bounded if the OS hasn't propagated EOF yet. +_SEARCH_DRAIN_JOIN_TIMEOUT: float = 2.0 +# Excluded directory patterns — hit by both backends. ripgrep also respects +# ``.gitignore`` and skips hidden directories by default, so most of these +# are belt-and-suspenders for the rg path; they're load-bearing for grep. +_SEARCH_EXCLUDE_DIRS: tuple[str, ...] = ( + ".git", + "node_modules", + "target", + "__pycache__", + ".mypy_cache", + ".ruff_cache", + ".pytest_cache", + "dist", + "build", + "*.egg-info", + ".tox", + ".venv", + "venv", + "vendor", +) + + +@functools.cache +def _detect_search_backend() -> str: + """Return ``'rg'`` if ripgrep is on PATH, else ``'grep'``. Cached.""" + return "rg" if shutil.which("rg") else "grep" + + +def _build_search_args(pattern: str, path: str, backend: str) -> list[str]: + """Build subprocess args for the chosen search backend. + + The ripgrep flag set is the load-bearing one: ``--max-columns`` + + ``--max-columns-preview`` bound per-line bytes natively (no Python-side + re-search needed), ``--max-filesize`` skips multi-MB JSONL/training + files entirely, and ``--max-count`` matches grep's ``-m`` per-file cap. + """ + if backend == "rg": + args = [ + "rg", + "-n", # line numbers + "-H", # always show filename + "--no-heading", # path:line:content format like grep + "--color=never", + "--no-config", # ignore ~/.ripgreprc for reproducibility + "--no-messages", # suppress filesystem error noise + "--max-count", + "100", + "--max-columns", + str(_MAX_SEARCH_LINE_LENGTH), + "--max-columns-preview", # show first N cols + omitted-marker + "--max-filesize", + _SEARCH_MAX_FILESIZE, + ] + for d in _SEARCH_EXCLUDE_DIRS: + args.extend(["-g", f"!{d}"]) + # ``-e`` protects the pattern from being parsed as a flag; ``--`` + # protects the path the same way. Without ``--`` an attacker who + # can prompt-inject the agent could pass ``path="--pre=COMMAND"`` + # and ripgrep would execute COMMAND as a per-file preprocessor. + args.extend(["-e", pattern, "--", path]) + return args + # grep fallback + args = ["grep", "-rn", "-I", "-E", "-m", "100", "--color=never"] + for d in _SEARCH_EXCLUDE_DIRS: + args.append(f"--exclude-dir={d}") + args.extend(["--", pattern, path]) + return args + + +def _parse_search_records(stdout: bytes) -> list[tuple[str, str, str]]: + """Parse ``path:lineno:content`` records from search backend stdout. + + Drops malformed lines (need ≥2 colons, numeric line-number, non-empty + path). Decodes bytes leniently for display. Lines that exceed the + per-line cap *plus* a small margin for backend-supplied truncation + markers are re-truncated with ``_SEARCH_TRUNCATION_SUFFIX``; this is + the load-bearing defense for the grep fallback (rg already enforces + ``--max-columns`` upstream). + """ + cap = _MAX_SEARCH_LINE_LENGTH + margin = _SEARCH_LINE_MARGIN + # Decode the whole buffer once rather than per-line — a cap-hit + # invocation can yield ~50K lines, and the per-line ``decode()`` was + # showing up in profiles. + text = stdout.decode("utf-8", errors="replace") + records: list[tuple[str, str, str]] = [] + for line in text.splitlines(): + path, sep1, rest = line.partition(":") + if not sep1 or not path: + continue + lineno, sep2, content = rest.partition(":") + if not sep2 or not lineno.isdigit(): + continue + if len(content) > cap + margin: + content = content[:cap] + _SEARCH_TRUNCATION_SUFFIX + records.append((path, lineno, content)) + return records + + +def _format_search_results( + records: list[tuple[str, str, str]], + capped: bool, +) -> str: + """Format match records with tiered degradation when output > budget. + + Tier 1: full ``path:line:content`` lines, stream-emitted with a running + cost check that short-circuits as soon as the budget would be exceeded. + Tier 2: K samples per file plus an ``…and N more in `` note, + stepping K down (5 → 3 → 1) until results fit. This guarantees every + file is at least mentioned, which prevents the alphabetic-bias dropout + that head+tail truncation produced. + Tier 3 (fallback): per-file counts only, sorted by descending count. + """ + by_file: dict[str, list[tuple[str, str]]] = {} + for path, lineno, content in records: + by_file.setdefault(path, []).append((lineno, content)) + total = len(records) + files = len(by_file) + if not total: + # Caller distinguishes "no matches" from "all malformed" via rc. + return _SEARCH_ALL_TRUNCATED_MSG + + summary = f"\n\n({total} matches across {files} files)" + if capped: + summary += " (raw output exceeded byte cap; results may be incomplete)" + + chunks: list[str] = [] + used = 0 + overflow = False + for path, matches in by_file.items(): + for lineno, content in matches: + line = f"{path}:{lineno}:{content}" + cost = len(line) + 1 + if used + cost + len(summary) > _SEARCH_OUTPUT_BUDGET: + overflow = True + break + chunks.append(line) + used += cost + if overflow: + break + if not overflow: + return "\n".join(chunks) + summary + + # Pick K analytically rather than iterating the K=5/3/1 ladder and + # rebuilding ``chunks2`` from scratch on each retry. Sample the first + # ~32 records for an emitted-line-length estimate, then divide the + # budget by ``files * (avg + 1)`` to get a K that should fit in one + # pass. Floor at 80 so a corpus of unusually short lines doesn't push + # K artificially high (the estimate would underweight the per-line + # newline + the trailing "...and N more in " notes). The fit + # check is preserved below — the estimate is approximate and Tier 3 + # remains the safety net. + sample = records[:32] + avg = max( + 80, + sum(len(p) + len(ln) + len(c) + 3 for p, ln, c in sample) // max(1, len(sample)), + ) + estimated_k = max(1, _SEARCH_OUTPUT_BUDGET // max(1, files * (avg + 1))) + k = min(_SEARCH_TIER2_SAMPLE_LADDER[0], estimated_k) + chunks2: list[str] = [] + used2 = 0 + fit = True + for path, matches in by_file.items(): + head = matches[:k] + for lineno, content in head: + line = f"{path}:{lineno}:{content}" + chunks2.append(line) + used2 += len(line) + 1 + if len(matches) > k: + note = f" ...and {len(matches) - k} more in {path}" + chunks2.append(note) + used2 += len(note) + 1 + if used2 > _SEARCH_OUTPUT_BUDGET: + fit = False + break + if fit: + header = ( + f"({total} matches across {files} files — " + f"showing first {k}/file. Narrow the query or read_file " + f"a specific path for full content.)" + ) + if capped: + header += " (raw output capped; counts may underreport.)" + return header + "\n\n" + "\n".join(chunks2) + + counts = sorted(by_file.items(), key=lambda kv: (-len(kv[1]), kv[0])) + body_lines: list[str] = [] + body_used = 0 + shown = 0 + for p, m in counts: + line = f"{p}: {len(m)} matches" + # Reserve room for the "(plus N more files)" tail line so we don't + # blow the budget when the count list itself is enormous. + if body_used + len(line) + 1 + _SEARCH_TIER3_TAIL_RESERVE > _SEARCH_OUTPUT_BUDGET: + break + body_lines.append(line) + body_used += len(line) + 1 + shown += 1 + if shown < files: + omitted_matches = sum(len(m) for _p, m in counts[shown:]) + body_lines.append( + f"(plus {files - shown} more files with {omitted_matches} matches between them)" + ) + body = "\n".join(body_lines) + header = ( + f"({total} matches across {files} files — too many to show inline. " + f"Counts only; narrow the query or read_file a specific path.)" + ) + if capped: + header += " (raw output capped; counts may underreport.)" + return header + "\n\n" + body + + # Memory scopes accepted by the ``memory`` tool's preparer + executor. # Single source of truth — every action validator imports this rather # than literal-listing the four values, so adding a fifth scope is a @@ -7862,69 +8118,153 @@ class ChatSession: self._report_tool_result(call_id, "read_file", f"image ({len(raw):,} bytes)") return call_id, content_parts + def _search_capture(self, args: list[str]) -> tuple[bytes, int, bytes, bool]: + """Run a search subprocess with a streaming, byte-capped stdout read. + + Returns ``(stdout, returncode, stderr, capped)``. Drains stderr in a + background thread to avoid pipe deadlock when the child writes a lot + to stderr while we're still reading stdout. + + The byte cap is the load-bearing defense for pathological inputs + (multi-GB JSONL, single-line minified bundles): on overflow, we + kill the child and trim to the last newline so the parser never + sees a partial trailing line. + """ + from turnstone.core.env import scrubbed_env + + proc = subprocess.Popen( + args, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + env=scrubbed_env(), + ) + + # Watchdog: ``proc.stdout.read`` is a blocking pipe read with no + # timeout, so a child stuck in kernel I/O (NFS, FUSE, broken + # backend) would hang forever despite ``tool_timeout``. The timer + # arms ``proc.kill`` after the deadline; we detect "child was + # killed by us, not by the byte cap" by setting ``timed_out`` + # before invoking kill. + timed_out = [False] + + def _watchdog() -> None: + timed_out[0] = True + with contextlib.suppress(Exception): + proc.kill() + + watchdog = threading.Timer(self.tool_timeout, _watchdog) + watchdog.daemon = True + watchdog.start() + + stderr_chunks: list[bytes] = [] + + def _drain_stderr() -> None: + if not proc.stderr: + return + total = 0 + try: + while total < _SEARCH_STDERR_CAP: + chunk = proc.stderr.read(_SEARCH_DRAIN_CHUNK) + if not chunk: + return + stderr_chunks.append(chunk) + total += len(chunk) + while proc.stderr.read(_SEARCH_DRAIN_CHUNK): + pass # discard tail so the child can finish writing + except Exception: + pass + + drain_thread = threading.Thread(target=_drain_stderr, daemon=True) + drain_thread.start() + + capped = False + try: + stdout = proc.stdout.read(_SEARCH_RAW_BYTE_CAP + 1) if proc.stdout else b"" + if len(stdout) > _SEARCH_RAW_BYTE_CAP: + capped = True + proc.kill() + stdout = stdout[:_SEARCH_RAW_BYTE_CAP] + last_nl = stdout.rfind(b"\n") + if last_nl >= 0: + stdout = stdout[:last_nl] + rc = proc.wait() + finally: + watchdog.cancel() + drain_thread.join(timeout=_SEARCH_DRAIN_JOIN_TIMEOUT) + for stream in (proc.stdout, proc.stderr): + try: + if stream is not None: + stream.close() + except Exception: + pass + + if timed_out[0]: + raise subprocess.TimeoutExpired(args, self.tool_timeout) + return stdout, rc, b"".join(stderr_chunks), capped + def _exec_search(self, item: dict[str, Any]) -> tuple[str, str]: - """Search file contents for a regex pattern using grep.""" + """Search file contents for a regex pattern via ripgrep (preferred) or grep.""" call_id = item["call_id"] pattern, path = item["pattern"], item["path"] try: - from turnstone.core.env import scrubbed_env + backend = _detect_search_backend() + args = _build_search_args(pattern, path, backend) + stdout, rc, stderr, capped = self._search_capture(args) - result = subprocess.run( - [ - "grep", - "-rn", - "-I", - "-E", - "-m", - "200", # max matches per file - "--color=never", # no ANSI codes in output - # Skip common build/vendor/VCS directories - "--exclude-dir=.git", - "--exclude-dir=node_modules", - "--exclude-dir=target", - "--exclude-dir=__pycache__", - "--exclude-dir=.mypy_cache", - "--exclude-dir=.ruff_cache", - "--exclude-dir=.pytest_cache", - "--exclude-dir=dist", - "--exclude-dir=build", - "--exclude-dir=*.egg-info", - "--exclude-dir=.tox", - "--exclude-dir=.venv", - "--exclude-dir=venv", - "--exclude-dir=vendor", - "--", - pattern, - path, # -- prevents pattern as flag - ], - capture_output=True, - text=True, - timeout=self.tool_timeout, - env=scrubbed_env(), - ) - output = result.stdout.strip() - if result.returncode == 1: - output = "(no matches)" - elif result.returncode > 1: - output = result.stderr.strip() or f"grep error (exit {result.returncode})" + # ripgrep and grep share rc semantics: 0 = matches, 1 = no + # matches, ≥2 = error. When ``capped`` is True we killed the + # child intentionally (byte-cap), so its negative rc is ours + # and we should treat the partial output as success. + if capped: + rc = 0 - # Count matches and files BEFORE truncation - match_count = output.count("\n") + 1 if result.returncode == 0 and output else 0 - if match_count: - files = {line.split(":", 1)[0] for line in output.splitlines() if ":" in line} - file_count = len(files) - else: - file_count = 0 + if rc == 1: + self._report_tool_result(call_id, "search", "no matches") + return call_id, "(no matches)" + if rc < 0: + # Signal-killed by something other than us (OOM killer, + # external SIGTERM). Surface it instead of parsing the + # truncated stdout as if the search had completed. + msg = f"{backend} killed by signal {-rc}" + self._report_tool_result(call_id, "search", msg, is_error=True) + return call_id, msg + if rc > 1: + err_text = stderr.decode("utf-8", errors="replace").strip() + msg = err_text or f"{backend} error (exit {rc})" + self._report_tool_result(call_id, "search", msg, is_error=True) + return call_id, msg - # Append summary footer before truncation so it counts toward the limit - original_len = len(output) - if match_count: - output += f"\n\n({match_count} matches across {file_count} files)" - output = self._truncate_output(output) + records = _parse_search_records(stdout) + original_len = len(stdout) - desc = f"{match_count} matches" if match_count else "no matches" + if not records: + if capped: + # The byte cap fired before any parseable line + # completed (typical shape: a single multi-MB line + # without a newline, e.g. minified bundle / training + # JSONL record). The malformed-output message would + # blame the query; surface the real cause instead. + msg = ( + "(search output exceeded the raw byte cap before " + "any parseable line completed — narrow your query " + "or restrict the path)" + ) + self._report_tool_result(call_id, "search", "byte cap hit", is_error=True) + return call_id, msg + # rc 0 with no parseable records means matches were found + # but every line was malformed. + self._report_tool_result(call_id, "search", "all matches malformed", is_error=True) + return call_id, _SEARCH_ALL_TRUNCATED_MSG + + output = _format_search_results(records, capped) + output = self._truncate_output(output) # belt-and-suspenders + + match_count = len(records) + desc = f"{match_count} matches" if original_len > 500: - desc += f" ({original_len} chars)" + desc += f" ({original_len} bytes raw)" + if capped: + desc += " [capped]" self._report_tool_result(call_id, "search", desc) return call_id, output