codemaster-cli 2.2.0__py3-none-any.whl

vibe/core/tools/builtins/bash.py

@@ -0,0 +1,357 @@
+from __future__ import annotations
+
+import asyncio
+from collections.abc import AsyncGenerator
+from functools import lru_cache
+import os
+import signal
+import sys
+from typing import ClassVar, Literal, final
+
+from pydantic import BaseModel, Field
+from tree_sitter import Language, Node, Parser
+import tree_sitter_bash as tsbash
+
+from vibe.core.tools.base import (
+    BaseTool,
+    BaseToolConfig,
+    BaseToolState,
+    InvokeContext,
+    ToolError,
+    ToolPermission,
+)
+from vibe.core.tools.ui import ToolCallDisplay, ToolResultDisplay, ToolUIData
+from vibe.core.types import ToolCallEvent, ToolResultEvent, ToolStreamEvent
+from vibe.core.utils import is_windows
+
+
+@lru_cache(maxsize=1)
+def _get_parser() -> Parser:
+    return Parser(Language(tsbash.language()))
+
+
+def _extract_commands(command: str) -> list[str]:
+    parser = _get_parser()
+    tree = parser.parse(command.encode("utf-8"))
+
+    commands: list[str] = []
+
+    def find_commands(node: Node) -> None:
+        if node.type == "command":
+            parts = []
+            for child in node.children:
+                if (
+                    child.type
+                    in {"command_name", "word", "string", "raw_string", "concatenation"}
+                    and child.text is not None
+                ):
+                    parts.append(child.text.decode("utf-8"))
+            if parts:
+                commands.append(" ".join(parts))
+
+        for child in node.children:
+            find_commands(child)
+
+    find_commands(tree.root_node)
+    return commands
+
+
+def _get_subprocess_encoding() -> str:
+    if sys.platform == "win32":
+        # Windows console uses OEM code page (e.g., cp850, cp1252)
+        import ctypes
+
+        return f"cp{ctypes.windll.kernel32.GetOEMCP()}"
+    return "utf-8"
+
+
+def _get_shell_executable() -> str | None:
+    if is_windows():
+        return None
+    return os.environ.get("SHELL")
+
+
+def _get_base_env() -> dict[str, str]:
+    base_env = {
+        **os.environ,
+        "CI": "true",
+        "NONINTERACTIVE": "1",
+        "NO_TTY": "1",
+        "NO_COLOR": "1",
+    }
+
+    if is_windows():
+        base_env["GIT_PAGER"] = "more"
+        base_env["PAGER"] = "more"
+    else:
+        base_env["TERM"] = "dumb"
+        base_env["DEBIAN_FRONTEND"] = "noninteractive"
+        base_env["GIT_PAGER"] = "cat"
+        base_env["PAGER"] = "cat"
+        base_env["LESS"] = "-FX"
+        base_env["LC_ALL"] = "en_US.UTF-8"
+
+    return base_env
+
+
+async def _kill_process_tree(proc: asyncio.subprocess.Process) -> None:
+    if proc.returncode is not None:
+        return
+
+    try:
+        if sys.platform == "win32":
+            try:
+                subprocess_proc = await asyncio.create_subprocess_exec(
+                    "taskkill",
+                    "/F",
+                    "/T",
+                    "/PID",
+                    str(proc.pid),
+                    stdout=asyncio.subprocess.DEVNULL,
+                    stderr=asyncio.subprocess.DEVNULL,
+                )
+                await subprocess_proc.wait()
+            except (FileNotFoundError, OSError):
+                proc.terminate()
+        else:
+            os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
+
+        await proc.wait()
+    except (ProcessLookupError, PermissionError, OSError):
+        pass
+
+
+def _get_default_allowlist() -> list[str]:
+    common = ["echo", "find", "git diff", "git log", "git status", "tree", "whoami"]
+
+    if is_windows():
+        return common + ["dir", "findstr", "more", "type", "ver", "where"]
+    else:
+        return common + [
+            "cat",
+            "file",
+            "head",
+            "ls",
+            "pwd",
+            "stat",
+            "tail",
+            "uname",
+            "wc",
+            "which",
+        ]
+
+
+def _get_default_denylist() -> list[str]:
+    common = ["gdb", "pdb", "passwd"]
+
+    if is_windows():
+        return common + ["cmd /k", "powershell -NoExit", "pwsh -NoExit", "notepad"]
+    else:
+        return common + [
+            "nano",
+            "vim",
+            "vi",
+            "emacs",
+            "bash -i",
+            "sh -i",
+            "zsh -i",
+            "fish -i",
+            "dash -i",
+            "screen",
+            "tmux",
+        ]
+
+
+def _get_default_denylist_standalone() -> list[str]:
+    common = ["python", "python3", "ipython"]
+
+    if is_windows():
+        return common + ["cmd", "powershell", "pwsh", "notepad"]
+    else:
+        return common + ["bash", "sh", "nohup", "vi", "vim", "emacs", "nano", "su"]
+
+
+class BashToolConfig(BaseToolConfig):
+    permission: ToolPermission = ToolPermission.ASK
+    max_output_bytes: int = Field(
+        default=16_000, description="Maximum bytes to capture from stdout and stderr."
+    )
+    default_timeout: int = Field(
+        default=300, description="Default timeout for commands in seconds."
+    )
+    allowlist: list[str] = Field(
+        default_factory=_get_default_allowlist,
+        description="Command prefixes that are automatically allowed",
+    )
+    denylist: list[str] = Field(
+        default_factory=_get_default_denylist,
+        description="Command prefixes that are automatically denied",
+    )
+    denylist_standalone: list[str] = Field(
+        default_factory=_get_default_denylist_standalone,
+        description="Commands that are denied only when run without arguments",
+    )
+
+
+class BashArgs(BaseModel):
+    command: str
+    timeout: int | None = Field(
+        default=None, description="Override the default command timeout."
+    )
+
+
+class BashResult(BaseModel):
+    command: str
+    stdout: str
+    stderr: str
+    returncode: int
+
+
+class Bash(
+    BaseTool[BashArgs, BashResult, BashToolConfig, BaseToolState],
+    ToolUIData[BashArgs, BashResult],
+):
+    description: ClassVar[str] = "Run a one-off bash command and capture its output."
+
+    @classmethod
+    def get_call_display(cls, event: ToolCallEvent) -> ToolCallDisplay:
+        if not isinstance(event.args, BashArgs):
+            return ToolCallDisplay(summary="bash")
+
+        return ToolCallDisplay(summary=f"bash: {event.args.command}")
+
+    @classmethod
+    def get_result_display(cls, event: ToolResultEvent) -> ToolResultDisplay:
+        if not isinstance(event.result, BashResult):
+            return ToolResultDisplay(
+                success=False, message=event.error or event.skip_reason or "No result"
+            )
+
+        return ToolResultDisplay(success=True, message=f"Ran {event.result.command}")
+
+    @classmethod
+    def get_status_text(cls) -> str:
+        return "Running command"
+
+    def check_allowlist_denylist(self, args: BashArgs) -> ToolPermission | None:
+        if is_windows():
+            return None
+
+        command_parts = _extract_commands(args.command)
+        if not command_parts:
+            return None
+
+        def is_denylisted(command: str) -> bool:
+            return any(command.startswith(pattern) for pattern in self.config.denylist)
+
+        def is_standalone_denylisted(command: str) -> bool:
+            parts = command.split()
+            if not parts:
+                return False
+
+            base_command = parts[0]
+            has_args = len(parts) > 1
+
+            if not has_args:
+                command_name = os.path.basename(base_command)
+                if command_name in self.config.denylist_standalone:
+                    return True
+                if base_command in self.config.denylist_standalone:
+                    return True
+
+            return False
+
+        def is_allowlisted(command: str) -> bool:
+            return any(command.startswith(pattern) for pattern in self.config.allowlist)
+
+        for part in command_parts:
+            if is_denylisted(part):
+                return ToolPermission.NEVER
+            if is_standalone_denylisted(part):
+                return ToolPermission.NEVER
+
+        if all(is_allowlisted(part) for part in command_parts):
+            return ToolPermission.ALWAYS
+
+        return None
+
+    @final
+    def _build_timeout_error(self, command: str, timeout: int) -> ToolError:
+        return ToolError(f"Command timed out after {timeout}s: {command!r}")
+
+    @final
+    def _build_result(
+        self, *, command: str, stdout: str, stderr: str, returncode: int
+    ) -> BashResult:
+        if returncode != 0:
+            error_msg = f"Command failed: {command!r}\n"
+            error_msg += f"Return code: {returncode}"
+            if stderr:
+                error_msg += f"\nStderr: {stderr}"
+            if stdout:
+                error_msg += f"\nStdout: {stdout}"
+            raise ToolError(error_msg.strip())
+
+        return BashResult(
+            command=command, stdout=stdout, stderr=stderr, returncode=returncode
+        )
+
+    async def run(
+        self, args: BashArgs, ctx: InvokeContext | None = None
+    ) -> AsyncGenerator[ToolStreamEvent | BashResult, None]:
+        timeout = args.timeout or self.config.default_timeout
+        max_bytes = self.config.max_output_bytes
+
+        proc = None
+        try:
+            # start_new_session is Unix-only, on Windows it's ignored
+            kwargs: dict[Literal["start_new_session"], bool] = (
+                {} if is_windows() else {"start_new_session": True}
+            )
+
+            proc = await asyncio.create_subprocess_shell(
+                args.command,
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.PIPE,
+                stdin=asyncio.subprocess.DEVNULL,
+                env=_get_base_env(),
+                executable=_get_shell_executable(),
+                **kwargs,
+            )
+
+            try:
+                stdout_bytes, stderr_bytes = await asyncio.wait_for(
+                    proc.communicate(), timeout=timeout
+                )
+            except TimeoutError:
+                await _kill_process_tree(proc)
+                raise self._build_timeout_error(args.command, timeout)
+
+            encoding = _get_subprocess_encoding()
+            stdout = (
+                stdout_bytes.decode(encoding, errors="replace")[:max_bytes]
+                if stdout_bytes
+                else ""
+            )
+            stderr = (
+                stderr_bytes.decode(encoding, errors="replace")[:max_bytes]
+                if stderr_bytes
+                else ""
+            )
+
+            returncode = proc.returncode or 0
+
+            yield self._build_result(
+                command=args.command,
+                stdout=stdout,
+                stderr=stderr,
+                returncode=returncode,
+            )
+
+        except (ToolError, asyncio.CancelledError):
+            raise
+        except Exception as exc:
+            raise ToolError(f"Error running command {args.command!r}: {exc}") from exc
+        finally:
+            if proc is not None:
+                await _kill_process_tree(proc)

vibe/core/tools/builtins/grep.py

@@ -0,0 +1,310 @@
+from __future__ import annotations
+
+import asyncio
+from collections.abc import AsyncGenerator
+from enum import StrEnum, auto
+from pathlib import Path
+import shutil
+from typing import TYPE_CHECKING, ClassVar
+
+from pydantic import BaseModel, Field
+
+from vibe.core.tools.base import (
+    BaseTool,
+    BaseToolConfig,
+    BaseToolState,
+    InvokeContext,
+    ToolError,
+    ToolPermission,
+)
+from vibe.core.tools.ui import ToolCallDisplay, ToolResultDisplay, ToolUIData
+from vibe.core.types import ToolStreamEvent
+
+if TYPE_CHECKING:
+    from vibe.core.types import ToolCallEvent, ToolResultEvent
+
+
+class GrepBackend(StrEnum):
+    RIPGREP = auto()
+    GNU_GREP = auto()
+
+
+class GrepToolConfig(BaseToolConfig):
+    permission: ToolPermission = ToolPermission.ALWAYS
+
+    max_output_bytes: int = Field(
+        default=64_000, description="Hard cap for the total size of matched lines."
+    )
+    default_max_matches: int = Field(
+        default=100, description="Default maximum number of matches to return."
+    )
+    default_timeout: int = Field(
+        default=60, description="Default timeout for the search command in seconds."
+    )
+    exclude_patterns: list[str] = Field(
+        default=[
+            ".venv/",
+            "venv/",
+            ".env/",
+            "env/",
+            "node_modules/",
+            ".git/",
+            "__pycache__/",
+            ".pytest_cache/",
+            ".mypy_cache/",
+            ".tox/",
+            ".nox/",
+            ".coverage/",
+            "htmlcov/",
+            "dist/",
+            "build/",
+            ".idea/",
+            ".vscode/",
+            "*.egg-info",
+            "*.pyc",
+            "*.pyo",
+            "*.pyd",
+            ".DS_Store",
+            "Thumbs.db",
+        ],
+        description="List of glob patterns to exclude from search (dirs should end with /).",
+    )
+    codeignore_file: str = Field(
+        default=".vibeignore",
+        description="Name of the file to read for additional exclusion patterns.",
+    )
+
+
+class GrepState(BaseToolState):
+    search_history: list[str] = Field(default_factory=list)
+
+
+class GrepArgs(BaseModel):
+    pattern: str
+    path: str = "."
+    max_matches: int | None = Field(
+        default=None, description="Override the default maximum number of matches."
+    )
+    use_default_ignore: bool = Field(
+        default=True, description="Whether to respect .gitignore and .ignore files."
+    )
+
+
+class GrepResult(BaseModel):
+    matches: str
+    match_count: int
+    was_truncated: bool = Field(
+        description="True if output was cut short by max_matches or max_output_bytes."
+    )
+
+
+class Grep(
+    BaseTool[GrepArgs, GrepResult, GrepToolConfig, GrepState],
+    ToolUIData[GrepArgs, GrepResult],
+):
+    description: ClassVar[str] = (
+        "Recursively search files for a regex pattern using ripgrep (rg) or grep. "
+        "Respects .gitignore and .codeignore files by default when using ripgrep."
+    )
+
+    def _detect_backend(self) -> GrepBackend:
+        if shutil.which("rg"):
+            return GrepBackend.RIPGREP
+        if shutil.which("grep"):
+            return GrepBackend.GNU_GREP
+        raise ToolError(
+            "Neither ripgrep (rg) nor grep is installed. "
+            "Please install ripgrep: https://github.com/BurntSushi/ripgrep#installation"
+        )
+
+    async def run(
+        self, args: GrepArgs, ctx: InvokeContext | None = None
+    ) -> AsyncGenerator[ToolStreamEvent | GrepResult, None]:
+        backend = self._detect_backend()
+        self._validate_args(args)
+        self.state.search_history.append(args.pattern)
+
+        exclude_patterns = self._collect_exclude_patterns()
+        cmd = self._build_command(args, exclude_patterns, backend)
+        stdout = await self._execute_search(cmd)
+
+        yield self._parse_output(
+            stdout, args.max_matches or self.config.default_max_matches
+        )
+
+    def _validate_args(self, args: GrepArgs) -> None:
+        if not args.pattern.strip():
+            raise ToolError("Empty search pattern provided.")
+
+        path_obj = Path(args.path).expanduser()
+        if not path_obj.is_absolute():
+            path_obj = Path.cwd() / path_obj
+
+        if not path_obj.exists():
+            raise ToolError(f"Path does not exist: {args.path}")
+
+    def _collect_exclude_patterns(self) -> list[str]:
+        patterns = list(self.config.exclude_patterns)
+
+        codeignore_path = Path.cwd() / self.config.codeignore_file
+        if codeignore_path.is_file():
+            patterns.extend(self._load_codeignore_patterns(codeignore_path))
+
+        return patterns
+
+    def _load_codeignore_patterns(self, codeignore_path: Path) -> list[str]:
+        patterns = []
+        try:
+            content = codeignore_path.read_text("utf-8")
+            for line in content.splitlines():
+                line = line.strip()
+                if line and not line.startswith("#"):
+                    patterns.append(line)
+        except OSError:
+            pass
+
+        return patterns
+
+    def _build_command(
+        self, args: GrepArgs, exclude_patterns: list[str], backend: GrepBackend
+    ) -> list[str]:
+        if backend == GrepBackend.RIPGREP:
+            return self._build_ripgrep_command(args, exclude_patterns)
+        return self._build_gnu_grep_command(args, exclude_patterns)
+
+    def _build_ripgrep_command(
+        self, args: GrepArgs, exclude_patterns: list[str]
+    ) -> list[str]:
+        max_matches = args.max_matches or self.config.default_max_matches
+
+        cmd = [
+            "rg",
+            "--line-number",
+            "--no-heading",
+            "--smart-case",
+            "--no-binary",
+            # Request one extra to detect truncation
+            "--max-count",
+            str(max_matches + 1),
+        ]
+
+        if not args.use_default_ignore:
+            cmd.append("--no-ignore")
+
+        for pattern in exclude_patterns:
+            cmd.extend(["--glob", f"!{pattern}"])
+
+        cmd.extend(["-e", args.pattern, args.path])
+
+        return cmd
+
+    def _build_gnu_grep_command(
+        self, args: GrepArgs, exclude_patterns: list[str]
+    ) -> list[str]:
+        max_matches = args.max_matches or self.config.default_max_matches
+
+        cmd = ["grep", "-r", "-n", "-I", "-E", f"--max-count={max_matches + 1}"]
+
+        if args.pattern.islower():
+            cmd.append("-i")
+
+        for pattern in exclude_patterns:
+            if pattern.endswith("/"):
+                dir_pattern = pattern.rstrip("/")
+                cmd.append(f"--exclude-dir={dir_pattern}")
+            else:
+                cmd.append(f"--exclude={pattern}")
+
+        cmd.extend(["-e", args.pattern, args.path])
+
+        return cmd
+
+    async def _execute_search(self, cmd: list[str]) -> str:
+        try:
+            proc = await asyncio.create_subprocess_exec(
+                *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE
+            )
+
+            try:
+                stdout_bytes, stderr_bytes = await asyncio.wait_for(
+                    proc.communicate(), timeout=self.config.default_timeout
+                )
+            except TimeoutError:
+                proc.kill()
+                await proc.wait()
+                raise ToolError(
+                    f"Search timed out after {self.config.default_timeout}s"
+                )
+
+            stdout = (
+                stdout_bytes.decode("utf-8", errors="ignore") if stdout_bytes else ""
+            )
+            stderr = (
+                stderr_bytes.decode("utf-8", errors="ignore") if stderr_bytes else ""
+            )
+
+            if proc.returncode not in {0, 1}:
+                error_msg = stderr or f"Process exited with code {proc.returncode}"
+                raise ToolError(f"grep error: {error_msg}")
+
+            return stdout
+
+        except ToolError:
+            raise
+        except Exception as exc:
+            raise ToolError(f"Error running grep: {exc}") from exc
+
+    def _parse_output(self, stdout: str, max_matches: int) -> GrepResult:
+        output_lines = stdout.splitlines() if stdout else []
+
+        truncated_lines = output_lines[:max_matches]
+        truncated_output = "\n".join(truncated_lines)
+
+        was_truncated = (
+            len(output_lines) > max_matches
+            or len(truncated_output) > self.config.max_output_bytes
+        )
+
+        final_output = truncated_output[: self.config.max_output_bytes]
+
+        return GrepResult(
+            matches=final_output,
+            match_count=len(truncated_lines),
+            was_truncated=was_truncated,
+        )
+
+    @classmethod
+    def get_call_display(cls, event: ToolCallEvent) -> ToolCallDisplay:
+        if not isinstance(event.args, GrepArgs):
+            return ToolCallDisplay(summary="grep")
+
+        summary = f"Grepping '{event.args.pattern}'"
+        if event.args.path != ".":
+            summary += f" in {event.args.path}"
+        if event.args.max_matches:
+            summary += f" (max {event.args.max_matches} matches)"
+        if not event.args.use_default_ignore:
+            summary += " [no-ignore]"
+
+        return ToolCallDisplay(summary=summary)
+
+    @classmethod
+    def get_result_display(cls, event: ToolResultEvent) -> ToolResultDisplay:
+        if not isinstance(event.result, GrepResult):
+            return ToolResultDisplay(
+                success=False, message=event.error or event.skip_reason or "No result"
+            )
+
+        message = f"Found {event.result.match_count} matches"
+        if event.result.was_truncated:
+            message += " (truncated)"
+
+        warnings = []
+        if event.result.was_truncated:
+            warnings.append("Output was truncated due to size/match limits")
+
+        return ToolResultDisplay(success=True, message=message, warnings=warnings)
+
+    @classmethod
+    def get_status_text(cls) -> str:
+        return "Searching files"