mirror of
https://github.com/AstrBotDevs/AstrBot
synced 2026-07-16 17:47:06 +08:00
309 lines
9.0 KiB
Python
309 lines
9.0 KiB
Python
from __future__ import annotations
|
|
|
|
import functools
|
|
import shlex
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
from shipyard import ShipyardClient, Spec
|
|
|
|
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 .shell_background import build_detached_shell_command
|
|
from .shipyard_search_file_util import search_files_via_shell
|
|
|
|
|
|
def _maybe_model_dump(value: Any) -> dict[str, Any]:
|
|
if isinstance(value, dict):
|
|
return value
|
|
if hasattr(value, "model_dump"):
|
|
dumped = value.model_dump()
|
|
if isinstance(dumped, dict):
|
|
return dumped
|
|
return {}
|
|
|
|
|
|
class ShipyardShellWrapper:
|
|
def __init__(self, _shipyard_shell: ShellComponent):
|
|
self._shell = _shipyard_shell
|
|
|
|
async def exec(
|
|
self,
|
|
command: str,
|
|
cwd: str | None = None,
|
|
env: dict[str, str] | None = None,
|
|
timeout: int | None = 300,
|
|
shell: bool = True,
|
|
background: bool = False,
|
|
session_id: str | None = None,
|
|
) -> dict[str, Any]:
|
|
_ = session_id
|
|
if not shell:
|
|
return {
|
|
"stdout": "",
|
|
"stderr": "error: only shell mode is supported in shipyard booter.",
|
|
"exit_code": 2,
|
|
"success": False,
|
|
}
|
|
|
|
run_command = command
|
|
if env:
|
|
env_prefix = " ".join(
|
|
f"{k}={shlex.quote(str(v))}" for k, v in sorted(env.items())
|
|
)
|
|
run_command = f"{env_prefix} {run_command}"
|
|
|
|
if background:
|
|
run_command = build_detached_shell_command(run_command)
|
|
|
|
result = await self._shell.exec(
|
|
run_command,
|
|
timeout=timeout or 300,
|
|
cwd=cwd,
|
|
)
|
|
payload = _maybe_model_dump(result)
|
|
|
|
stdout = payload.get("output", payload.get("stdout", "")) or ""
|
|
stderr = payload.get("error", payload.get("stderr", "")) or ""
|
|
exit_code = payload.get("exit_code")
|
|
if background:
|
|
pid: int | None = None
|
|
try:
|
|
pid = int(str(stdout).strip().splitlines()[-1])
|
|
except Exception:
|
|
pid = None
|
|
return {
|
|
"pid": pid,
|
|
"stdout": (
|
|
f"Command is running in the background. pid={pid}"
|
|
if pid is not None
|
|
else "Command was submitted in the background."
|
|
),
|
|
"stderr": stderr,
|
|
"exit_code": exit_code,
|
|
"success": bool(payload.get("success", not stderr)),
|
|
"execution_id": payload.get("execution_id"),
|
|
"execution_time_ms": payload.get("execution_time_ms"),
|
|
"command": payload.get("command"),
|
|
}
|
|
|
|
return {
|
|
"stdout": stdout,
|
|
"stderr": stderr,
|
|
"exit_code": exit_code,
|
|
"success": bool(payload.get("success", not stderr)),
|
|
"execution_id": payload.get("execution_id"),
|
|
"execution_time_ms": payload.get("execution_time_ms"),
|
|
"command": payload.get("command"),
|
|
}
|
|
|
|
|
|
class ShipyardFileSystemWrapper:
|
|
def __init__(
|
|
self,
|
|
_shipyard_fs: FileSystemComponent,
|
|
_shipyard_shell: ShellComponent,
|
|
):
|
|
self._fs = _shipyard_fs
|
|
self._shell = _shipyard_shell
|
|
|
|
async def create_file(
|
|
self,
|
|
path: str,
|
|
content: str = "",
|
|
mode: int = 420,
|
|
) -> dict[str, Any]:
|
|
return await self._fs.create_file(path=path, content=content, mode=mode)
|
|
|
|
async def read_file(
|
|
self,
|
|
path: str,
|
|
encoding: str = "utf-8",
|
|
offset: int | None = None,
|
|
limit: int | None = None,
|
|
) -> dict[str, Any]:
|
|
return await self._fs.read_file(
|
|
path=path,
|
|
encoding=encoding,
|
|
offset=offset,
|
|
limit=limit,
|
|
)
|
|
|
|
async def write_file(
|
|
self,
|
|
path: str,
|
|
content: str,
|
|
mode: str = "w",
|
|
encoding: str = "utf-8",
|
|
) -> dict[str, Any]:
|
|
return await self._fs.write_file(
|
|
path=path,
|
|
content=content,
|
|
mode=mode,
|
|
encoding=encoding,
|
|
)
|
|
|
|
async def list_dir(
|
|
self,
|
|
path: str = ".",
|
|
show_hidden: bool = False,
|
|
) -> dict[str, Any]:
|
|
return await self._fs.list_dir(path=path, show_hidden=show_hidden)
|
|
|
|
async def delete_file(self, path: str) -> dict[str, Any]:
|
|
return await self._fs.delete_file(path=path)
|
|
|
|
async def search_files(
|
|
self,
|
|
pattern: str,
|
|
path: str | None = None,
|
|
glob: str | None = None,
|
|
after_context: int | None = None,
|
|
before_context: int | None = None,
|
|
) -> dict[str, Any]:
|
|
return await search_files_via_shell(
|
|
self._shell,
|
|
pattern=pattern,
|
|
path=path,
|
|
glob=glob,
|
|
after_context=after_context,
|
|
before_context=before_context,
|
|
)
|
|
|
|
async def edit_file(
|
|
self,
|
|
path: str,
|
|
old_string: str,
|
|
new_string: str,
|
|
replace_all: bool = False,
|
|
encoding: str = "utf-8",
|
|
) -> dict[str, Any]:
|
|
return await self._fs.edit_file(
|
|
path=path,
|
|
old_string=old_string,
|
|
new_string=new_string,
|
|
replace_all=replace_all,
|
|
encoding=encoding,
|
|
)
|
|
|
|
|
|
class ShipyardBooter(ComputerBooter):
|
|
@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())
|
|
|
|
def __init__(
|
|
self,
|
|
endpoint_url: str,
|
|
access_token: str,
|
|
ttl: int = 3600,
|
|
session_num: int = 10,
|
|
) -> None:
|
|
self._sandbox_client = ShipyardClient(
|
|
endpoint_url=endpoint_url,
|
|
access_token=access_token,
|
|
)
|
|
self._ttl = ttl
|
|
self._session_num = session_num
|
|
|
|
async def boot(self, session_id: str) -> None:
|
|
ship = await self._sandbox_client.create_ship(
|
|
ttl=self._ttl,
|
|
spec=Spec(cpus=1.0, memory="512m"),
|
|
max_session_num=self._session_num,
|
|
session_id=session_id,
|
|
)
|
|
logger.info(
|
|
"[Computer] sandbox_created booter=shipyard ship_id=%s session=%s",
|
|
ship.id,
|
|
session_id,
|
|
)
|
|
self._ship = ship
|
|
self._shell = ShipyardShellWrapper(self._ship.shell) # type: ignore[arg-type]
|
|
self._fs = ShipyardFileSystemWrapper(self._ship.fs, self._shell) # type: ignore[arg-type]
|
|
|
|
async def shutdown(self, **kwargs) -> None:
|
|
logger.info("[Computer] booter_shutdown booter=shipyard status=done")
|
|
|
|
@property
|
|
def fs(self) -> FileSystemComponent:
|
|
return self._ship.fs # type: ignore[return-value]
|
|
|
|
@property
|
|
def python(self) -> PythonComponent:
|
|
return self._ship.python
|
|
|
|
@property
|
|
def shell(self) -> ShellComponent:
|
|
return self._shell
|
|
|
|
async def upload_file(self, path: str, file_name: str) -> dict:
|
|
"""Upload file to sandbox"""
|
|
result = await self._ship.upload_file(path, file_name)
|
|
logger.info(
|
|
"[Computer] file_upload booter=shipyard remote_path=%s",
|
|
file_name,
|
|
)
|
|
return result
|
|
|
|
async def download_file(self, remote_path: str, local_path: str):
|
|
"""Download file from sandbox."""
|
|
result = await self._ship.download_file(remote_path, local_path)
|
|
logger.info(
|
|
"[Computer] file_download booter=shipyard remote_path=%s local_path=%s",
|
|
remote_path,
|
|
local_path,
|
|
)
|
|
return result
|
|
|
|
async def available(self) -> bool:
|
|
"""Check if the sandbox is available."""
|
|
try:
|
|
ship_id = self._ship.id
|
|
data = await self._sandbox_client.get_ship(ship_id)
|
|
if not data:
|
|
logger.debug(
|
|
"[Computer] health_check booter=shipyard ship_id=%s healthy=false reason=no_data",
|
|
ship_id,
|
|
)
|
|
return False
|
|
health = bool(data.get("status", 0) == 1)
|
|
logger.debug(
|
|
"[Computer] health_check booter=shipyard ship_id=%s healthy=%s",
|
|
ship_id,
|
|
health,
|
|
)
|
|
return health
|
|
except Exception:
|
|
logger.exception(
|
|
"[Computer] health_check_failed booter=shipyard ship_id=%s",
|
|
getattr(getattr(self, "_ship", None), "id", "unknown"),
|
|
)
|
|
return False
|