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://61.48.133.26:9999"

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


@tool(
    display={
        "zh": {
            "name": "创建沙盒环境",
            "params": {
                "image": "Docker 镜像",
                "mem_limit": "内存限制",
                "nano_cpus": "CPU 限制",
                "ports": "端口列表"
            }
        },
        "en": {
            "name": "Create Sandbox",
            "params": {
                "image": "Docker image",
                "mem_limit": "Memory limit",
                "nano_cpus": "CPU limit",
                "ports": "Port list"
            }
        }
    }
)
async def sandbox_create_environment(
    image: str = "agent-sandbox:latest",
    mem_limit: str = "512m",
    nano_cpus: int = 500000000,
    ports: Optional[List[int]] = None,
    uid: str = ""
) -> 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]。
        uid: 用户ID（自动注入）

    Returns:
        ToolResult 包含：
            - sandbox_id (str): 沙盒唯一标识，后续操作需要用到
            - message (str): 提示信息
            - port_mapping (dict): 端口映射关系，如 {8080: 32001}
            - access_urls (list): 访问 URL 列表
    """
    try:
        url = f"{SANDBOX_SERVER_URL}/api/create_environment"
        payload = {
            "image": image,
            "mem_limit": mem_limit,
            "nano_cpus": nano_cpus
        }
        if ports:
            payload["ports"] = ports

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

        return ToolResult(
            title="Sandbox Created",
            output=json.dumps(data, ensure_ascii=False, indent=2),
            long_term_memory=f"Created sandbox: {data.get('sandbox_id', 'unknown')}"
        )
    except httpx.HTTPStatusError as e:
        return ToolResult(
            title="Create Sandbox Failed",
            output="",
            error=f"HTTP error {e.response.status_code}: {e.response.text}"
        )
    except Exception as e:
        return ToolResult(
            title="Create Sandbox Failed",
            output="",
            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,
    uid: str = ""
) -> ToolResult:
    """
    在指定的沙盒中执行 Shell 命令。

    Args:
        sandbox_id: 沙盒 ID，由 create_environment 返回。
        command: 要执行的 Shell 命令，如 "pip install flask" 或 "python app.py"。
        is_background: 是否后台执行，默认为 False。
            - False：前台执行，等待命令完成并返回输出
            - True：后台执行，适合启动长期运行的服务
        timeout: 前台命令的超时时间（秒），默认 120 秒。后台命令不受此限制。
        uid: 用户ID（自动注入）

    Returns:
        ToolResult 包含：
            前台执行：
                - exit_code (int): 命令退出码
                - stdout (str): 标准输出
                - stderr (str): 标准错误
            后台执行：
                - status (str): 状态
                - message (str): 提示信息
                - log_file (str): 日志文件路径
    """
    try:
        url = f"{SANDBOX_SERVER_URL}/api/run_shell"
        payload = {
            "sandbox_id": sandbox_id,
            "command": command,
            "is_background": is_background,
            "timeout": timeout
        }

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

        return ToolResult(
            title=f"Shell: {command[:50]}{'...' if len(command) > 50 else ''}",
            output=json.dumps(data, ensure_ascii=False, indent=2),
            long_term_memory=f"Executed in sandbox {sandbox_id}: {command[:100]}"
        )
    except httpx.HTTPStatusError as e:
        return ToolResult(
            title="Run Shell Failed",
            output="",
            error=f"HTTP error {e.response.status_code}: {e.response.text}"
        )
    except Exception as e:
        return ToolResult(
            title="Run Shell Failed",
            output="",
            error=str(e)
        )


@tool(
    display={
        "zh": {
            "name": "重建沙盒端口",
            "params": {
                "sandbox_id": "沙盒 ID",
                "ports": "端口列表",
                "mem_limit": "内存限制",
                "nano_cpus": "CPU 限制"
            }
        },
        "en": {
            "name": "Rebuild Sandbox Ports",
            "params": {
                "sandbox_id": "Sandbox ID",
                "ports": "Port list",
                "mem_limit": "Memory limit",
                "nano_cpus": "CPU limit"
            }
        }
    }
)
async def sandbox_rebuild_with_ports(
    sandbox_id: str,
    ports: List[int],
    mem_limit: str = "1g",
    nano_cpus: int = 1000000000,
    uid: str = ""
) -> ToolResult:
    """
    重建沙盒并应用新的端口映射。

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

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

    Args:
        sandbox_id: 当前沙盒 ID。
        ports: 需要映射的端口列表，如 [8080, 3306, 6379]。
        mem_limit: 容器最大内存限制，默认为 "1g"。
        nano_cpus: 容器最大 CPU 限制（纳秒），默认为 1000000000（1 CPU）。
        uid: 用户ID（自动注入）

    Returns:
        ToolResult 包含：
            - old_sandbox_id (str): 旧沙盒 ID（已销毁）
            - new_sandbox_id (str): 新沙盒 ID（后续使用这个）
            - port_mapping (dict): 端口映射关系
            - access_urls (list): 访问 URL 列表
    """
    try:
        url = f"{SANDBOX_SERVER_URL}/api/rebuild_with_ports"
        payload = {
            "sandbox_id": sandbox_id,
            "ports": ports,
            "mem_limit": mem_limit,
            "nano_cpus": nano_cpus
        }

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

        return ToolResult(
            title="Sandbox Rebuilt",
            output=json.dumps(data, ensure_ascii=False, indent=2),
            long_term_memory=f"Rebuilt sandbox {sandbox_id} -> {data.get('new_sandbox_id', 'unknown')} with ports {ports}"
        )
    except httpx.HTTPStatusError as e:
        return ToolResult(
            title="Rebuild Sandbox Failed",
            output="",
            error=f"HTTP error {e.response.status_code}: {e.response.text}"
        )
    except Exception as e:
        return ToolResult(
            title="Rebuild Sandbox Failed",
            output="",
            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,
    uid: str = ""
) -> ToolResult:
    """
    销毁沙盒环境，释放资源。

    Args:
        sandbox_id: 沙盒 ID。
        uid: 用户ID（自动注入）

    Returns:
        ToolResult 包含：
            - status (str): 操作状态，如 "success"
            - message (str): 提示信息
            - removed_tools (list): 被移除的工具列表（如有关联的已注册工具）
    """
    try:
        url = f"{SANDBOX_SERVER_URL}/api/destroy_environment"
        payload = {
            "sandbox_id": sandbox_id
        }

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

        return ToolResult(
            title="Sandbox Destroyed",
            output=json.dumps(data, ensure_ascii=False, indent=2),
            long_term_memory=f"Destroyed sandbox: {sandbox_id}"
        )
    except httpx.HTTPStatusError as e:
        return ToolResult(
            title="Destroy Sandbox Failed",
            output="",
            error=f"HTTP error {e.response.status_code}: {e.response.text}"
        )
    except Exception as e:
        return ToolResult(
            title="Destroy Sandbox Failed",
            output="",
            error=str(e)
        )


@tool(
    display={
        "zh": {
            "name": "注册沙盒工具",
            "params": {
                "tool_name": "工具名称",
                "description": "工具描述",
                "input_schema": "参数定义",
                "sandbox_id": "沙盒 ID",
                "internal_port": "内部端口",
                "endpoint_path": "API 路径",
                "http_method": "HTTP 方法"
            }
        },
        "en": {
            "name": "Register Sandbox Tool",
            "params": {
                "tool_name": "Tool name",
                "description": "Tool description",
                "input_schema": "Input schema",
                "sandbox_id": "Sandbox ID",
                "internal_port": "Internal port",
                "endpoint_path": "API path",
                "http_method": "HTTP method"
            }
        }
    }
)
async def sandbox_register_tool(
    tool_name: str,
    description: str,
    input_schema: Dict[str, Any],
    sandbox_id: str,
    internal_port: int,
    endpoint_path: str = "/",
    http_method: str = "POST",
    metadata: Optional[Dict[str, Any]] = None,
    uid: str = ""
) -> ToolResult:
    """
    将部署好的服务注册为工具。

    注册后，该工具会出现在统一 MCP Server 的工具列表中，可被上游服务调用。

    Args:
        tool_name: 工具唯一标识（字母开头，只能包含字母、数字、下划线），
                   如 "rendercv_api"。
        description: 工具描述，描述该工具的功能。
        input_schema: JSON Schema 格式的参数定义，定义工具接收的参数。
        sandbox_id: 服务所在的沙盒 ID。
        internal_port: 服务在容器内的端口。
        endpoint_path: API 路径，默认 "/"。
        http_method: HTTP 方法，默认 "POST"。
        metadata: 额外元数据（可选）。
        uid: 用户ID（自动注入）

    Returns:
        ToolResult 包含：
            - status (str): 操作状态，"success" 或 "error"
            - message (str): 提示信息
            - tool_info (dict): 工具信息（成功时）
    """
    try:
        url = f"{SANDBOX_SERVER_URL}/api/register_tool"
        payload = {
            "tool_name": tool_name,
            "description": description,
            "input_schema": input_schema,
            "sandbox_id": sandbox_id,
            "internal_port": internal_port,
            "endpoint_path": endpoint_path,
            "http_method": http_method
        }
        if metadata:
            payload["metadata"] = metadata

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

        return ToolResult(
            title=f"Tool Registered: {tool_name}",
            output=json.dumps(data, ensure_ascii=False, indent=2),
            long_term_memory=f"Registered tool '{tool_name}' on sandbox {sandbox_id}:{internal_port}"
        )
    except httpx.HTTPStatusError as e:
        return ToolResult(
            title="Register Tool Failed",
            output="",
            error=f"HTTP error {e.response.status_code}: {e.response.text}"
        )
    except Exception as e:
        return ToolResult(
            title="Register Tool Failed",
            output="",
            error=str(e)
        )