Agent/agent/tools/builtin/sandbox.py

"""
Sandbox Tools (Async)
通过 HTTP 异步调用沙盒管理服务的客户端库。
"""

import json
import httpx
from typing import Optional, List, Dict, Any

from agent.tools import tool, ToolResult, ToolContext


# 服务地址，可根据实际部署情况修改
# SANDBOX_SERVER_URL = "http://192.168.100.20:9998"
SANDBOX_SERVER_URL = "http://61.48.133.26:9998"

# 默认超时时间（秒）
DEFAULT_TIMEOUT = 300.0


@tool(
    display={
        "zh": {
            "name": "创建沙盒环境",
            "params": {
                "image": "Docker 镜像",
                "mem_limit": "内存限制",
                "nano_cpus": "CPU 限制",
                "ports": "端口列表",
                "use_gpu": "启用 GPU",
                "gpu_count": "GPU 数量"
            }
        },
        "en": {
            "name": "Create Sandbox",
            "params": {
                "image": "Docker image",
                "mem_limit": "Memory limit",
                "nano_cpus": "CPU limit",
                "ports": "Port list",
                "use_gpu": "Enable GPU",
                "gpu_count": "GPU count"
            }
        }
    }
)
async def sandbox_create_environment(
    image: str = "agent-sandbox:latest",
    mem_limit: str = "512m",
    nano_cpus: int = 500000000,
    ports: Optional[List[int]] = None,
    use_gpu: bool = False,
    gpu_count: int = -1,
    server_url: str = None,
    timeout: float = DEFAULT_TIMEOUT,
    context: Optional[ToolContext] = None,
) -> ToolResult:
    """
    创建一个隔离的 Docker 开发环境。

    Args:
        image: Docker 镜像名称，默认为 "agent-sandbox:latest"。
               可以使用其他镜像如 "python:3.12-slim", "node:18-slim" 等。
        mem_limit: 容器最大内存限制，默认为 "512m"。
        nano_cpus: 容器最大 CPU 限制（纳秒），默认为 500000000（0.5 CPU）。
        ports: 需要映射的端口列表，如 [8080, 3000]。
        use_gpu: 是否启用 GPU 支持，默认为 False。需要宿主机安装 nvidia-container-toolkit。
        gpu_count: 使用的 GPU 数量，-1 表示使用所有可用 GPU，默认为 -1。
        server_url: 服务地址，默认使用全局配置 SANDBOX_SERVER_URL。
        timeout: 请求超时时间（秒），默认 300 秒。
        context: 工具上下文

    Returns:
        ToolResult: 包含沙盒创建结果
    """
    url = f"{server_url or SANDBOX_SERVER_URL}/api/create_environment"
    payload = {
        "image": image,
        "mem_limit": mem_limit,
        "nano_cpus": nano_cpus,
        "use_gpu": use_gpu,
        "gpu_count": gpu_count
    }
    if ports:
        payload["ports"] = ports

    try:
        async with httpx.AsyncClient(timeout=timeout) as client:
            response = await client.post(url, json=payload)
            response.raise_for_status()
            data = response.json()

        sandbox_id = data.get("sandbox_id", "")
        port_mapping = data.get("port_mapping", {})
        access_urls = data.get("access_urls", [])

        output_parts = [f"沙盒 ID: {sandbox_id}"]
        if port_mapping:
            output_parts.append(f"端口映射: {json.dumps(port_mapping)}")
        if access_urls:
            output_parts.append(f"访问地址: {', '.join(access_urls)}")

        return ToolResult(
            title="沙盒环境创建成功",
            output="\n".join(output_parts),
            metadata=data
        )
    except httpx.HTTPStatusError as e:
        return ToolResult(
            title="沙盒创建失败",
            output=f"HTTP 错误: {e.response.status_code}",
            error=str(e)
        )
    except httpx.RequestError as e:
        return ToolResult(
            title="沙盒创建失败",
            output=f"网络请求失败: {str(e)}",
            error=str(e)
        )


@tool(
    display={
        "zh": {
            "name": "执行沙盒命令",
            "params": {
                "sandbox_id": "沙盒 ID",
                "command": "Shell 命令",
                "is_background": "后台执行",
                "timeout": "超时时间"
            }
        },
        "en": {
            "name": "Run Shell in Sandbox",
            "params": {
                "sandbox_id": "Sandbox ID",
                "command": "Shell command",
                "is_background": "Run in background",
                "timeout": "Timeout"
            }
        }
    }
)
async def sandbox_run_shell(
    sandbox_id: str,
    command: str,
    is_background: bool = False,
    timeout: int = 120,
    server_url: str = None,
    request_timeout: float = DEFAULT_TIMEOUT,
    context: Optional[ToolContext] = None,
) -> ToolResult:
    """
    在指定的沙盒中执行 Shell 命令。

    Args:
        sandbox_id: 沙盒 ID，由 create_environment 返回。
        command: 要执行的 Shell 命令，如 "pip install flask" 或 "python app.py"。
        is_background: 是否后台执行，默认为 False。
            - False：前台执行，等待命令完成并返回输出
            - True：后台执行，适合启动长期运行的服务
        timeout: 前台命令的超时时间（秒），默认 120 秒。后台命令不受此限制。
        server_url: 服务地址，默认使用全局配置 SANDBOX_SERVER_URL。
        request_timeout: HTTP 请求超时时间（秒），默认 300 秒。
        context: 工具上下文

    Returns:
        ToolResult: 命令执行结果
    """
    url = f"{server_url or SANDBOX_SERVER_URL}/api/run_shell"
    payload = {
        "sandbox_id": sandbox_id,
        "command": command,
        "is_background": is_background,
        "timeout": timeout
    }

    try:
        async with httpx.AsyncClient(timeout=request_timeout) as client:
            response = await client.post(url, json=payload)
            response.raise_for_status()
            data = response.json()

        if is_background:
            status = data.get("status", "")
            message = data.get("message", "")
            log_file = data.get("log_file", "")
            output = f"状态: {status}\n消息: {message}"
            if log_file:
                output += f"\n日志文件: {log_file}"
            return ToolResult(
                title=f"后台命令已启动: {command[:50]}",
                output=output,
                metadata=data
            )
        else:
            exit_code = data.get("exit_code", -1)
            stdout = data.get("stdout", "")
            stderr = data.get("stderr", "")

            output_parts = []
            if stdout:
                output_parts.append(stdout)
            if stderr:
                if output_parts:
                    output_parts.append("\n--- stderr ---")
                output_parts.append(stderr)
            if not output_parts:
                output_parts.append("(命令无输出)")

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

            return ToolResult(
                title=title,
                output="\n".join(output_parts),
                metadata=data,
                error=None if success else f"Command failed with exit code {exit_code}"
            )
    except httpx.HTTPStatusError as e:
        return ToolResult(
            title="命令执行失败",
            output=f"HTTP 错误: {e.response.status_code}",
            error=str(e)
        )
    except httpx.RequestError as e:
        return ToolResult(
            title="命令执行失败",
            output=f"网络请求失败: {str(e)}",
            error=str(e)
        )


@tool(
    display={
        "zh": {
            "name": "重建沙盒端口",
            "params": {
                "sandbox_id": "沙盒 ID",
                "ports": "端口列表",
                "mem_limit": "内存限制",
                "nano_cpus": "CPU 限制",
                "use_gpu": "启用 GPU",
                "gpu_count": "GPU 数量"
            }
        },
        "en": {
            "name": "Rebuild Sandbox Ports",
            "params": {
                "sandbox_id": "Sandbox ID",
                "ports": "Port list",
                "mem_limit": "Memory limit",
                "nano_cpus": "CPU limit",
                "use_gpu": "Enable GPU",
                "gpu_count": "GPU count"
            }
        }
    }
)
async def sandbox_rebuild_with_ports(
    sandbox_id: str,
    ports: List[int],
    mem_limit: str = "1g",
    nano_cpus: int = 1000000000,
    use_gpu: bool = False,
    gpu_count: int = -1,
    server_url: str = None,
    timeout: float = DEFAULT_TIMEOUT,
    context: Optional[ToolContext] = None,
) -> ToolResult:
    """
    重建沙盒并应用新的端口映射。

    使用场景：先创建沙盒克隆项目，阅读 README 后才知道需要暴露哪些端口，
    此时调用此函数重建沙盒，应用正确的端口映射。

    注意：重建后会返回新的 sandbox_id，后续操作需要使用新 ID。
    容器内的所有文件（克隆的代码、安装的依赖等）都会保留。

    Args:
        sandbox_id: 当前沙盒 ID。
        ports: 需要映射的端口列表，如 [8080, 3306, 6379]。
        mem_limit: 容器最大内存限制，默认为 "1g"。
        nano_cpus: 容器最大 CPU 限制（纳秒），默认为 1000000000（1 CPU）。
        use_gpu: 是否启用 GPU 支持，默认为 False。需要宿主机安装 nvidia-container-toolkit。
        gpu_count: 使用的 GPU 数量，-1 表示使用所有可用 GPU，默认为 -1。
        server_url: 服务地址，默认使用全局配置 SANDBOX_SERVER_URL。
        timeout: 请求超时时间（秒），默认 300 秒。
        context: 工具上下文

    Returns:
        ToolResult: 重建结果
    """
    url = f"{server_url or SANDBOX_SERVER_URL}/api/rebuild_with_ports"
    payload = {
        "sandbox_id": sandbox_id,
        "ports": ports,
        "mem_limit": mem_limit,
        "nano_cpus": nano_cpus,
        "use_gpu": use_gpu,
        "gpu_count": gpu_count
    }

    try:
        async with httpx.AsyncClient(timeout=timeout) as client:
            response = await client.post(url, json=payload)
            response.raise_for_status()
            data = response.json()

        old_id = data.get("old_sandbox_id", "")
        new_id = data.get("new_sandbox_id", "")
        port_mapping = data.get("port_mapping", {})
        access_urls = data.get("access_urls", [])

        output_parts = [
            f"旧沙盒 ID: {old_id} (已销毁)",
            f"新沙盒 ID: {new_id}"
        ]
        if port_mapping:
            output_parts.append(f"端口映射: {json.dumps(port_mapping)}")
        if access_urls:
            output_parts.append(f"访问地址: {', '.join(access_urls)}")

        return ToolResult(
            title="沙盒重建成功",
            output="\n".join(output_parts),
            metadata=data
        )
    except httpx.HTTPStatusError as e:
        return ToolResult(
            title="沙盒重建失败",
            output=f"HTTP 错误: {e.response.status_code}",
            error=str(e)
        )
    except httpx.RequestError as e:
        return ToolResult(
            title="沙盒重建失败",
            output=f"网络请求失败: {str(e)}",
            error=str(e)
        )


@tool(
    requires_confirmation=True,
    display={
        "zh": {
            "name": "销毁沙盒环境",
            "params": {
                "sandbox_id": "沙盒 ID"
            }
        },
        "en": {
            "name": "Destroy Sandbox",
            "params": {
                "sandbox_id": "Sandbox ID"
            }
        }
    }
)
async def sandbox_destroy_environment(
    sandbox_id: str,
    server_url: str = None,
    timeout: float = DEFAULT_TIMEOUT,
    context: Optional[ToolContext] = None,
) -> ToolResult:
    """
    销毁沙盒环境，释放资源。

    Args:
        sandbox_id: 沙盒 ID。
        server_url: 服务地址，默认使用全局配置 SANDBOX_SERVER_URL。
        timeout: 请求超时时间（秒），默认 300 秒。
        context: 工具上下文

    Returns:
        ToolResult: 销毁结果
    """
    url = f"{server_url or SANDBOX_SERVER_URL}/api/destroy_environment"
    payload = {
        "sandbox_id": sandbox_id
    }

    try:
        async with httpx.AsyncClient(timeout=timeout) as client:
            response = await client.post(url, json=payload)
            response.raise_for_status()
            data = response.json()

        status = data.get("status", "")
        message = data.get("message", "")
        removed_tools = data.get("removed_tools", [])

        output_parts = [f"状态: {status}", f"消息: {message}"]
        if removed_tools:
            output_parts.append(f"已移除的工具: {', '.join(removed_tools)}")

        return ToolResult(
            title="沙盒环境已销毁",
            output="\n".join(output_parts),
            metadata=data
        )
    except httpx.HTTPStatusError as e:
        return ToolResult(
            title="沙盒销毁失败",
            output=f"HTTP 错误: {e.response.status_code}",
            error=str(e)
        )
    except httpx.RequestError as e:
        return ToolResult(
            title="沙盒销毁失败",
            output=f"网络请求失败: {str(e)}",
            error=str(e)
        )