__sandbox_async_template.py

"""Sandbox 高层 API / Sandbox High-Level API

此模块定义沙箱资源的高级API。
This module defines the high-level API for sandbox resources.
"""

from typing import (
    Any,
    Dict,
    List,
    Literal,
    Optional,
    overload,
    TYPE_CHECKING,
    Union,
)

from pydantic import Field

from agentrun.sandbox.model import TemplateType
from agentrun.utils.config import Config
from agentrun.utils.model import BaseModel

if TYPE_CHECKING:
    from agentrun.sandbox.aio_sandbox import AioSandbox
    from agentrun.sandbox.browser_sandbox import BrowserSandbox
    from agentrun.sandbox.code_interpreter_sandbox import CodeInterpreterSandbox
    from agentrun.sandbox.custom_sandbox import CustomSandbox
    from agentrun.sandbox.model import (
        ListSandboxesInput,
        ListSandboxesOutput,
        NASConfig,
        OSSMountConfig,
        PageableInput,
        PolarFsConfig,
        TemplateInput,
    )
    from agentrun.sandbox.template import Template


class Sandbox(BaseModel):
    """Sandbox 实例

    封装了 Sandbox 的基本信息和操作方法
    """

    _template_type: Optional[TemplateType]

    created_at: Optional[str] = None
    """沙箱创建时间 / Sandbox creation time"""
    ended_at: Optional[str] = None
    """沙箱结束时间 / Sandbox end time"""
    last_updated_at: Optional[str] = None
    """最后更新时间 / Last updated time"""
    metadata: Optional[Dict[str, Any]] = None
    """元数据 / Metadata"""
    sandbox_arn: Optional[str] = None
    """沙箱全局唯一资源名称 / Sandbox ARN"""
    sandbox_id: Optional[str] = None
    """沙箱 ID / Sandbox ID"""
    sandbox_idle_ttlin_seconds: Optional[int] = Field(
        None, alias="sandboxIdleTTLInSeconds"
    )
    """沙箱空闲 TTL（秒） / Sandbox Idle TTL (seconds)"""
    sandbox_idle_timeout_seconds: Optional[int] = None
    """沙箱空闲超时时间（秒） / Sandbox Idle Timeout (seconds)"""
    status: Optional[str] = None
    """沙箱状态 / Sandbox status"""
    template_id: Optional[str] = None
    """模板 ID / Template ID"""
    template_name: Optional[str] = None
    """模板名称 / Template name"""
    _config: Optional[Config] = None
    """配置对象，用于子类的 data_api 初始化 / Config object for data_api initialization"""

    @classmethod
    def __get_client(cls):
        """获取 Sandbox 客户端"""
        from .client import SandboxClient

        return SandboxClient()

    @classmethod
    @overload
    async def create_async(
        cls,
        template_type: Literal[TemplateType.CODE_INTERPRETER],
        template_name: Optional[str] = None,
        sandbox_idle_timeout_seconds: Optional[int] = 600,
        nas_config: Optional["NASConfig"] = None,
        oss_mount_config: Optional["OSSMountConfig"] = None,
        polar_fs_config: Optional["PolarFsConfig"] = None,
        config: Optional[Config] = None,
    ) -> "CodeInterpreterSandbox":
        ...

    @classmethod
    @overload
    async def create_async(
        cls,
        template_type: Literal[TemplateType.BROWSER],
        template_name: Optional[str] = None,
        sandbox_idle_timeout_seconds: Optional[int] = 600,
        nas_config: Optional["NASConfig"] = None,
        oss_mount_config: Optional["OSSMountConfig"] = None,
        polar_fs_config: Optional["PolarFsConfig"] = None,
        config: Optional[Config] = None,
    ) -> "BrowserSandbox":
        ...

    @classmethod
    @overload
    async def create_async(
        cls,
        template_type: Literal[TemplateType.AIO],
        template_name: Optional[str] = None,
        sandbox_idle_timeout_seconds: Optional[int] = 600,
        nas_config: Optional["NASConfig"] = None,
        oss_mount_config: Optional["OSSMountConfig"] = None,
        polar_fs_config: Optional["PolarFsConfig"] = None,
        config: Optional[Config] = None,
    ) -> "AioSandbox":
        ...

    @classmethod
    @overload
    async def create_async(
        cls,
        template_type: Literal[TemplateType.CUSTOM],
        template_name: Optional[str] = None,
        sandbox_idle_timeout_seconds: Optional[int] = 600,
        nas_config: Optional["NASConfig"] = None,
        oss_mount_config: Optional["OSSMountConfig"] = None,
        polar_fs_config: Optional["PolarFsConfig"] = None,
        config: Optional[Config] = None,
    ) -> "CustomSandbox":
        ...

    @classmethod
    async def create_async(
        cls,
        template_type: TemplateType,
        template_name: Optional[str] = None,
        sandbox_idle_timeout_seconds: Optional[int] = 600,
        nas_config: Optional["NASConfig"] = None,
        oss_mount_config: Optional["OSSMountConfig"] = None,
        polar_fs_config: Optional["PolarFsConfig"] = None,
        config: Optional[Config] = None,
    ) -> Union[
        "CodeInterpreterSandbox",
        "BrowserSandbox",
        "AioSandbox",
        "CustomSandbox",
    ]:

        if template_name is None:
            # todo 可以考虑为用户创建一个模板？
            raise ValueError("template_name is required")

        # 先根据传入的template_name，获取template的类型
        template = await cls.get_template_async(template_name, config=config)

        # 根据 template 类型创建相应的 Sandbox 子类
        from agentrun.sandbox.aio_sandbox import AioSandbox
        from agentrun.sandbox.browser_sandbox import BrowserSandbox
        from agentrun.sandbox.code_interpreter_sandbox import (
            CodeInterpreterSandbox,
        )
        from agentrun.sandbox.custom_sandbox import CustomSandbox

        if template_type != template.template_type:
            raise ValueError(
                f"template_type of {template_name} is {template.template_type},"
                f" not {template_type}"
            )

        # 创建 Sandbox（返回基类实例）
        base_sandbox = await cls.__get_client().create_sandbox_async(
            template_name=template_name,
            sandbox_idle_timeout_seconds=sandbox_idle_timeout_seconds,
            nas_config=nas_config,
            oss_mount_config=oss_mount_config,
            polar_fs_config=polar_fs_config,
        )

        # 根据 template 类型转换为对应的子类实例
        sandbox = None
        if template.template_type == TemplateType.CODE_INTERPRETER:
            sandbox = CodeInterpreterSandbox.model_validate(
                base_sandbox.model_dump(by_alias=False)
            )
        elif template.template_type == TemplateType.BROWSER:
            sandbox = BrowserSandbox.model_validate(
                base_sandbox.model_dump(by_alias=False)
            )
        elif template.template_type == TemplateType.AIO:
            sandbox = AioSandbox.model_validate(
                base_sandbox.model_dump(by_alias=False)
            )
        elif template.template_type == TemplateType.CUSTOM:
            sandbox = CustomSandbox.model_validate(
                base_sandbox.model_dump(by_alias=False)
            )
        else:
            raise ValueError(
                f"template_type {template.template_type} is not supported"
            )

        sandbox._config = config
        return sandbox

    @classmethod
    async def stop_by_id_async(cls, sandbox_id: str):
        """通过 ID 停止 Sandbox（异步）

        Args:
            sandbox_id: Sandbox ID
            config: 配置对象

        Returns:
            Sandbox: Sandbox 对象
        """
        if sandbox_id is None:
            raise ValueError("sandbox_id is required")
        # todo 后续适配后使用 stop()
        return await cls.__get_client().stop_sandbox_async(sandbox_id)

    @classmethod
    async def delete_by_id_async(cls, sandbox_id: str):
        """通过 ID 删除 Sandbox（异步）

        Args:
            sandbox_id: Sandbox ID
            config: 配置对象

        Returns:
            Sandbox: Sandbox 对象
        """
        if sandbox_id is None:
            raise ValueError("sandbox_id is required")
        return await cls.__get_client().delete_sandbox_async(sandbox_id)

    @classmethod
    async def list_async(
        cls,
        input: Optional["ListSandboxesInput"] = None,
        config: Optional[Config] = None,
    ) -> "ListSandboxesOutput":
        """列出 Sandboxes（异步）

        Args:
            input: 列表查询配置
            config: 配置对象

        Returns:
            ListSandboxesOutput: Sandbox 列表结果
        """
        return await cls.__get_client().list_sandboxes_async(input, config)

    @classmethod
    @overload
    async def connect_async(
        cls,
        sandbox_id: str,
        template_type: Literal[TemplateType.CODE_INTERPRETER],
        config: Optional[Config] = None,
    ) -> "CodeInterpreterSandbox":
        ...

    @classmethod
    @overload
    async def connect_async(
        cls,
        sandbox_id: str,
        template_type: Literal[TemplateType.BROWSER],
        config: Optional[Config] = None,
    ) -> "BrowserSandbox":
        ...

    @classmethod
    @overload
    async def connect_async(
        cls,
        sandbox_id: str,
        template_type: Literal[TemplateType.AIO],
        config: Optional[Config] = None,
    ) -> "AioSandbox":
        ...

    @classmethod
    @overload
    async def connect_async(
        cls,
        sandbox_id: str,
        template_type: None = None,
        config: Optional[Config] = None,
    ) -> Union["CodeInterpreterSandbox", "BrowserSandbox", "AioSandbox"]:
        ...

    @classmethod
    async def connect_async(
        cls,
        sandbox_id: str,
        template_type: Optional[TemplateType] = None,
        config: Optional[Config] = None,
    ) -> Union["CodeInterpreterSandbox", "BrowserSandbox", "AioSandbox"]:
        """连接一个SandBox（异步）

        Args:
            sandbox_id: Sandbox ID
            type: 可选的类型参数，用于类型提示和运行时验证
            config: 配置对象

        Returns:
            Sandbox: 根据模板类型返回对应的 Sandbox 子类对象

        Raises:
            ValueError: 如果模板类型不支持或与预期类型不匹配
        """
        if sandbox_id is None:
            raise ValueError("sandbox_id is required")

        # 先获取 sandbox 信息
        sandbox = await cls.__get_client().get_sandbox_async(
            sandbox_id, config=config
        )

        # 根据 template_name 获取 template 类型
        if sandbox.template_name is None:
            raise ValueError(f"Sandbox {sandbox_id} has no template_name")

        template = await cls.get_template_async(
            sandbox.template_name, config=config
        )

        # 如果提供了 type 参数，验证类型是否匹配
        if (
            template_type is not None
            and template.template_type != template_type
        ):
            raise ValueError(
                f"Sandbox {sandbox_id} has template type"
                f" {template.template_type}, but expected {template_type}"
            )

        # 根据 template 类型创建相应的 Sandbox 子类
        from agentrun.sandbox.aio_sandbox import AioSandbox
        from agentrun.sandbox.browser_sandbox import BrowserSandbox
        from agentrun.sandbox.code_interpreter_sandbox import (
            CodeInterpreterSandbox,
        )

        result = None
        if template.template_type == TemplateType.CODE_INTERPRETER:
            result = CodeInterpreterSandbox.model_validate(
                sandbox.model_dump(by_alias=False)
            )
        elif template.template_type == TemplateType.BROWSER:
            result = BrowserSandbox.model_validate(
                sandbox.model_dump(by_alias=False)
            )
        elif template.template_type == TemplateType.AIO:
            result = AioSandbox.model_validate(
                sandbox.model_dump(by_alias=False)
            )
        else:
            raise ValueError(
                f"Unsupported template type: {template.template_type}. "
                "Expected 'code-interpreter', 'browser' or 'aio'"
            )

        result._config = config
        return result

    # ==================== Template 相关类方法 ====================

    @classmethod
    async def create_template_async(
        cls, input: "TemplateInput", config: Optional[Config] = None
    ) -> "Template":
        """创建 Template（异步）

        Args:
            input: Template 配置
            config: 配置对象

        Returns:
            Template: 创建的 Template 对象
        """
        if input.template_type is None:
            raise ValueError("template_type is required")
        return await cls.__get_client().create_template_async(
            input, config=config
        )

    @classmethod
    async def get_template_async(
        cls, template_name: str, config: Optional[Config] = None
    ) -> "Template":
        """获取 Template（异步）

        Args:
            template_name: Template 名称
            config: 配置对象

        Returns:
            Template: Template 对象
        """
        if template_name is None:
            raise ValueError("template_name is required")
        return await cls.__get_client().get_template_async(
            template_name, config=config
        )

    @classmethod
    async def update_template_async(
        cls,
        template_name: str,
        input: "TemplateInput",
        config: Optional[Config] = None,
    ) -> "Template":
        """更新 Template（异步）

        Args:
            template_name: Template 名称
            input: Template 更新配置
            config: 配置对象

        Returns:
            Template: 更新后的 Template 对象
        """
        if template_name is None:
            raise ValueError("template_name is required")
        return await cls.__get_client().update_template_async(
            template_name, input, config=config
        )

    @classmethod
    async def delete_template_async(
        cls, template_name: str, config: Optional[Config] = None
    ) -> "Template":
        """删除 Template（异步）

        Args:
            template_name: Template 名称
            config: 配置对象

        Returns:
            Template: 删除的 Template 对象
        """
        if template_name is None:
            raise ValueError("template_name is required")
        return await cls.__get_client().delete_template_async(
            template_name, config=config
        )

    @classmethod
    async def list_templates_async(
        cls,
        input: Optional["PageableInput"] = None,
        config: Optional[Config] = None,
    ) -> List["Template"]:
        """列出 Templates（异步）

        Args:
            input: 分页配置
            config: 配置对象

        Returns:
            List[Template]: Template 列表
        """
        return await cls.__get_client().list_templates_async(
            input, config=config
        )

    async def get_async(self):
        if self.sandbox_id is None:
            raise ValueError("sandbox_id is required to get a Sandbox")

        return await self.connect_async(self.sandbox_id)

    async def delete_async(self):
        if self.sandbox_id is None:
            raise ValueError("sandbox_id is required to delete a Sandbox")

        return await self.delete_by_id_async(self.sandbox_id)

    async def stop_async(self):
        if self.sandbox_id is None:
            raise ValueError("sandbox_id is required to stop a Sandbox")
        # todo 后续适配后使用 stop()
        return await self.stop_by_id_async(self.sandbox_id)