from __future__ import annotations import asyncio import functools import random from typing import TYPE_CHECKING, Any, cast import aiohttp import anyio import boxlite from shipyard.filesystem import FileSystemComponent as ShipyardFileSystemComponent from shipyard.python import PythonComponent as ShipyardPythonComponent from shipyard.shell import ShellComponent as ShipyardShellComponent from astrbot.api import logger if TYPE_CHECKING: from astrbot.core.agent.tool import ToolSchema from astrbot.core.computer.olayer import ( FileSystemComponent, PythonComponent, ShellComponent, ) from .base import ComputerBooter from .shipyard import ShipyardFileSystemWrapper, ShipyardShellWrapper class MockShipyardSandboxClient: def __init__(self, sb_url: str) -> None: self.sb_url = sb_url.rstrip("/") async def _exec_operation( self, ship_id: str, operation_type: str, payload: dict[str, Any], session_id: str, ) -> dict[str, Any]: async with aiohttp.ClientSession() as session: headers = {"X-SESSION-ID": session_id} async with session.post( f"{self.sb_url}/{operation_type}", json=payload, headers=headers, ) as response: if response.status == 200: return await response.json() error_text = await response.text() raise Exception( f"Failed to exec operation: {response.status} {error_text}", ) async def upload_file(self, path: str, remote_path: str) -> dict: """Upload a file to the sandbox""" url = f"http://{self.sb_url}/upload" try: # Read file content async with await anyio.open_file(path, "rb") as f: file_content = await f.read() # Create multipart form data data = aiohttp.FormData() data.add_field( "file", file_content, filename=remote_path.rsplit("/", maxsplit=1)[-1], content_type="application/octet-stream", ) data.add_field("file_path", remote_path) timeout = aiohttp.ClientTimeout(total=120) # 2 minutes for file upload async with aiohttp.ClientSession(timeout=timeout) as session: async with session.post(url, data=data) as response: if response.status == 200: logger.info( "[Computer] file_upload booter=boxlite remote_path=%s", remote_path, ) return { "success": True, "message": "File uploaded successfully", "file_path": remote_path, } error_text = await response.text() logger.warning( "[Computer] file_upload_failed booter=boxlite error=http_status status=%s remote_path=%s", response.status, remote_path, ) return { "success": False, "error": f"Server returned {response.status}: {error_text}", "message": "File upload failed", } except aiohttp.ClientError as e: logger.error("[Computer] file_upload_failed booter=boxlite error=%s", e) return { "success": False, "error": f"Connection error: {e!s}", "message": "File upload failed", } except TimeoutError: logger.warning( "[Computer] file_upload_failed booter=boxlite error=timeout remote_path=%s", remote_path, ) return { "success": False, "error": "File upload timeout", "message": "File upload failed", } except FileNotFoundError: logger.error( "[Computer] file_upload_failed booter=boxlite error=file_not_found path=%s", path, ) return { "success": False, "error": f"File not found: {path}", "message": "File upload failed", } except Exception as exc: logger.exception( "[Computer] file_upload_failed booter=boxlite error=unexpected", ) return { "success": False, "error": f"Internal error: {exc!s}", "message": "File upload failed", } async def wait_healthy(self, ship_id: str, session_id: str) -> None: """Mock wait healthy""" loop = 60 while loop > 0: try: logger.debug( "[Computer] health_check booter=boxlite ship_id=%s session=%s endpoint=%s attempt=%s healthy=pending", ship_id, session_id, self.sb_url, 61 - loop, ) url = f"{self.sb_url}/health" async with aiohttp.ClientSession() as session: async with session.get(url) as response: if response.status == 200: logger.debug( "[Computer] health_check booter=boxlite ship_id=%s session=%s endpoint=%s healthy=true", ship_id, session_id, self.sb_url, ) return await asyncio.sleep(1) loop -= 1 except Exception: await asyncio.sleep(1) loop -= 1 logger.warning( "[Computer] health_check_timeout booter=boxlite ship_id=%s session=%s endpoint=%s", ship_id, session_id, self.sb_url, ) class BoxliteBooter(ComputerBooter): async def boot(self, session_id: str) -> None: logger.info( "[Computer] booter_boot booter=boxlite session=%s status=starting", session_id, ) random_port = random.randint(20000, 30000) SimpleBox = vars(boxlite)["SimpleBox"] self.box = SimpleBox( image="soulter/shipyard-ship", memory_mib=512, cpus=1, ports=[ { "host_port": random_port, "guest_port": 8123, }, ], ) await self.box.start() logger.info( "[Computer] booter_boot booter=boxlite session=%s status=ready ship_id=%s", session_id, self.box.id, ) self.mocked = MockShipyardSandboxClient( sb_url=f"http://127.0.0.1:{random_port}", ) raw_fs = ShipyardFileSystemComponent( client=cast("Any", self.mocked), ship_id=self.box.id, session_id=session_id, ) self._python = ShipyardPythonComponent( client=cast("Any", self.mocked), ship_id=self.box.id, session_id=session_id, ) raw_shell = ShipyardShellComponent( client=cast("Any", self.mocked), ship_id=self.box.id, session_id=session_id, ) self._shell = ShipyardShellWrapper(cast("Any", raw_shell)) self._fs = ShipyardFileSystemWrapper(cast("Any", raw_fs), self._shell) await self.mocked.wait_healthy(self.box.id, session_id) async def shutdown(self, **kwargs) -> None: logger.info( "[Computer] booter_shutdown booter=boxlite ship_id=%s status=starting", self.box.id, ) self.box.shutdown() logger.info( "[Computer] booter_shutdown booter=boxlite ship_id=%s status=done", self.box.id, ) async def available(self) -> bool: return hasattr(self, "box") @property def fs(self) -> FileSystemComponent: return self._fs @property def python(self) -> PythonComponent: return self._python @property def shell(self) -> ShellComponent: return self._shell async def upload_file(self, path: str, file_name: str) -> dict: """Upload file to sandbox""" return await self.mocked.upload_file(path, file_name) @classmethod @functools.cache def _default_tools(cls) -> tuple[ToolSchema, ...]: from astrbot.core.computer.tools import ( ExecuteShellTool, FileDownloadTool, FileUploadTool, PythonTool, ) return ( ExecuteShellTool(), PythonTool(), FileUploadTool(), FileDownloadTool(), ) @classmethod def get_default_tools(cls) -> list[ToolSchema]: return list(cls._default_tools())