__client_async_template.py

"""Sandbox客户端模板 / Sandbox Client Template

此模板用于生成沙箱客户端代码。
This template is used to generate sandbox client code.
"""

import time
from typing import Any, Dict, List, Optional, TYPE_CHECKING

from alibabacloud_agentrun20250910.models import (
    CreateTemplateInput,
    ListSandboxesRequest,
    ListTemplatesRequest,
    UpdateTemplateInput,
)

from agentrun.sandbox.api import SandboxControlAPI, SandboxDataAPI
from agentrun.sandbox.model import (
    ListSandboxesInput,
    ListSandboxesOutput,
    NASConfig,
    OSSMountConfig,
    PageableInput,
    PolarFsConfig,
    TemplateInput,
)
from agentrun.utils.config import Config
from agentrun.utils.exception import (
    AgentRunError,
    ClientError,
    ResourceNotExistError,
)

from .sandbox import Sandbox

if TYPE_CHECKING:
    from agentrun.sandbox.template import Template


class SandboxClient:
    """Sandbox 客户端 / Sandbox Client

    用于管理 Sandbox 和 Template。
    Used for managing Sandboxes and Templates.
    """

    def __init__(self, config: Optional[Config] = None):
        """初始化 Sandbox 客户端 / Initialize Sandbox client

        Args:
            config: 配置对象,可选 / Configuration object, optional
        """
        self.__control_api = SandboxControlAPI(config=config)
        self.__sandbox_data_api = SandboxDataAPI(config=config)

    async def _wait_template_ready_async(
        self,
        template_name: str,
        config: Optional[Config] = None,
        interval_seconds: int = 5,
        timeout_seconds: int = 300,
    ) -> "Template":
        """Wait for Template to be ready (async)

        Args:
            template_name: Template name
            config: Config object
            interval_seconds: Polling interval in seconds
            timeout_seconds: Timeout in seconds

        Returns:
            Template: Ready Template object

        Raises:
            TimeoutError: Timeout error
            ClientError: Client error
        """
        import asyncio

        start_time = time.time()
        while True:
            template = await self.get_template_async(
                template_name, config=config
            )

            # Check if ready
            if template.status == "READY":
                return template

            # Check if failed
            if (
                template.status == "CREATE_FAILED"
                or template.status == "UPDATE_FAILED"
            ):
                raise AgentRunError(
                    f"Template {template_name} creation failed, status:"
                    f" {template.status}"
                )

            # Check timeout
            if time.time() - start_time > timeout_seconds:
                raise TimeoutError(
                    f"Timeout waiting for Template {template_name} to be ready,"
                    f" current status: {template.status}"
                )

            await asyncio.sleep(interval_seconds)

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

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

        Returns:
            Template: 创建的 Template 对象

        Raises:
            ClientError: 客户端错误
            ServerError: 服务器错误
        """
        from agentrun.sandbox.template import Template

        # 转换为 SDK 需要的格式
        sdk_input = CreateTemplateInput().from_map(
            input.model_dump(by_alias=True)
        )
        result = await self.__control_api.create_template_async(
            sdk_input, config=config
        )
        template = Template.from_inner_object(result)

        # Poll and wait for Template to be ready
        template = await self._wait_template_ready_async(
            template.template_name or "", config=config
        )

        return template

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

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

        Returns:
            Template: 删除的 Template 对象

        Raises:
            ResourceNotExistError: Template 不存在
            ClientError: 客户端错误
            ServerError: 服务器错误
        """
        from agentrun.sandbox.template import Template

        try:
            result = await self.__control_api.delete_template_async(
                template_name, config=config
            )
            return Template.from_inner_object(result)
        except ClientError as e:
            if e.status_code == 404:
                raise ResourceNotExistError("Template", template_name) from e
            raise e

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

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

        Returns:
            Template: 更新后的 Template 对象

        Raises:
            ResourceNotExistError: Template 不存在
            ClientError: 客户端错误
            ServerError: 服务器错误
        """
        from agentrun.sandbox.template import Template

        try:
            # 转换为 SDK 需要的格式
            sdk_input = UpdateTemplateInput().from_map(
                input.model_dump(by_alias=True, exclude_none=True)
            )
            result = await self.__control_api.update_template_async(
                template_name, sdk_input, config=config
            )
            return Template.from_inner_object(result)
        except ClientError as e:
            if e.status_code == 404:
                raise ResourceNotExistError("Template", template_name) from e
            raise e

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

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

        Returns:
            Template: Template 对象

        Raises:
            ResourceNotExistError: Template 不存在
            ClientError: 客户端错误
            ServerError: 服务器错误
        """
        from agentrun.sandbox.template import Template

        try:
            result = await self.__control_api.get_template_async(
                template_name, config=config
            )
            return Template.from_inner_object(result)
        except ClientError as e:
            if e.status_code == 404:
                raise ResourceNotExistError("Template", template_name) from e
            raise e

    async def list_templates_async(
        self,
        input: Optional[PageableInput] = None,
        config: Optional[Config] = None,
    ) -> List["Template"]:
        """枚举 Templates（异步）

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

        Returns:
            List[Template]: Template 列表

        Raises:
            ClientError: 客户端错误
            ServerError: 服务器错误
            TimeoutError: Timeout waiting for Template to be ready
        """
        from agentrun.sandbox.template import Template

        if input is None:
            input = PageableInput()

        # 转换为 SDK 需要的格式
        sdk_input = ListTemplatesRequest().from_map(
            input.model_dump(by_alias=True)
        )
        results = await self.__control_api.list_templates_async(
            sdk_input, config=config
        )
        return (
            [Template.from_inner_object(item) for item in results.items]
            if results.items
            else []
        )

    async def create_sandbox_async(
        self,
        template_name: str,
        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,
    ) -> Sandbox:
        """创建 Sandbox（异步） / Create Sandbox (async)

        Args:
            template_name: 模板名称 / Template name
            sandbox_idle_timeout_seconds: 沙箱空闲超时时间（秒） / Sandbox idle timeout (seconds)
            nas_config: NAS 配置 / NAS configuration
            oss_mount_config: OSS 挂载配置 / OSS mount configuration
            polar_fs_config: PolarFS 配置 / PolarFS configuration
            config: 配置对象 / Config object

        Returns:
            Sandbox: 创建的 Sandbox 对象 / Created Sandbox object

        Raises:
            ClientError: 客户端错误 / Client error
            ServerError: 服务器错误 / Server error
        """
        # 将配置对象转换为字典格式
        nas_config_dict: Optional[Dict[str, Any]] = None
        if nas_config is not None:
            nas_config_dict = nas_config.model_dump(by_alias=True)

        oss_mount_config_dict: Optional[Dict[str, Any]] = None
        if oss_mount_config is not None:
            oss_mount_config_dict = oss_mount_config.model_dump(by_alias=True)

        polar_fs_config_dict: Optional[Dict[str, Any]] = None
        if polar_fs_config is not None:
            polar_fs_config_dict = polar_fs_config.model_dump(by_alias=True)

        result = await self.__sandbox_data_api.create_sandbox_async(
            template_name=template_name,
            sandbox_idle_timeout_seconds=sandbox_idle_timeout_seconds,
            nas_config=nas_config_dict,
            oss_mount_config=oss_mount_config_dict,
            polar_fs_config=polar_fs_config_dict,
            config=config,
        )

        # 判断返回结果是否成功
        if result.get("code") != "SUCCESS":
            raise ClientError(
                status_code=0,
                message=(
                    "Failed to create sandbox:"
                    f" {result.get('message', 'Unknown error')}"
                ),
            )

        # 从 data 字段中提取数据并实例化（使用 model_validate 从字典创建）
        data = result.get("data", {})
        return Sandbox.model_validate(data, by_alias=True)

    async def stop_sandbox_async(
        self, sandbox_id: str, config: Optional[Config] = None
    ) -> Sandbox:
        """停止 Sandbox（异步）

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

        Returns:
            Sandbox: 停止后的 Sandbox 对象

        Raises:
            ResourceNotExistError: Sandbox 不存在
            ClientError: 客户端错误
            ServerError: 服务器错误
        """
        try:
            result = await self.__sandbox_data_api.stop_sandbox_async(
                sandbox_id
            )

            # 判断返回结果是否成功
            if result.get("code") != "SUCCESS":
                raise ClientError(
                    status_code=0,
                    message=(
                        "Failed to stop sandbox:"
                        f" {result.get('message', 'Unknown error')}"
                    ),
                )

            # 从 data 字段中提取数据并实例化（使用 model_validate 从字典创建）
            data = result.get("data", {})
            return Sandbox.model_validate(data, by_alias=True)
        except ClientError as e:
            if e.status_code == 404:
                raise ResourceNotExistError("Sandbox", sandbox_id) from e
            raise e

    async def delete_sandbox_async(
        self, sandbox_id: str, config: Optional[Config] = None
    ) -> Sandbox:
        """删除 Sandbox（异步）

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

        Returns:
            Sandbox: 停止后的 Sandbox 对象

        Raises:
            ResourceNotExistError: Sandbox 不存在
            ClientError: 客户端错误
            ServerError: 服务器错误
        """
        try:
            result = await self.__sandbox_data_api.delete_sandbox_async(
                sandbox_id
            )

            # 判断返回结果是否成功
            if result.get("code") != "SUCCESS":
                raise ClientError(
                    status_code=0,
                    message=(
                        "Failed to stop sandbox:"
                        f" {result.get('message', 'Unknown error')}"
                    ),
                )

            # 从 data 字段中提取数据并实例化（使用 model_validate 从字典创建）
            data = result.get("data", {})
            return Sandbox.model_validate(data, by_alias=True)
        except ClientError as e:
            if e.status_code == 404:
                raise ResourceNotExistError("Sandbox", sandbox_id) from e
            raise e

    async def get_sandbox_async(
        self,
        sandbox_id: str,
        config: Optional[Config] = None,
    ) -> Sandbox:
        """获取 Sandbox（异步）

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

        Returns:
            Sandbox: Sandbox 对象

        Raises:
            ResourceNotExistError: Sandbox 不存在
            ClientError: 客户端错误
            ServerError: 服务器错误
        """
        try:
            result = await self.__sandbox_data_api.get_sandbox_async(sandbox_id)

            # 判断返回结果是否成功
            if result.get("code") != "SUCCESS":
                raise ClientError(
                    status_code=0,
                    message=(
                        "Failed to get sandbox:"
                        f" {result.get('message', 'Unknown error')}"
                    ),
                )

            # 从 data 字段中提取数据并实例化（使用 model_validate 从字典创建）
            data = result.get("data", {})
            return Sandbox.model_validate(data, by_alias=True)
        except ClientError as e:
            if e.status_code == 404:
                raise ResourceNotExistError("Sandbox", sandbox_id) from e
            raise e

    async def list_sandboxes_async(
        self,
        input: Optional[ListSandboxesInput] = None,
        config: Optional[Config] = None,
    ) -> ListSandboxesOutput:
        """枚举 Sandboxes（异步）

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

        Returns:
            List[Sandbox]: Sandbox 列表

        Raises:
            ClientError: 客户端错误
            ServerError: 服务器错误
        """
        if input is None:
            input = ListSandboxesInput()

        # 转换为 SDK 需要的格式
        sdk_input = ListSandboxesRequest().from_map(
            input.model_dump(by_alias=True)
        )

        results = await self.__control_api.list_sandboxes_async(
            sdk_input, config=config
        )
        return ListSandboxesOutput(
            sandboxes=[
                Sandbox.from_inner_object(item) for item in results.items
            ],
            next_token=results.next_token,
        )