refactor(sandbox): center provider-owned metadata

Author: zouyonghe

astrbot/core/computer/computer_client.py

@@ -10,6 +10,7 @@
 
 from astrbot.api import logger
 from astrbot.core.computer.cua_registry import CuaSandboxRegistry
+from astrbot.core.computer.cua_sandbox_provider import CuaSandboxProvider
 from astrbot.core.computer.sandbox_manager import SANDBOX_LEASE_SECONDS, SandboxManager
 from astrbot.core.skills.skill_manager import SANDBOX_SKILLS_ROOT, SkillManager
 from astrbot.core.star.context import Context
@@ -38,29 +39,7 @@ class _CUAIdleState:
 cua_idle_state: dict[str, _CUAIdleState] = {}
 
 
-class _ComputerClientCuaProvider:
-    provider_id = "cua"
-
-    def build_create_config(self, context: Context, session_id: str) -> dict:
-        from .booters.cua import build_cua_booter_kwargs
-
-        config = context.get_config(umo=session_id)
-        sandbox_cfg = config.get("provider_settings", {}).get("sandbox", {})
-        return build_cua_booter_kwargs(sandbox_cfg)
-
-    async def create_booter(
-        self, context: Context, session_id: str, sandbox_id: str, config: dict
-    ) -> ComputerBooter:
-        return await _boot_managed_cua_sandbox(context, session_id, sandbox_id, config)
-
-    async def destroy_booter(self, booter: ComputerBooter, record: dict) -> None:
-        await booter.shutdown()
-
-
-sandbox_manager = SandboxManager(
-    registry=cua_registry,
-    providers={"cua": _ComputerClientCuaProvider()},
-)
+sandbox_manager: SandboxManager
 
 
 def _sync_sandbox_manager_refs() -> None:
@@ -70,9 +49,6 @@ def _sync_sandbox_manager_refs() -> None:
     sandbox_manager.boot_locks = _cua_boot_locks
 
 
-_sync_sandbox_manager_refs()
-
-
 def _save_cua_registry() -> None:
     _sync_sandbox_manager_refs()
     sandbox_manager.save_registry()
@@ -184,6 +160,22 @@ async def _boot_managed_cua_sandbox(
     return client
 
 
+async def _boot_managed_cua_sandbox_hook(
+    context: Context,
+    session_id: str,
+    sandbox_id: str,
+    cua_kwargs: dict,
+) -> ComputerBooter:
+    return await _boot_managed_cua_sandbox(context, session_id, sandbox_id, cua_kwargs)
+
+
+sandbox_manager = SandboxManager(
+    registry=cua_registry,
+    providers={"cua": CuaSandboxProvider(boot_hook=_boot_managed_cua_sandbox_hook)},
+)
+_sync_sandbox_manager_refs()
+
+
 def _new_managed_cua_sandbox_id() -> str:
     return f"cua-{uuid.uuid4().hex[:12]}"
 

astrbot/core/computer/cua_sandbox_provider.py

@@ -2,12 +2,15 @@
 
 import time
 import uuid
+from collections.abc import Awaitable, Callable
 
 from astrbot.api import logger
 from astrbot.core.computer.booters import cua as cua_booter
 from astrbot.core.computer.booters.base import ComputerBooter
 from astrbot.core.star.context import Context
 
+BootHook = Callable[[Context, str, str, dict], Awaitable[ComputerBooter]]
+
 
 async def _sync_skills_to_sandbox(booter: ComputerBooter) -> None:
     from astrbot.core.computer.computer_client import _sync_skills_to_sandbox as sync
@@ -18,18 +21,41 @@ async def _sync_skills_to_sandbox(booter: ComputerBooter) -> None:
 class CuaSandboxProvider:
     provider_id = "cua"
 
+    def __init__(self, boot_hook: BootHook | None = None) -> None:
+        self._boot_hook = boot_hook
+
     def build_create_config(self, context: Context, session_id: str) -> dict:
         config = context.get_config(umo=session_id)
         sandbox_cfg = config.get("provider_settings", {}).get("sandbox", {})
         return cua_booter.build_cua_booter_kwargs(sandbox_cfg)
 
+    def build_connect_info(self, sandbox_name: str, config: dict) -> dict:
+        return {
+            "name": sandbox_name,
+            "local": config.get("local", True),
+            "image": config.get("image"),
+            "os_type": config.get("os_type"),
+        }
+
+    def get_idle_timeout(self, context: Context, session_id: str) -> float:
+        config = context.get_config(umo=session_id)
+        sandbox_cfg = config.get("provider_settings", {}).get("sandbox", {})
+        value = sandbox_cfg.get("cua_idle_timeout", 0)
+        try:
+            timeout = float(value)
+        except (TypeError, ValueError):
+            return 0.0
+        return max(timeout, 0.0)
+
     async def create_booter(
         self,
         context: Context,
         session_id: str,
         sandbox_id: str,
         config: dict,
     ) -> ComputerBooter:
+        if self._boot_hook is not None:
+            return await self._boot_hook(context, session_id, sandbox_id, config)
         uuid_str = uuid.uuid5(uuid.NAMESPACE_DNS, session_id).hex
         client = cua_booter.CuaBooter(**config)
         started_at = time.monotonic()

astrbot/core/computer/sandbox_manager.py

@@ -50,13 +50,8 @@ def drop_boot_lock(self, sandbox_id: str) -> None:
         self.boot_locks.pop(sandbox_id, None)
 
     def get_idle_timeout(self, config: dict, provider_id: str) -> float:
-        sandbox_cfg = config.get("provider_settings", {}).get("sandbox", {})
-        value = sandbox_cfg.get(f"{provider_id}_idle_timeout", 0)
-        try:
-            timeout = float(value)
-        except (TypeError, ValueError):
-            return 0.0
-        return max(timeout, 0.0)
+        _ = config, provider_id
+        return 0.0
 
     def build_record_payload(
         self,
@@ -67,8 +62,10 @@ def build_record_payload(
         provider_id: str,
         config: dict,
         idle_timeout: float,
+        connect_info: dict,
         is_default: bool = False,
     ) -> dict:
+        _ = config
         return {
             "sandbox_id": sandbox_id,
             "sandbox_name": sandbox_name,
@@ -78,12 +75,7 @@ def build_record_payload(
             "created_by_astrbot": True,
             "owner_user_id": session_id,
             "owner_session_id": session_id,
-            "connect_info": {
-                "name": sandbox_name,
-                "local": config.get("local", True),
-                "image": config.get("image"),
-                "os_type": config.get("os_type"),
-            },
+            "connect_info": connect_info,
             "is_default": is_default,
             "idle_timeout": idle_timeout,
         }
@@ -153,6 +145,7 @@ async def get_or_create_booter(
     ) -> ComputerBooter:
         provider = self.get_provider(provider_id)
         create_config = provider.build_create_config(context, session_id)
+        idle_timeout = provider.get_idle_timeout(context, session_id)
 
         current_sandbox_id = self.registry.get_current_sandbox_id(session_id)
         current_record = (
@@ -186,9 +179,7 @@ async def get_or_create_booter(
                 self.save_registry()
                 self.schedule_idle_cleanup(
                     current_sandbox_id,
-                    self.get_idle_timeout(
-                        context.get_config(umo=session_id), provider_id
-                    ),
+                    idle_timeout,
                 )
                 return booter
             self.session_booter.pop(current_sandbox_id, None)
@@ -205,8 +196,9 @@ async def get_or_create_booter(
                     session_id=session_id,
                     provider_id=provider_id,
                     config=create_config,
-                    idle_timeout=self.get_idle_timeout(
-                        context.get_config(umo=session_id), provider_id
+                    idle_timeout=idle_timeout,
+                    connect_info=provider.build_connect_info(
+                        target_sandbox_id, create_config
                     ),
                     is_default=True,
                 )
@@ -285,7 +277,7 @@ async def get_or_create_booter(
         self.save_registry()
         self.schedule_idle_cleanup(
             target_sandbox_id,
-            self.get_idle_timeout(context.get_config(umo=session_id), provider_id),
+            idle_timeout,
         )
         return self.session_booter[target_sandbox_id]
 
@@ -298,10 +290,9 @@ async def create_sandbox_uncontrolled(
     ) -> dict:
         provider = self.get_provider(provider_id)
         create_config = provider.build_create_config(context, session_id)
-        config = context.get_config(umo=session_id)
         sandbox_id = self.new_sandbox_id(provider_id)
         sandbox_name = sandbox_name or sandbox_id
-        idle_timeout = self.get_idle_timeout(config, provider_id)
+        idle_timeout = provider.get_idle_timeout(context, session_id)
         record = self.registry.upsert_sandbox(
             **self.build_record_payload(
                 sandbox_id=sandbox_id,
@@ -310,6 +301,7 @@ async def create_sandbox_uncontrolled(
                 provider_id=provider_id,
                 config=create_config,
                 idle_timeout=idle_timeout,
+                connect_info=provider.build_connect_info(sandbox_name, create_config),
             )
         )
         try:
@@ -391,6 +383,7 @@ def update_sandbox_config(
     def _upsert_new_sandbox_record(
         self, context: Context, session_id: str, provider_id: str, create_config: dict
     ) -> str:
+        provider = self.get_provider(provider_id)
         sandbox_id = self.new_sandbox_id(provider_id)
         self.registry.upsert_sandbox(
             **self.build_record_payload(
@@ -399,9 +392,8 @@ def _upsert_new_sandbox_record(
                 session_id=session_id,
                 provider_id=provider_id,
                 config=create_config,
-                idle_timeout=self.get_idle_timeout(
-                    context.get_config(umo=session_id), provider_id
-                ),
+                idle_timeout=provider.get_idle_timeout(context, session_id),
+                connect_info=provider.build_connect_info(sandbox_id, create_config),
             )
         )
         self.save_registry()

astrbot/core/computer/sandbox_provider.py

@@ -11,6 +11,10 @@ class SandboxProvider(Protocol):
 
     def build_create_config(self, context: Context, session_id: str) -> dict: ...
 
+    def build_connect_info(self, sandbox_name: str, config: dict) -> dict: ...
+
+    def get_idle_timeout(self, context: Context, session_id: str) -> float: ...
+
     async def create_booter(
         self,
         context: Context,

astrbot/core/computer/sandbox_registry.py

@@ -7,6 +7,7 @@
 from typing import Any
 
 from astrbot.api import logger
+from astrbot.core.computer.sandbox_models import SandboxRecord
 from astrbot.core.utils.astrbot_path import get_astrbot_data_path
 
 _UNSET = object()
@@ -104,6 +105,7 @@ def upsert_sandbox(
                 "notes": notes,
             }
         )
+        record = SandboxRecord.from_dict(record).to_dict()
         self._payload["sandboxes"][sandbox_id] = record
         if is_default or (managed and self._payload["default_sandbox_id"] is None):
             self.set_default_sandbox_id(sandbox_id)
@@ -250,9 +252,24 @@ def load(self) -> None:
         self._payload = _default_registry_payload()
         self._payload.update(payload)
         sandboxes = self._payload.get("sandboxes", {})
+        valid_sandboxes = {}
         for record in sandboxes.values():
+            try:
+                normalized = SandboxRecord.from_dict(record).to_dict()
+            except (KeyError, TypeError, ValueError) as exc:
+                logger.warning(
+                    "[Computer] Skip invalid sandbox registry record: %s",
+                    exc,
+                )
+                continue
+            record = normalized
             if not record.get("managed"):
+                valid_sandboxes[record["sandbox_id"]] = record
                 continue
             record["controller_user_id"] = None
             record["controller_session_id"] = None
             record["lease_expires_at"] = None
+            valid_sandboxes[record["sandbox_id"]] = record
+        self._payload["sandboxes"] = valid_sandboxes
+        if self._payload["default_sandbox_id"] not in valid_sandboxes:
+            self._payload["default_sandbox_id"] = None

tests/unit/test_sandbox_manager.py

@@ -25,6 +25,12 @@ def __init__(self):
     def build_create_config(self, context, session_id):
         return {"image": "fake"}
 
+    def build_connect_info(self, sandbox_name, config):
+        return {"kind": "fake", "name": sandbox_name}
+
+    def get_idle_timeout(self, context, session_id):
+        return 7.0
+
     async def create_booter(self, context, session_id, sandbox_id, config):
         self.boots.append((session_id, sandbox_id, config))
         return FakeBooter(sandbox_id, self)
@@ -99,6 +105,11 @@ async def test_sandbox_manager_creates_default_and_current_sandbox(tmp_path):
 
     assert booter.sandbox_id == registry.default_sandbox_id
     assert registry.get_current_sandbox_id("session-a") == booter.sandbox_id
+    assert registry.get_sandbox(booter.sandbox_id)["connect_info"] == {
+        "kind": "fake",
+        "name": booter.sandbox_id,
+    }
+    assert registry.get_sandbox(booter.sandbox_id)["idle_timeout"] == 7.0
     assert provider.boots == [("session-a", booter.sandbox_id, {"image": "fake"})]