zhenxun_bot/zhenxun/services/ai/sandbox/models.py

"""
沙箱相关核心类型定义
"""

from abc import ABC, abstractmethod
import asyncio
import hashlib
import json
from pathlib import Path
import re
from typing import TYPE_CHECKING, Annotated, Any, Literal

from pydantic import BaseModel, ConfigDict, Field

from zhenxun.services.ai.utils.logger import log_sandbox as logger
from zhenxun.utils.pydantic_compat import model_dump

if TYPE_CHECKING:
    from .drivers.base import BaseSandboxSession


class LanguageProfile(BaseModel):
    """语言执行配置模板"""

    language: str
    """语言名称"""
    aliases: list[str] = Field(default_factory=list)
    """语言别名列表"""
    source_ext: str
    """源文件后缀名"""
    compile_cmd: str | None = None
    """编译命令"""
    run_cmd: str
    """运行命令"""
    deps_install_cmd: str | None = None
    """依赖安装命令"""


class BaseEntry(BaseModel, ABC):
    """沙箱物化项(Artifacts)抽象基类"""

    @abstractmethod
    async def apply(self, session: "BaseSandboxSession", dest: str) -> None:
        """将当前项物化(应用)到沙箱中的指定目标路径"""
        pass


class MemoryFile(BaseEntry):
    """基于内存字符串或二进制的数据物化"""

    type: Literal["memory_file"] = "memory_file"
    content: str | bytes

    async def apply(self, session: "BaseSandboxSession", dest: str) -> None:
        data = (
            self.content.encode("utf-8")
            if isinstance(self.content, str)
            else self.content
        )
        await session.write(dest, data)


class LocalFile(BaseEntry):
    """基于宿主机本地单文件的物化"""

    type: Literal["local_file"] = "local_file"
    src: str

    async def apply(self, session: "BaseSandboxSession", dest: str) -> None:
        local_path = Path(self.src)
        if await asyncio.to_thread(local_path.is_file):
            content = await asyncio.to_thread(local_path.read_bytes)
            await session.write(dest, content)
        else:
            logger.warning(f"[Sandbox] 本地文件 {self.src} 不存在，已跳过物化。")


class LocalDir(BaseEntry):
    """基于宿主机本地目录的物化 (会自动打包上传)"""

    type: Literal["local_dir"] = "local_dir"
    src: str

    async def apply(self, session: "BaseSandboxSession", dest: str) -> None:
        success = await session.upload_raw_dir(self.src, dest)
        if not success:
            logger.warning(f"[Sandbox] 本地目录 {self.src} 上传失败，已跳过物化。")


EntryUnion = Annotated[MemoryFile | LocalFile | LocalDir, Field(discriminator="type")]


class BaseSetupStep(BaseModel, ABC):
    """环境装配步骤基类，多态支持任何第三方语言的安装"""

    model_config = ConfigDict(arbitrary_types_allowed=True)

    @abstractmethod
    async def apply(self, session: Any) -> None:
        """在沙箱会话中执行该装配步骤"""
        pass


class AptSetup(BaseSetupStep):
    type: Literal["apt"] = "apt"
    packages: list[str]

    async def apply(self, session: Any) -> None:
        if not self.packages:
            return

        res = await session.run_process("dpkg -l", timeout=10)
        installed_pkgs = set()
        if res.exit_code == 0:
            for line in res.stdout.splitlines():
                if line.startswith("ii "):
                    parts = line.split()
                    if len(parts) >= 2:
                        pkg_name = parts[1].split(":")[0].strip().lower()
                        installed_pkgs.add(pkg_name)

        missing_packages = [
            pkg for pkg in self.packages if pkg.strip().lower() not in installed_pkgs
        ]
        if not missing_packages:
            logger.debug(f"[Sandbox] Apt 依赖已全部满足，跳过安装: {self.packages}")
            return

        pkg_str = " ".join(missing_packages)
        logger.info(f"[Sandbox] 正在安装缺失的系统级依赖 (Apt): {pkg_str}")
        await session.run_process(
            "sudo DEBIAN_FRONTEND=noninteractive apt-get update -qq && "
            f"sudo DEBIAN_FRONTEND=noninteractive apt-get install -y -qq {pkg_str}",
            timeout=300,
        )


class PythonSetup(BaseSetupStep):
    type: Literal["python"] = "python"
    packages: list[str]

    async def apply(self, session: Any) -> None:
        if not self.packages:
            return

        res = await session.run_process("python3 -m pip freeze", timeout=10)
        installed_pkgs = set()
        if res.exit_code == 0:
            for line in res.stdout.splitlines():
                if "==" not in line and "@" not in line:
                    continue
                base_name = re.split(r"[=><~!@\[]", line)[0].strip().lower()
                installed_pkgs.add(base_name)

        missing_packages = [
            pkg
            for pkg in self.packages
            if re.split(r"[=><~!@\[]", pkg)[0].strip().lower() not in installed_pkgs
        ]
        if not missing_packages:
            logger.debug(f"[Sandbox] Python 依赖已全部满足，跳过安装: {self.packages}")
            return

        pkg_str = " ".join(missing_packages)
        check_uv = await session.run_process("command -v uv")
        if check_uv.exit_code != 0:
            logger.info("[Sandbox] 未检测到 uv，正在极速下载 uv...")
            await session.run_process(
                "pip install uv --disable-pip-version-check -q", timeout=60
            )
        logger.info(f"[Sandbox] 正在安装缺失的 Python 依赖: {pkg_str}")
        res = await session.run_process(f"uv pip install {pkg_str}", timeout=120)
        if res.exit_code != 0:
            logger.warning(f"[Sandbox] uv 安装报错，尝试降级使用 pip: {res.stderr}")
            await session.run_process(f"pip install {pkg_str} -q", timeout=180)


class NodeSetup(BaseSetupStep):
    type: Literal["node"] = "node"
    packages: list[str]

    async def apply(self, session: Any) -> None:
        if not self.packages:
            return

        res = await session.run_process("npm list -g --depth=0", timeout=15)
        installed_pkgs = set()
        if res.exit_code == 0:
            for line in res.stdout.splitlines():
                match = re.search(r"[├└]──\s+((?:@[^/]+/)?[^@\s]+)", line)
                if match:
                    installed_pkgs.add(match.group(1).strip().lower())

        missing_packages = []
        for pkg in self.packages:
            base_name = (
                ("@" + pkg[1:].split("@")[0])
                if pkg.startswith("@")
                else pkg.split("@")[0]
            )
            if base_name.strip().lower() not in installed_pkgs:
                missing_packages.append(pkg)

        if not missing_packages:
            logger.debug(f"[Sandbox] Node 依赖已全部满足，跳过安装: {self.packages}")
            return

        pkg_str = " ".join(missing_packages)
        logger.info(f"[Sandbox] 正在安装缺失的 Node 依赖: {pkg_str}")
        res = await session.run_process(
            f"npm install -g {pkg_str} --no-fund --no-audit --cache /tmp/npm-cache",
            timeout=180,
        )
        if res.exit_code != 0:
            from zhenxun.services.ai.core.exceptions import SandboxFatalError

            raise SandboxFatalError(
                f"Node 依赖安装失败 (Exit Code: {res.exit_code}): "
                f"{res.stderr or res.error}"
            )


class ShellSetup(BaseSetupStep):
    type: Literal["shell"] = "shell"
    scripts: list[str]

    async def apply(self, session: Any) -> None:
        for script in self.scripts:
            logger.info(f"[Sandbox] 正在执行自定义装配脚本: {script}")
            res = await session.run_process(script, timeout=300)
            if res.exit_code != 0:
                from zhenxun.services.ai.core.exceptions import SandboxFatalError

                raise SandboxFatalError(
                    f"自定义装配脚本执行失败: {res.stderr or res.error}"
                )


class BindMount(BaseModel):
    """宿主机物理目录绑定映射配置"""

    host_path: str
    """宿主机绝对路径"""
    sandbox_path: str = "/workspace"
    """沙箱内目标路径，默认覆盖 /workspace"""
    read_only: bool = False
    """是否以只读模式挂载"""


class SandboxBlueprint(BaseModel):
    """沙箱大一统声明式蓝图配置"""

    sandbox_type: str = Field(default="auto")
    """强制驱动类型，如 docker 或 local"""
    enable_network: bool = Field(default=False)
    """是否允许访问外网"""
    needs_state: bool = Field(default=False)
    """是否需要持久化状态"""
    image: str | None = Field(default=None)
    """自定义镜像名称。为空则使用全局默认配置"""
    container_name: str = Field(default="global")
    """容器/命名空间名称。相同名称将共享同一个物理容器和依赖环境"""

    setup_steps: list[BaseSetupStep] = Field(default_factory=list)
    """多态环境装配图元管线"""

    entries: dict[str, EntryUnion] = Field(default_factory=dict)
    """沙箱物化节点树(Artifacts)，键为沙箱内的相对目标路径"""
    env: dict[str, str] = Field(default_factory=dict)
    """环境变量"""

    required_extensions: list[str] = Field(default_factory=list)
    """需要自动挂载的扩展"""

    bind_mounts: list[BindMount] = Field(default_factory=list)
    """宿主机目录物理挂载映射列表 (Bind Mounts)"""

    def with_sandbox_type(self, sandbox_type: str) -> "SandboxBlueprint":
        """设置沙箱驱动类型"""
        self.sandbox_type = sandbox_type
        return self

    def with_network(self, enable: bool = True) -> "SandboxBlueprint":
        """设置是否启用网络访问"""
        self.enable_network = enable
        return self

    def with_state(self, enable: bool = True) -> "SandboxBlueprint":
        """设置是否启用状态持久化"""
        self.needs_state = enable
        return self

    def with_image(self, image: str) -> "SandboxBlueprint":
        """声明使用的自定义 Docker 镜像"""
        self.image = image
        return self

    def with_container_name(self, name: str) -> "SandboxBlueprint":
        """声明沙箱使用的物理容器命名空间（实现环境隔离）"""
        self.container_name = name
        return self

    def with_setup_step(self, step: BaseSetupStep) -> "SandboxBlueprint":
        """声明一个多态环境装配步骤"""
        self.setup_steps.append(step)
        return self

    def with_python_packages(self, packages: list[str]) -> "SandboxBlueprint":
        """追加预置 Python 依赖包"""
        self.setup_steps.append(PythonSetup(packages=packages))
        return self

    def with_system_packages(self, packages: list[str]) -> "SandboxBlueprint":
        """追加预置系统包依赖"""
        self.setup_steps.append(AptSetup(packages=packages))
        return self

    def with_node_packages(self, packages: list[str]) -> "SandboxBlueprint":
        """追加预置 Node 依赖包"""
        self.setup_steps.append(NodeSetup(packages=packages))
        return self

    def with_install_scripts(self, scripts: list[str]) -> "SandboxBlueprint":
        """追加自定义安装脚本"""
        self.setup_steps.append(ShellSetup(scripts=scripts))
        return self

    def with_file(self, path: str, content: str | bytes) -> "SandboxBlueprint":
        """声明预置内存字符串或二进制文件"""
        self.entries[path] = MemoryFile(content=content)
        return self

    def with_local_file(self, path: str, local_path: str) -> "SandboxBlueprint":
        """声明预置本地宿主机单文件"""
        self.entries[path] = LocalFile(src=local_path)
        return self

    def with_local_dir(self, path: str, local_dir: str) -> "SandboxBlueprint":
        """声明预置本地宿主机完整目录"""

        for bm in self.bind_mounts:
            if path.startswith(bm.sandbox_path) or bm.sandbox_path.startswith(path):
                logger.warning(
                    f"⚠️ [Sandbox] 路径 '{path}' 和物理挂载映射 "
                    f"'{bm.sandbox_path}' 存在重叠！\n"
                    "继续使用 with_local_dir 打包上传可能会引发冗余上传"
                    "或双向覆盖冲突。建议直接使用物理挂载。"
                )
        self.entries[path] = LocalDir(src=local_dir)
        return self

    def with_env(self, key: str, value: str) -> "SandboxBlueprint":
        """设置环境变量"""
        self.env[key] = value
        return self

    def with_extension(self, extension: str) -> "SandboxBlueprint":
        """声明需要挂载的扩展"""
        if extension not in self.required_extensions:
            self.required_extensions.append(extension)
        return self

    def with_workspace(
        self,
        local_path: str | Path,
        remote_path: str = "/workspace",
        read_only: bool = False,
    ) -> "SandboxBlueprint":
        """声明宿主机物理目录双向挂载 (Bind Mount)"""
        from pathlib import Path

        if remote_path in self.entries:
            logger.warning(
                f"⚠️ [Sandbox] 目标沙箱路径 '{remote_path}' "
                "已包含普通实体映射(如 LocalDir/File)。\n"
                "强制绑定物理挂载 (Bind Mount) 将会遮蔽原有实体。"
            )

        abs_path = Path(local_path).resolve().as_posix()
        self.bind_mounts.append(
            BindMount(host_path=abs_path, sandbox_path=remote_path, read_only=read_only)
        )
        return self

    def calculate_hash(self) -> str:
        """计算当前环境配置的 MD5 指纹，用于缓存命中"""
        entries_dict = {
            k: model_dump(self.entries[k]) for k in sorted(self.entries.keys())
        }

        data_to_hash = {
            "image": self.image,
            "container_name": self.container_name,
            "steps": [model_dump(s) for s in self.setup_steps],
            "entries": entries_dict,
            "bind_mounts": [model_dump(m) for m in self.bind_mounts],
        }
        json_str = json.dumps(data_to_hash, separators=(",", ":"))
        return hashlib.md5(json_str.encode("utf-8")).hexdigest()

    def merge(self, other: "SandboxBlueprint") -> "SandboxBlueprint":
        """合并另一个蓝图配置"""
        if not other:
            return self
        self.enable_network = self.enable_network or other.enable_network
        self.needs_state = self.needs_state or other.needs_state
        self.sandbox_type = (
            other.sandbox_type if other.sandbox_type != "auto" else self.sandbox_type
        )
        self.image = other.image or self.image
        self.container_name = (
            other.container_name
            if other.container_name != "global"
            else self.container_name
        )

        self.setup_steps.extend(other.setup_steps)
        self.entries.update(other.entries)
        self.env.update(other.env)
        self.required_extensions = list(
            dict.fromkeys(self.required_extensions + other.required_extensions)
        )
        self.bind_mounts.extend(other.bind_mounts)
        return self


class SandboxExecutionResult(BaseModel):
    """沙箱统一的执行结果数据结构"""

    stdout: str = Field(default="")
    """标准输出流"""
    stderr: str = Field(default="")
    """标准错误流"""
    exit_code: int = Field(default=0)
    """进程退出码"""
    error: str | None = Field(default=None)
    """框架执行错误信息"""
    is_timeout: bool = Field(default=False)
    """进程是否因超时被挂起(仍在运行)"""
    images: list[str] = Field(default_factory=list)
    """Base64 图片列表"""
    artifacts: dict[str, bytes] = Field(default_factory=dict)
    """生成的文件工件"""

    @property
    def is_success(self) -> bool:
        return self.exit_code == 0 and self.error is None


class SandboxSessionState(BaseModel):
    """沙箱会话的序列化状态，用于无状态恢复"""

    session_id: str
    """会话唯一标识符"""
    backend_id: str
    """后端驱动分配的容器/沙箱 ID"""
    container_name: str = "global"
    """底层所关联的物理容器名称"""
    sandbox_type: str = "docker"
    """沙箱驱动类型"""
    workspace_root_ready: bool = False
    """工作区目录是否已初始化完成"""
    extra_data: dict[str, Any] = Field(default_factory=dict)
    """会话的其它序列化元数据"""