Agent/agent/tools/builtin/bash.py

"""
Bash Tool - 命令执行工具

参考 OpenCode bash.ts 完整实现。

核心功能：
- 执行 shell 命令
- 超时控制
- 工作目录设置
- 环境变量传递
"""

import os
import signal
import asyncio
from pathlib import Path
from typing import Optional, Dict

from agent.tools import tool, ToolResult, ToolContext

# 常量
DEFAULT_TIMEOUT = 120  # 2 分钟
MAX_OUTPUT_LENGTH = 50000  # 最大输出长度
GRACEFUL_KILL_WAIT = 3  # SIGTERM 后等几秒再 SIGKILL


def _kill_process_tree(pid: int) -> None:
    """先 SIGTERM 整个进程组，等 GRACEFUL_KILL_WAIT 秒后 SIGKILL 兜底。"""
    import time

    try:
        pgid = os.getpgid(pid)
    except ProcessLookupError:
        return

    # 先优雅终止
    try:
        os.killpg(pgid, signal.SIGTERM)
    except ProcessLookupError:
        return

    # 等一小段时间让进程自行退出
    time.sleep(GRACEFUL_KILL_WAIT)

    # 强制杀
    try:
        os.killpg(pgid, signal.SIGKILL)
    except ProcessLookupError:
        pass  # 已退出


@tool(description="执行 bash 命令")
async def bash_command(
    command: str,
    timeout: Optional[int] = None,
    workdir: Optional[str] = None,
    env: Optional[Dict[str, str]] = None,
    description: str = "",
    context: Optional[ToolContext] = None
) -> ToolResult:
    """
    执行 bash 命令

    Args:
        command: 要执行的命令
        timeout: 超时时间（秒），默认 120 秒
        workdir: 工作目录，默认为当前目录
        env: 环境变量字典（会合并到系统环境变量）
        description: 命令描述（5-10 个词）
        context: 工具上下文

    Returns:
        ToolResult: 命令输出
    """
    # 参数验证
    if timeout is not None and timeout < 0:
        return ToolResult(
            title="参数错误",
            output=f"无效的 timeout 值: {timeout}。必须是正数。",
            error="Invalid timeout"
        )

    timeout_sec = timeout or DEFAULT_TIMEOUT

    # 工作目录
    cwd = Path(workdir) if workdir else Path.cwd()
    if not cwd.exists():
        return ToolResult(
            title="目录不存在",
            output=f"工作目录不存在: {workdir}",
            error="Directory not found"
        )

    # 准备环境变量
    process_env = os.environ.copy()
    if env:
        process_env.update(env)

    # 执行命令
    try:
        process = await asyncio.create_subprocess_shell(
            command,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            cwd=str(cwd),
            env=process_env,
            start_new_session=True,  # 新进程组，超时时可杀整棵进程树
        )

        # 等待命令完成（带超时）
        try:
            stdout, stderr = await asyncio.wait_for(
                process.communicate(),
                timeout=timeout_sec
            )
        except asyncio.TimeoutError:
            # 超时，杀整个进程组（shell + 所有子进程）
            _kill_process_tree(process.pid)
            try:
                await asyncio.wait_for(process.wait(), timeout=GRACEFUL_KILL_WAIT + 2)
            except asyncio.TimeoutError:
                pass
            return ToolResult(
                title="命令超时",
                output=f"命令执行超时（>{timeout_sec}s）: {command[:100]}",
                error="Timeout",
                metadata={"command": command, "timeout": timeout_sec}
            )

        # 解码输出
        stdout_text = stdout.decode('utf-8', errors='replace') if stdout else ""
        stderr_text = stderr.decode('utf-8', errors='replace') if stderr else ""

        # 截断过长输出
        truncated = False
        if len(stdout_text) > MAX_OUTPUT_LENGTH:
            stdout_text = stdout_text[:MAX_OUTPUT_LENGTH] + f"\n\n(输出被截断，总长度: {len(stdout_text)} 字符)"
            truncated = True

        # 组合输出
        output = ""
        if stdout_text:
            output += stdout_text

        if stderr_text:
            if output:
                output += "\n\n--- stderr ---\n"
            output += stderr_text

        if not output:
            output = "(命令无输出)"

        # 检查退出码
        exit_code = process.returncode
        success = exit_code == 0

        title = description or f"命令: {command[:50]}"
        if not success:
            title += f" (exit code: {exit_code})"

        return ToolResult(
            title=title,
            output=output,
            metadata={
                "exit_code": exit_code,
                "success": success,
                "truncated": truncated,
                "command": command,
                "cwd": str(cwd)
            },
            error=None if success else f"Command failed with exit code {exit_code}"
        )

    except Exception as e:
        return ToolResult(
            title="执行错误",
            output=f"命令执行失败: {str(e)}",
            error=str(e),
            metadata={"command": command}
        )