minimalcode 1.5.2__py3-none-any.whl

minimalcode/tools/search.py

@@ -0,0 +1,521 @@
+"""Search primitives for the grep and find tools.
+
+The same ignore rules and regex semantics are applied by two engines: ripgrep
+when it can be spawned, and a bounded Python fallback otherwise. This module
+owns process launch and output parsing so the runtime only has to choose an
+engine and shape the result.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import queue
+import re
+import signal
+import subprocess
+import sys
+import tempfile
+import threading
+import time
+from pathlib import Path
+from typing import Any
+
+from ..errors import ToolError
+from ..sandbox import SandboxRunner
+
+# Shared by the Python fallback walk and the ripgrep invocations so both
+# implementations search the same set of files.
+IGNORED_DIRECTORIES: tuple[str, ...] = (
+    ".git",
+    ".hg",
+    ".svn",
+    ".mypy_cache",
+    ".pytest_cache",
+    ".ruff_cache",
+    "__pycache__",
+    ".venv",
+    "build",
+    "dist",
+    "node_modules",
+    "venv",
+)
+IGNORED_GLOBS: tuple[str, ...] = tuple(
+    glob for name in IGNORED_DIRECTORIES for glob in (f"!{name}/**", f"!**/{name}/**")
+)
+
+# Match ripgrep's --max-columns output width and keep the Python fallback from
+# returning a single minified line that swallows the whole context budget.
+_MAX_GREP_LINE_CHARS = 300
+_MAX_READ_LINE_CHARS = 1_000_000
+_SEARCH_TIMEOUT_SECONDS = 30
+_SEARCH_MAX_OUTPUT_CHARS = 1_000_000
+
+# The fallback runs in a child interpreter so a pathological regex cannot block
+# the agent process. It receives one JSON document on stdin and emits one JSON
+# document on stdout; line counting deliberately ignores fragments of a line
+# longer than the per-line cap.
+_PYTHON_GREP_SCRIPT = r"""
+import json
+import os
+import re
+import sys
+
+spec = json.load(sys.stdin)
+try:
+    regex = re.compile(
+        spec["pattern"], re.IGNORECASE if spec["ignore_case"] else 0
+    )
+except re.error as exc:
+    print(json.dumps({"error": f"invalid regular expression: {exc}"}))
+    raise SystemExit(0)
+
+matches = []
+for path in spec["files"]:
+    try:
+        if os.path.getsize(path) > spec["max_bytes"]:
+            continue
+        with open(path, "r", encoding="utf-8", errors="replace") as handle:
+            number = 0
+            while True:
+                line = handle.readline(spec["max_line_chars"])
+                if not line:
+                    break
+                if line.endswith("\r") and len(line) >= spec["max_line_chars"]:
+                    # The limit may land between CR and LF; consume the LF so
+                    # it cannot be counted as an extra empty line.
+                    position = handle.tell()
+                    following = handle.read(1)
+                    if following == "\n":
+                        line += following
+                    else:
+                        handle.seek(position)
+                number += 1
+                cut = len(line) >= spec["max_line_chars"] and not line.endswith(
+                    ("\n", "\r")
+                )
+                if cut:
+                    # A line of exactly the limit may still be complete; look at
+                    # the following character before dropping its tail.
+                    position = handle.tell()
+                    following = handle.read(1)
+                    if not following:
+                        cut = False
+                    elif following == "\n":
+                        line += following
+                        cut = False
+                    elif following == "\r":
+                        line += following
+                        peek = handle.read(1)
+                        if peek == "\n":
+                            line += peek
+                        elif peek:
+                            handle.seek(position + 1)
+                        cut = False
+                    else:
+                        handle.seek(position)
+                if cut:
+                    # Swallow the rest of this physical line so the next
+                    # readline starts a new line and `number` stays aligned
+                    # with the file instead of counting each fragment.
+                    while True:
+                        rest = handle.readline(spec["max_line_chars"])
+                        if not rest:
+                            break
+                        if rest.endswith("\n"):
+                            break
+                        if rest.endswith("\r") and len(rest) >= spec["max_line_chars"]:
+                            rest_position = handle.tell()
+                            rest_following = handle.read(1)
+                            if rest_following != "\n":
+                                handle.seek(rest_position)
+                                break
+                            break
+                        if rest.endswith("\r"):
+                            break
+                if "\x00" in line:
+                    continue
+                # Match the line content, not its terminator. `$` and `\s` would
+                # otherwise see the newline and report matches ripgrep never
+                # finds (`\s+$`, `[^0-9]$`), and a CRLF line would hide a real
+                # `c$` match.
+                text = line.rstrip("\r\n")
+                if not regex.search(text):
+                    continue
+                if cut or len(text) > spec["max_columns"]:
+                    text = text[: spec["max_columns"]] + " [... omitted end of long line]"
+                matches.append(f"{path}:{number}:{text}")
+                # One match past the limit proves truncation; an exact fit is a
+                # complete result, matching the ripgrep and inline engines.
+                if len(matches) > spec["limit"]:
+                    print(
+                        json.dumps(
+                            {"matches": matches[: spec["limit"]], "truncated": True}
+                        )
+                    )
+                    raise SystemExit(0)
+    except (OSError, UnicodeError):
+        continue
+
+print(json.dumps({"matches": matches, "truncated": False}))
+"""
+
+
+def smart_case_ignore(pattern: str) -> bool:
+    """Return True when ripgrep's --smart-case would ignore case."""
+    literal: list[str] = []
+    index = 0
+    while index < len(pattern):
+        if pattern[index] == "\\" and index + 1 < len(pattern):
+            # An escaped character carries no case of its own. Skipping the
+            # pair keeps `\\D` (an escaped backslash plus a literal `D`) visible
+            # to the check, which stripping every letter escape hid.
+            index += 2
+            continue
+        literal.append(pattern[index])
+        index += 1
+    return not any(char.isupper() for char in literal)
+
+
+def drain_physical_line(handle: Any, limit: int) -> None:
+    """Consume the rest of a line that was longer than ``limit``."""
+    while True:
+        rest = handle.readline(limit)
+        if not rest:
+            return
+        if rest.endswith("\n"):
+            return
+        if rest.endswith("\r") and len(rest) >= limit:
+            # A CRLF split at the limit must consume the LF too; otherwise it
+            # becomes an empty line and shifts every later line number.
+            position = handle.tell()
+            following = handle.read(1)
+            if following != "\n":
+                handle.seek(position)
+                return
+            return
+        if rest.endswith("\r"):
+            return
+
+
+def _kill_process_tree(process: subprocess.Popen[bytes]) -> None:
+    """Kill a search child and any descendant that inherited its stdout."""
+    if os.name == "nt":
+        try:
+            subprocess.run(
+                ["taskkill", "/F", "/T", "/PID", str(process.pid)],
+                stdin=subprocess.DEVNULL,
+                stdout=subprocess.DEVNULL,
+                stderr=subprocess.DEVNULL,
+                timeout=10,
+                check=False,
+            )
+        except (OSError, subprocess.TimeoutExpired):
+            pass
+    else:
+        try:
+            os.killpg(os.getpgid(process.pid), signal.SIGKILL)
+        except OSError:
+            pass
+    try:
+        process.kill()
+    except OSError:
+        pass
+
+
+def inline_guard_available() -> bool:
+    """Return True when the in-process regex fallback can enforce a deadline.
+
+    ``inline_python_grep`` uses ``SIGALRM``, which exists only on POSIX and is
+    delivered only to the main thread. The agent runs adjacent read-only tools
+    in worker threads, so callers must not pick that fallback when the guard is
+    unavailable.
+    """
+    return hasattr(signal, "SIGALRM") and (
+        threading.current_thread() is threading.main_thread()
+    )
+
+
+def run_bounded_command(
+    command: list[str],
+    *,
+    cwd: Path,
+    limit: int,
+    timeout: float = _SEARCH_TIMEOUT_SECONDS,
+    sandbox: SandboxRunner | None = None,
+) -> tuple[int, list[str], bool, str]:
+    """Run a search process while keeping only ``limit`` output lines.
+
+    The result is ``(returncode, lines, truncated, stderr)``. Reading stops as
+    soon as one extra line proves the result is truncated, then the child is
+    killed so a repository-wide search cannot outrun the caller.
+    """
+    if sandbox is not None:
+        result = sandbox.run_argv(
+            command,
+            cwd=cwd,
+            timeout=timeout,
+            max_output_chars=_SEARCH_MAX_OUTPUT_CHARS,
+        )
+        if result.get("status") == "timeout":
+            raise ToolError(f"search timed out after {timeout:g} seconds")
+        if result.get("status") == "error" and result.get("exit_code") is None:
+            detail = str(
+                result.get("message") or result.get("stderr") or "search failed"
+            )
+            raise ToolError(detail)
+        # `str.splitlines()` also breaks on form feeds, U+2028, and a lone CR
+        # inside a matched line, so the tail of the line was dropped as if it
+        # were noise. Search output is newline-delimited.
+        lines = [
+            line.rstrip("\r")
+            for line in str(result.get("stdout") or "").split("\n")
+            if line
+        ]
+        truncated = bool(result.get("truncated")) or len(lines) > limit
+        stderr = str(result.get("stderr") or "").strip()
+        return int(result.get("exit_code", 0) or 0), lines[:limit], truncated, stderr
+
+    timed_out = threading.Event()
+    process: subprocess.Popen[bytes] | None = None
+    with tempfile.TemporaryFile() as err_file:
+        try:
+            process = subprocess.Popen(
+                command,
+                cwd=cwd,
+                stdin=subprocess.DEVNULL,
+                stdout=subprocess.PIPE,
+                stderr=err_file,
+                # A process group lets the timeout kill grandchildren that keep
+                # the stdout pipe open after the direct child exited.
+                start_new_session=os.name != "nt",
+            )
+        except (OSError, ValueError) as exc:
+            raise ToolError(f"search failed: {exc}") from exc
+
+        assert process.stdout is not None
+        # Reading in a daemon thread keeps a grandchild that inherited stdout
+        # from blocking the caller past the deadline. The direct child can exit
+        # first, and on Windows the surviving grandchild is no longer reachable
+        # through taskkill's process tree.
+        output: queue.Queue[bytes | None] = queue.Queue()
+        reader_done = threading.Event()
+
+        def read_output() -> None:
+            try:
+                for _ in range(limit + 2):
+                    raw = process.stdout.readline()
+                    if not raw:
+                        break
+                    output.put(raw)
+            except (OSError, ValueError):
+                pass
+            finally:
+                output.put(None)
+                reader_done.set()
+
+        reader = threading.Thread(target=read_output, daemon=True)
+        reader.start()
+        lines: list[str] = []
+        truncated = False
+        deadline = time.monotonic() + timeout
+        while len(lines) <= limit:
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                timed_out.set()
+                break
+            try:
+                raw = output.get(timeout=remaining)
+            except queue.Empty:
+                timed_out.set()
+                break
+            if raw is None:
+                break
+            lines.append(raw.decode("utf-8", errors="replace").rstrip("\r\n"))
+        if len(lines) > limit:
+            truncated = True
+            lines = lines[:limit]
+            _kill_process_tree(process)
+        elif timed_out.is_set():
+            _kill_process_tree(process)
+        try:
+            process.wait(timeout=5)
+        except subprocess.TimeoutExpired:
+            timed_out.set()
+            _kill_process_tree(process)
+        # Closing the pipe is safe only once the reader stopped using it; a
+        # blocked read is left to its daemon thread instead of deadlocking.
+        if reader_done.wait(1.0):
+            process.stdout.close()
+        err_file.seek(0)
+        stderr = err_file.read(4000).decode("utf-8", errors="replace").strip()
+        if timed_out.is_set():
+            raise ToolError(f"search timed out after {timeout:g} seconds")
+        return process.returncode, lines, truncated, stderr
+
+
+def _parse_python_grep(stdout: str, limit: int) -> dict[str, Any]:
+    """Validate the fallback child's JSON and cap its match list."""
+    try:
+        parsed = json.loads(stdout)
+    except json.JSONDecodeError as exc:
+        raise ToolError("search returned an invalid result") from exc
+    if not isinstance(parsed, dict):
+        raise ToolError("search returned an invalid result")
+    if parsed.get("error"):
+        raise ToolError(str(parsed["error"]))
+    matches = parsed.get("matches")
+    if not isinstance(matches, list):
+        raise ToolError("search returned an invalid result")
+    return {
+        "status": "success",
+        "matches": [str(item) for item in matches[:limit]],
+        "truncated": bool(parsed.get("truncated")),
+    }
+
+
+def run_python_grep(
+    *,
+    pattern: str,
+    files: list[Path],
+    limit: int,
+    max_bytes: int,
+    sandbox: SandboxRunner | None = None,
+    cwd: Path | None = None,
+) -> dict[str, Any]:
+    """Run the regex fallback in a child process so a bad pattern can time out."""
+    spec = {
+        "pattern": pattern,
+        "ignore_case": smart_case_ignore(pattern),
+        "files": [str(path) for path in files],
+        "limit": limit,
+        "max_bytes": max_bytes,
+        "max_line_chars": min(max_bytes + 1, _MAX_READ_LINE_CHARS),
+        "max_columns": _MAX_GREP_LINE_CHARS,
+    }
+    script = [sys.executable, "-c", _PYTHON_GREP_SCRIPT]
+    if sandbox is not None:
+        # The child emits one bounded line per match, so the JSON envelope is
+        # larger than the default search cap. A truncated envelope is invalid
+        # JSON and used to surface as a bogus "invalid result" error.
+        output_cap = max(
+            _SEARCH_MAX_OUTPUT_CHARS,
+            limit * (_MAX_GREP_LINE_CHARS + 4096) + 4096,
+        )
+        result = sandbox.run_argv(
+            script,
+            cwd=cwd or Path.cwd(),
+            timeout=_SEARCH_TIMEOUT_SECONDS,
+            max_output_chars=output_cap,
+            stdin_text=json.dumps(spec),
+        )
+        if result.get("status") == "timeout":
+            raise ToolError("search timed out")
+        if result.get("status") == "error" and result.get("exit_code") is None:
+            detail = str(
+                result.get("message") or result.get("stderr") or "search failed"
+            )
+            raise ToolError(detail)
+        returncode = int(result.get("exit_code", 0) or 0)
+        stdout = str(result.get("stdout") or "")
+        stderr = str(result.get("stderr") or "").strip()
+    else:
+        try:
+            completed = subprocess.run(
+                script,
+                input=json.dumps(spec),
+                capture_output=True,
+                text=True,
+                encoding="utf-8",
+                timeout=_SEARCH_TIMEOUT_SECONDS,
+                check=False,
+            )
+        except subprocess.TimeoutExpired as exc:
+            raise ToolError("search timed out") from exc
+        except OSError as exc:
+            raise ToolError(f"search failed: {exc}") from exc
+        returncode = completed.returncode
+        stdout = completed.stdout
+        stderr = completed.stderr.strip()
+    if returncode != 0:
+        detail = stderr or f"search exited {returncode}"
+        raise ToolError(detail)
+    return _parse_python_grep(stdout, limit)
+
+
+def inline_python_grep(
+    *,
+    pattern: str,
+    files: list[Path],
+    limit: int,
+    max_bytes: int,
+) -> dict[str, Any]:
+    """Pure-Python fallback used when the backend cannot run argv children.
+
+    A backend without argv support runs the regular expression in this process,
+    so a bounded alarm is the only thing that keeps a catastrophic pattern from
+    hanging the agent. Windows has no SIGALRM; the docker backend that needs
+    this path is POSIX-only, so the guard is best-effort elsewhere.
+    """
+    flags = re.IGNORECASE if smart_case_ignore(pattern) else 0
+    try:
+        regex = re.compile(pattern, flags)
+    except re.error as exc:
+        raise ToolError(f"invalid regular expression: {exc}") from exc
+
+    use_alarm = (
+        hasattr(signal, "SIGALRM")
+        and threading.current_thread() is threading.main_thread()
+    )
+    previous_handler: Any = None
+    if use_alarm:
+
+        def _timeout(_signum: int, _frame: Any) -> None:
+            raise ToolError(
+                f"search timed out after {_SEARCH_TIMEOUT_SECONDS:g} seconds"
+            )
+
+        previous_handler = signal.signal(signal.SIGALRM, _timeout)
+        signal.setitimer(signal.ITIMER_REAL, _SEARCH_TIMEOUT_SECONDS)
+
+    matches: list[str] = []
+    try:
+        for path in files:
+            try:
+                if path.stat().st_size > max_bytes:
+                    continue
+                with path.open("r", encoding="utf-8", errors="replace") as handle:
+                    for number, line in enumerate(handle, start=1):
+                        if "\x00" in line:
+                            continue
+                        # Match the content without the terminator so `$` and
+                        # `\s` behave like the ripgrep engine.
+                        text = line.rstrip("\r\n")
+                        if not regex.search(text):
+                            continue
+                        if len(text) > _MAX_GREP_LINE_CHARS:
+                            text = (
+                                text[:_MAX_GREP_LINE_CHARS]
+                                + " [... omitted end of long line]"
+                            )
+                        matches.append(f"{path}:{number}:{text}")
+                        if len(matches) > limit:
+                            return {
+                                "status": "success",
+                                "matches": matches[:limit],
+                                "truncated": True,
+                            }
+            except (OSError, UnicodeError):
+                # An unreadable file must not abort a repository-wide search.
+                continue
+    finally:
+        if use_alarm:
+            signal.setitimer(signal.ITIMER_REAL, 0)
+            # `getsignal` returns None when the handler was installed outside
+            # Python; passing that back raises TypeError and would leave our
+            # timeout handler installed for the rest of the process.
+            signal.signal(
+                signal.SIGALRM,
+                previous_handler if previous_handler is not None else signal.SIG_DFL,
+            )
+    return {"status": "success", "matches": matches[:limit], "truncated": False}