Agent/agent/tools/builtin/sandbox.py

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

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

from agent import tool
from agent.tools.models import ToolResult


# 服务地址，可根据实际部署情况修改
# 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
) -> Dict[str, Any]:
    """
    创建一个隔离的 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 秒。

    Returns:
        dict: 包含以下字段：
            - sandbox_id (str): 沙盒唯一标识，后续操作需要用到
            - message (str): 提示信息
            - port_mapping (dict): 端口映射关系，如 {8080: 32001}
            - access_urls (list): 访问 URL 列表

    Raises:
        httpx.HTTPStatusError: HTTP 请求返回错误状态码时抛出
        httpx.RequestError: 网络请求失败时抛出
    """
    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

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


@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
) -> Dict[str, Any]:
    """
    在指定的沙盒中执行 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 秒。

    Returns:
        dict: 根据执行方式返回不同内容：
            前台执行：
                - exit_code (int): 命令退出码
                - stdout (str): 标准输出
                - stderr (str): 标准错误
            后台执行：
                - status (str): 状态
                - message (str): 提示信息
                - log_file (str): 日志文件路径

    Raises:
        httpx.HTTPStatusError: HTTP 请求返回错误状态码时抛出
        httpx.RequestError: 网络请求失败时抛出
    """
    url = f"{server_url or SANDBOX_SERVER_URL}/api/run_shell"
    payload = {
        "sandbox_id": sandbox_id,
        "command": command,
        "is_background": is_background,
        "timeout": timeout
    }

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


@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
) -> Dict[str, Any]:
    """
    重建沙盒并应用新的端口映射。

    使用场景：先创建沙盒克隆项目，阅读 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 秒。

    Returns:
        dict: 包含以下字段：
            - old_sandbox_id (str): 旧沙盒 ID（已销毁）
            - new_sandbox_id (str): 新沙盒 ID（后续使用这个）
            - port_mapping (dict): 端口映射关系
            - access_urls (list): 访问 URL 列表

    Raises:
        httpx.HTTPStatusError: HTTP 请求返回错误状态码时抛出
        httpx.RequestError: 网络请求失败时抛出
    """
    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
    }

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


@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
) -> Dict[str, Any]:
    """
    销毁沙盒环境，释放资源。

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

    Returns:
        dict: 包含以下字段：
            - status (str): 操作状态，如 "success"
            - message (str): 提示信息
            - removed_tools (list): 被移除的工具列表（如有关联的已注册工具）

    Raises:
        httpx.HTTPStatusError: HTTP 请求返回错误状态码时抛出
        httpx.RequestError: 网络请求失败时抛出
    """
    url = f"{server_url or SANDBOX_SERVER_URL}/api/destroy_environment"
    payload = {
        "sandbox_id": sandbox_id
    }

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