plugins/sandbox/mcp.py

"""
沙盒MCP插件 MCP工具定义
"""

from app.mcp.decorators import mcp_tool
from pydantic import BaseModel, Field
from typing import List

from .main import plugin
from .sandbox.security import SecurityController
from .sandbox.file_ops import FileOperations
from .sandbox.command_exec import CommandExecutor


class ReadFileInput(BaseModel):
    """文件读取工具输入参数"""
    path: str = Field(description="文件相对路径（相对于沙盒根目录）")
    encoding: str = Field(default="utf-8", description="文件编码")


class ReadFileOutput(BaseModel):
    """文件读取工具输出结果"""
    success: bool = Field(description="操作是否成功")
    content: str | None = Field(default=None, description="文件内容")
    size: int | None = Field(default=None, description="文件大小（字节）")
    error: str | None = Field(default=None, description="错误信息")


class WriteFileInput(BaseModel):
    """文件写入工具输入参数"""
    path: str = Field(description="文件相对路径（相对于沙盒根目录）")
    content: str = Field(description="文件内容")
    mode: str = Field(default="write", description="写入模式：write 或 append")


class WriteFileOutput(BaseModel):
    """文件写入工具输出结果"""
    success: bool = Field(description="操作是否成功")
    path: str | None = Field(default=None, description="写入的文件路径")
    size: int | None = Field(default=None, description="文件大小（字节）")
    error: str | None = Field(default=None, description="错误信息")


class ListFilesInput(BaseModel):
    """文件列表工具输入参数"""
    path: str = Field(default="", description="目录相对路径（空表示沙盒根目录）")
    pattern: str = Field(default="*", description="匹配模式（glob pattern）")


class ListFilesOutput(BaseModel):
    """文件列表工具输出结果"""
    success: bool = Field(description="操作是否成功")
    files: List[dict] | None = Field(default=None, description="文件列表")
    error: str | None = Field(default=None, description="错误信息")


class ExecuteInput(BaseModel):
    """命令执行工具输入参数"""
    command: str = Field(description="要执行的命令")
    timeout: int = Field(default=30, description="超时时间（秒）")
    cwd: str = Field(default="", description="工作目录（相对路径）")


class ExecuteOutput(BaseModel):
    """命令执行工具输出结果"""
    success: bool = Field(description="操作是否成功")
    stdout: str = Field(default="", description="标准输出")
    stderr: str = Field(default="", description="标准错误")
    exit_code: int = Field(default=-1, description="退出码")
    error: str | None = Field(default=None, description="错误信息")


@mcp_tool(
    name="sandbox-read",
    title="沙盒文件读取",
    description="从沙盒目录读取文件内容",
    annotations={
        "readOnlyHint": True,
        "destructiveHint": False,
    }
)
async def sandbox_read(path: str, encoding: str = "utf-8") -> ReadFileOutput:
    """从沙盒目录读取文件内容。

    Args:
        params: 包含文件路径和编码参数

    Returns:
        ReadFileOutput: 文件内容和相关信息
    """
    security = SecurityController()
    file_ops = FileOperations(security)

    result = file_ops.read_file(path, encoding)

    return ReadFileOutput(
        success=result.get("success", False),
        content=result.get("content"),
        size=result.get("size"),
        error=result.get("error"),
    )


@mcp_tool(
    name="sandbox-write",
    title="沙盒文件写入",
    description="向沙盒目录写入文件内容",
    annotations={
        "readOnlyHint": False,
        "destructiveHint": False,
    }
)
async def sandbox_write(
    content: str,
    path: str | None = None,
    filename: str | None = None,
    mode: str = "write",
) -> WriteFileOutput:
    """向沙盒目录写入文件内容。

    Args:
        params: 包含文件路径、内容和写入模式参数

    Returns:
        WriteFileOutput: 写入结果
    """
    security = SecurityController()
    file_ops = FileOperations(security)

    target_path = path or filename
    if not target_path:
        return WriteFileOutput(
            success=False,
            error="缺少文件路径，请传入 path 或 filename",
        )

    result = file_ops.write_file(target_path, content, mode)

    return WriteFileOutput(
        success=result.get("success", False),
        path=result.get("path"),
        size=result.get("size"),
        error=result.get("error"),
    )


@mcp_tool(
    name="sandbox-list",
    title="沙盒文件列表",
    description="列出沙盒目录中的文件和目录",
    annotations={
        "readOnlyHint": True,
        "destructiveHint": False,
    }
)
async def sandbox_list(path: str = "", pattern: str = "*") -> ListFilesOutput:
    """列出沙盒目录中的文件和目录。

    Args:
        params: 包含目录路径和匹配模式参数

    Returns:
        ListFilesOutput: 文件列表
    """
    security = SecurityController()
    file_ops = FileOperations(security)

    result = file_ops.list_files(path, pattern)

    return ListFilesOutput(
        success=result.get("success", False),
        files=result.get("files"),
        error=result.get("error"),
    )


@mcp_tool(
    name="sandbox-exec",
    title="沙盒命令执行",
    description="在沙盒环境中安全执行命令",
    annotations={
        "readOnlyHint": False,
        "destructiveHint": True,
    }
)
async def sandbox_exec(command: str, timeout: int = 30, cwd: str = "") -> ExecuteOutput:
    """在沙盒环境中安全执行命令。

    Args:
        params: 包含命令、超时时间和工作目录参数

    Returns:
        ExecuteOutput: 命令执行结果
    """
    security = SecurityController()
    cmd_exec = CommandExecutor(security)

    result = await cmd_exec.execute(command, timeout, cwd)

    return ExecuteOutput(
        success=result.get("success", False),
        stdout=result.get("stdout", ""),
        stderr=result.get("stderr", ""),
        exit_code=result.get("exit_code", -1),
        error=result.get("error"),
    )