fix(cua): harden sandbox lifecycle controls

Author: zouyonghe

astrbot/core/computer/computer_client.py

@@ -25,6 +25,7 @@
 _MANAGED_SKILLS_FILE = ".astrbot_managed_skills.json"
 cua_registry = CuaSandboxRegistry()
 CUA_LEASE_SECONDS = 300
+_cua_boot_locks: dict[str, asyncio.Lock] = {}
 
 
 @dataclass(slots=True)
@@ -59,6 +60,14 @@ def _clear_cua_idle_state(sandbox_id: str) -> None:
         state.task.cancel()
 
 
+def _cua_boot_lock(sandbox_id: str) -> asyncio.Lock:
+    lock = _cua_boot_locks.get(sandbox_id)
+    if lock is None:
+        lock = asyncio.Lock()
+        _cua_boot_locks[sandbox_id] = lock
+    return lock
+
+
 def _schedule_cua_idle_cleanup(sandbox_id: str, timeout: float) -> None:
     _clear_cua_idle_state(sandbox_id)
     if timeout <= 0:
@@ -112,6 +121,11 @@ async def _expire_when_idle() -> None:
                         )
                     finally:
                         session_booter.pop(sandbox_id, None)
+                if record.get("retention_policy") == "persistent":
+                    cua_registry.update_sandbox_status(sandbox_id, "stopped")
+                else:
+                    cua_registry.delete_sandbox(sandbox_id)
+                _save_cua_registry()
                 return
         except asyncio.CancelledError:
             raise
@@ -316,12 +330,13 @@ async def _get_or_create_cua_booter(
         )
         _save_cua_registry()
 
-    if target_sandbox_id not in session_booter:
-        client = await _boot_managed_cua_sandbox(
-            context, session_id, target_sandbox_id, cua_kwargs
-        )
-        setattr(client, "sandbox_id", target_sandbox_id)
-        session_booter[target_sandbox_id] = client
+    async with _cua_boot_lock(target_sandbox_id):
+        if target_sandbox_id not in session_booter:
+            client = await _boot_managed_cua_sandbox(
+                context, session_id, target_sandbox_id, cua_kwargs
+            )
+            setattr(client, "sandbox_id", target_sandbox_id)
+            session_booter[target_sandbox_id] = client
 
     _acquire_cua_sandbox_lease(target_sandbox_id, session_id)
     cua_registry.touch_sandbox(target_sandbox_id)
@@ -472,6 +487,9 @@ def switch_current_cua_sandbox(session_id: str, sandbox_id: str) -> dict:
     record = cua_registry.get_sandbox(sandbox_id)
     if record is None or not record.get("managed"):
         raise RuntimeError(f"Sandbox {sandbox_id} not found")
+    booter = session_booter.get(sandbox_id)
+    if booter is None:
+        raise RuntimeError(f"Sandbox {sandbox_id} is not running")
     if not _acquire_cua_sandbox_lease(sandbox_id, session_id):
         raise RuntimeError(f"Sandbox {sandbox_id} is busy")
     cua_registry.set_current_sandbox_id(session_id, sandbox_id)

astrbot/core/computer/cua_registry.py

@@ -154,6 +154,15 @@ def update_sandbox_config(
             record["retention_policy"] = retention_policy
         return deepcopy(record)
 
+    def update_sandbox_status(
+        self, sandbox_id: str, status: str
+    ) -> dict[str, Any] | None:
+        record = self._payload["sandboxes"].get(sandbox_id)
+        if record is None:
+            return None
+        record["status"] = status
+        return deepcopy(record)
+
     def acquire_lease(
         self,
         *,

astrbot/core/core_lifecycle.py

@@ -400,6 +400,7 @@ async def stop(self) -> None:
 
     async def restart(self) -> None:
         """重启 AstrBot 核心生命周期管理类, 终止各个管理器并重新加载平台实例"""
+        await cleanup_managed_cua_sandboxes()
         await self.provider_manager.terminate()
         await self.platform_manager.terminate()
         await self.kb_manager.terminate()

astrbot/core/tools/computer_tools/cua_sandbox.py

@@ -226,6 +226,8 @@ async def call(
         context: ContextWrapper[AstrAgentContext],
         sandbox_id: str,
     ) -> ToolExecResult:
+        if err := check_admin_permission(context, "Taking over CUA sandboxes"):
+            return err
         try:
             sandbox = takeover_cua_sandbox(_session_id(context), sandbox_id)
             return _to_json({"success": True, "sandbox": sandbox})
@@ -294,6 +296,8 @@ async def call(
         send_to_user: bool = True,
         return_image_to_llm: bool = True,
     ) -> ToolExecResult:
+        if err := check_admin_permission(context, "Taking CUA sandbox screenshots"):
+            return err
         try:
             booter = await get_cua_sandbox_observer_booter_by_id(sandbox_id)
             gui = getattr(booter, "gui", None)

dashboard/src/views/SandboxManagementPage.vue

@@ -85,7 +85,7 @@
                   {{ tm('actions.console') }}
                   <v-tooltip activator="parent" location="top">{{ tm('tooltips.console') }}</v-tooltip>
                 </v-btn>
-                <v-btn size="small" variant="tonal" :disabled="!isCua(item) || !hasController(item)" @click="releaseSandbox(item)">{{ tm('actions.release') }}</v-btn>
+                <v-btn size="small" variant="tonal" :disabled="!canReleaseFromDashboard(item)" @click="releaseSandbox(item)">{{ tm('actions.release') }}</v-btn>
                 <v-btn size="small" variant="tonal" :disabled="!isCua(item)" @click="screenshotSandbox(item)">{{ tm('actions.screenshot') }}</v-btn>
                 <v-btn size="small" color="error" variant="tonal" :disabled="!isCua(item)" @click="openDestroyConfirm(item)">{{ tm('actions.destroy') }}</v-btn>
               </div>
@@ -375,6 +375,10 @@ function hasController(item: SandboxRecord) {
   return !!item.controller_session_id
 }
 
+function canReleaseFromDashboard(item: SandboxRecord) {
+  return isCua(item) && item.controller_session_id === 'dashboard'
+}
+
 function formatTime(value?: number | null) {
   if (!value) return '-'
   return new Date(value * 1000).toLocaleString()
@@ -517,6 +521,7 @@ async function confirmDestroySandbox() {
 
 async function screenshotSandbox(item: SandboxRecord) {
   const data = await postAction('/api/sandboxes/screenshot', { sandbox_id: item.sandbox_id })
+  if (!data) return
   const screenshot = data?.screenshot
   const legacyResult = data?.result
   const mimeType = screenshot?.mime_type || legacyResult?.mime_type || 'image/png'
@@ -567,7 +572,7 @@ function quoteForShell(value: string) {
 
 function buildConsoleShellCommand(command: string, cwd: string) {
   const prefix = cwd && cwd !== '~' ? `cd ${quoteForShell(cwd)} && ` : ''
-  return `${prefix}${command}; printf '\n__ASTRBOT_CWD__%s\n' "$PWD"`
+  return `${prefix}{ ${command}; __astrbot_status=$?; }; printf '\n__ASTRBOT_CWD__%s\n' "$PWD"; exit $__astrbot_status`
 }
 
 function parseConsoleShellResult(stdout: string, fallbackCwd: string) {

tests/unit/test_core_lifecycle.py

@@ -942,6 +942,31 @@ async def test_restart_terminates_managers_and_starts_thread(
             mock_thread.assert_called_once()
             mock_thread.return_value.start.assert_called_once()
 
+    @pytest.mark.asyncio
+    async def test_restart_cleans_up_managed_cua_sandboxes(
+        self, mock_log_broker, mock_db
+    ):
+        lifecycle = AstrBotCoreLifecycle(mock_log_broker, mock_db)
+        lifecycle.provider_manager = MagicMock()
+        lifecycle.provider_manager.terminate = AsyncMock()
+        lifecycle.platform_manager = MagicMock()
+        lifecycle.platform_manager.terminate = AsyncMock()
+        lifecycle.kb_manager = MagicMock()
+        lifecycle.kb_manager.terminate = AsyncMock()
+        lifecycle.dashboard_shutdown_event = asyncio.Event()
+        lifecycle.astrbot_updator = MagicMock()
+
+        with (
+            patch("astrbot.core.core_lifecycle.threading.Thread"),
+            patch(
+                "astrbot.core.core_lifecycle.cleanup_managed_cua_sandboxes",
+                new_callable=AsyncMock,
+            ) as cleanup,
+        ):
+            await lifecycle.restart()
+
+        cleanup.assert_awaited_once()
+
 
 class TestAstrBotCoreLifecycleLoadPipelineScheduler:
     """Tests for AstrBotCoreLifecycle.load_pipeline_scheduler method."""

tests/unit/test_cua_computer_use.py

@@ -546,6 +546,92 @@ async def available(self):
         await computer_client.get_booter(ctx, "session-b")
 
 
+@pytest.mark.asyncio
+async def test_get_booter_uses_selected_running_sandbox_instead_of_default(
+    monkeypatch,
+    tmp_path,
+):
+    from astrbot.core.computer import computer_client
+    from astrbot.core.computer.cua_registry import CuaSandboxRegistry
+
+    class FakeBooter:
+        def __init__(self, sandbox_id: str):
+            self.sandbox_id = sandbox_id
+
+        async def available(self):
+            return True
+
+    registry = CuaSandboxRegistry(storage_path=tmp_path / "registry.json")
+    registry.upsert_sandbox(
+        sandbox_id="sb-default",
+        sandbox_name="default",
+        booter_type="cua",
+        provider="cua",
+        managed=True,
+        created_by_astrbot=True,
+        owner_user_id="session-a",
+        owner_session_id="session-a",
+        connect_info={},
+        is_default=True,
+    )
+    registry.upsert_sandbox(
+        sandbox_id="sb-selected",
+        sandbox_name="selected",
+        booter_type="cua",
+        provider="cua",
+        managed=True,
+        created_by_astrbot=True,
+        owner_user_id="session-a",
+        owner_session_id="session-a",
+        connect_info={},
+    )
+    registry.set_current_sandbox_id("session-a", "sb-selected")
+    monkeypatch.setattr(computer_client, "cua_registry", registry)
+    computer_client.session_booter.clear()
+    computer_client.session_booter["sb-default"] = FakeBooter("sb-default")
+    computer_client.session_booter["sb-selected"] = FakeBooter("sb-selected")
+
+    booter = await computer_client.get_booter(
+        FakeContext(
+            {
+                "provider_settings": {
+                    "computer_use_runtime": "sandbox",
+                    "sandbox": {"booter": "cua"},
+                }
+            }
+        ),
+        "session-a",
+    )
+
+    assert booter.sandbox_id == "sb-selected"
+
+
+@pytest.mark.asyncio
+async def test_switch_current_rejects_non_running_sandbox(monkeypatch, tmp_path):
+    from astrbot.core.computer import computer_client
+    from astrbot.core.computer.cua_registry import CuaSandboxRegistry
+
+    registry = CuaSandboxRegistry(storage_path=tmp_path / "registry.json")
+    registry.upsert_sandbox(
+        sandbox_id="sb-stale",
+        sandbox_name="stale",
+        booter_type="cua",
+        provider="cua",
+        managed=True,
+        created_by_astrbot=True,
+        owner_user_id="session-a",
+        owner_session_id="session-a",
+        connect_info={},
+    )
+    monkeypatch.setattr(computer_client, "cua_registry", registry)
+    computer_client.session_booter.clear()
+
+    with pytest.raises(RuntimeError, match="not running"):
+        computer_client.switch_current_cua_sandbox("session-a", "sb-stale")
+
+    assert registry.get_current_sandbox_id("session-a") is None
+
+
 @pytest.mark.asyncio
 async def test_busy_sandbox_screenshot_is_allowed_without_taking_control(
     monkeypatch, tmp_path
@@ -935,6 +1021,53 @@ async def available(self):
     )
 
 
+@pytest.mark.asyncio
+async def test_get_booter_serializes_concurrent_default_boot(monkeypatch, tmp_path):
+    from astrbot.core.computer import computer_client
+    from astrbot.core.computer.cua_registry import CuaSandboxRegistry
+
+    boot_started = asyncio.Event()
+    release_boot = asyncio.Event()
+    boots = []
+
+    class FakeBooter:
+        def __init__(self, sandbox_id: str):
+            self.sandbox_id = sandbox_id
+
+        async def available(self):
+            return True
+
+    async def fake_boot_managed(ctx, session_id, sandbox_id, cua_kwargs):
+        boots.append(sandbox_id)
+        boot_started.set()
+        await release_boot.wait()
+        return FakeBooter(sandbox_id)
+
+    registry = CuaSandboxRegistry(storage_path=tmp_path / "registry.json")
+    monkeypatch.setattr(computer_client, "cua_registry", registry)
+    monkeypatch.setattr(computer_client, "_boot_managed_cua_sandbox", fake_boot_managed)
+    computer_client.session_booter.clear()
+
+    ctx = FakeContext(
+        {
+            "provider_settings": {
+                "computer_use_runtime": "sandbox",
+                "sandbox": {"booter": "cua"},
+            }
+        }
+    )
+    task_one = asyncio.create_task(computer_client.get_booter(ctx, "session-a"))
+    await boot_started.wait()
+    task_two = asyncio.create_task(computer_client.get_booter(ctx, "session-a"))
+    await asyncio.sleep(0)
+    release_boot.set()
+
+    booter_one, booter_two = await asyncio.gather(task_one, task_two)
+
+    assert booter_one is booter_two
+    assert len(boots) == 1
+
+
 @pytest.mark.asyncio
 async def test_get_booter_shuts_down_client_when_skill_sync_fails(monkeypatch):
     from astrbot.core.computer import computer_client
@@ -1036,6 +1169,48 @@ async def shutdown(self):
 
     assert shutdowns == [booter.session_id]
     assert booter.sandbox_id not in computer_client.session_booter
+    assert registry.get_sandbox(booter.sandbox_id) is None
+
+
+@pytest.mark.asyncio
+async def test_cua_idle_timeout_marks_persistent_sandbox_stopped_after_lease_release(
+    monkeypatch,
+):
+    from astrbot.core.computer import computer_client
+    from astrbot.core.computer.cua_registry import CuaSandboxRegistry
+
+    shutdowns = []
+
+    class FakeBooter:
+        async def shutdown(self):
+            shutdowns.append("sb-persistent")
+
+    registry = CuaSandboxRegistry(storage_path=Path("/tmp/cua-idle-persistent.json"))
+    registry.upsert_sandbox(
+        sandbox_id="sb-persistent",
+        sandbox_name="persistent",
+        booter_type="cua",
+        provider="cua",
+        managed=True,
+        created_by_astrbot=True,
+        owner_user_id="session-a",
+        owner_session_id="session-a",
+        connect_info={},
+        retention_policy="persistent",
+        idle_timeout=0.05,
+    )
+    monkeypatch.setattr(computer_client, "cua_registry", registry)
+    computer_client.session_booter.clear()
+    computer_client.session_booter["sb-persistent"] = FakeBooter()
+
+    computer_client._schedule_cua_idle_cleanup("sb-persistent", 0.05)
+    await asyncio.sleep(0.1)
+
+    record = registry.get_sandbox("sb-persistent")
+    assert shutdowns == ["sb-persistent"]
+    assert record is not None
+    assert record["status"] == "stopped"
+    assert "sb-persistent" not in computer_client.session_booter
 
 
 @pytest.mark.asyncio