fix(sandbox): restore persistent sandboxes asynchronously

Author: zouyonghe

astrbot/core/computer/sandbox_manager.py

@@ -985,10 +985,7 @@ async def destroy_sandbox(self, session_id: str, sandbox_id: str) -> dict:
             finally:
                 self.session_booter.pop(sandbox_id, None)
         self.clear_idle_state(sandbox_id)
-        if record.get("retention_policy") == "persistent":
-            self.registry.update_sandbox_status(sandbox_id, SandboxStatus.STOPPED)
-        else:
-            self.registry.delete_sandbox(sandbox_id)
+        self.registry.delete_sandbox(sandbox_id)
         self.drop_boot_lock(sandbox_id)
         await self.save_registry_async()
 
@@ -1121,7 +1118,14 @@ async def reconcile_on_startup(self) -> None:
 
         await self.save_registry_async()
 
-    async def restore_persistent_sandboxes(self, context: Context) -> None:
+    async def restore_persistent_sandboxes(
+        self,
+        context: Context,
+        *,
+        per_sandbox_timeout: float | None = None,
+    ) -> tuple[int, int]:
+        restored = 0
+        deleted = 0
         for record in self.registry.list_sandboxes():
             sandbox_id = record["sandbox_id"]
             if not record.get("managed"):
@@ -1134,18 +1138,35 @@ async def restore_persistent_sandboxes(self, context: Context) -> None:
             }:
                 continue
             try:
-                await self._revive_persistent_booter_if_needed(
-                    record,
+                restore_coro = self._revive_persistent_booter_if_needed(
+                    record=record,
+                    sandbox_id=sandbox_id,
+                    session_id=str(record.get("owner_session_id") or "dashboard"),
+                    context=context,
+                )
+                if per_sandbox_timeout is None:
+                    await restore_coro
+                else:
+                    await asyncio.wait_for(restore_coro, timeout=per_sandbox_timeout)
+                restored += 1
+            except asyncio.TimeoutError:
+                self.session_booter.pop(sandbox_id, None)
+                self.clear_idle_state(sandbox_id)
+                self.registry.delete_sandbox(sandbox_id)
+                self.drop_boot_lock(sandbox_id)
+                await self.save_registry_async()
+                deleted += 1
+                logger.warning(
+                    "[Computer] Persistent sandbox restore timed out; removed stale record: %s",
                     sandbox_id,
-                    str(record.get("owner_session_id") or "dashboard"),
-                    context,
                 )
             except Exception as exc:
                 logger.warning(
                     "[Computer] Failed to restore persistent sandbox %s: %s",
                     sandbox_id,
                     exc,
                 )
+        return restored, deleted
 
     async def cleanup_managed_sandboxes(self) -> None:
         for sandbox_id in list(self.pending_boot_tasks):

astrbot/core/core_lifecycle.py

@@ -61,6 +61,7 @@ def __init__(self, log_broker: LogBroker, db: BaseDatabase) -> None:
         self.cron_manager: CronJobManager | None = None
         self.temp_dir_cleaner: TempDirCleaner | None = None
         self._default_chat_provider_warning_emitted = False
+        self._persistent_restore_task: asyncio.Task | None = None
 
         # 设置代理
         proxy_config = self.astrbot_config.get("http_proxy", "")
@@ -254,8 +255,9 @@ async def initialize(self) -> None:
                 exc_info=True,
             )
 
-        await computer_client.sandbox_manager.restore_persistent_sandboxes(
-            self.star_context
+        self._persistent_restore_task = asyncio.create_task(
+            self._restore_persistent_sandboxes_background(),
+            name="persistent-sandbox-restore",
         )
 
         # 根据配置实例化各个 Provider
@@ -292,6 +294,28 @@ async def initialize(self) -> None:
 
         asyncio.create_task(update_llm_metadata())
 
+    async def _restore_persistent_sandboxes_background(self) -> None:
+        try:
+            restored, deleted = (
+                await computer_client.sandbox_manager.restore_persistent_sandboxes(
+                    self.star_context,
+                    per_sandbox_timeout=30.0,
+                )
+            )
+            logger.info(
+                "Persistent sandbox restore finished: restored=%d deleted=%d",
+                restored,
+                deleted,
+            )
+        except asyncio.CancelledError:
+            raise
+        except Exception as e:
+            logger.warning(
+                "Persistent sandbox restore failed: %s",
+                e,
+                exc_info=True,
+            )
+
     def _load(self) -> None:
         """加载事件总线和任务并初始化."""
         # 创建一个异步任务来执行事件总线的 dispatch() 方法
@@ -384,6 +408,15 @@ async def stop(self) -> None:
         if self.cron_manager:
             await self.cron_manager.shutdown()
 
+        persistent_restore_task = getattr(self, "_persistent_restore_task", None)
+        if persistent_restore_task is not None:
+            persistent_restore_task.cancel()
+            try:
+                await persistent_restore_task
+            except asyncio.CancelledError:
+                pass
+            self._persistent_restore_task = None
+
         try:
             await computer_client.cleanup_managed_sandboxes()
         except Exception as e:

astrbot/core/initial_loader.py

@@ -25,33 +25,35 @@ def __init__(self, db: BaseDatabase, log_broker: LogBroker) -> None:
 
     async def start(self) -> None:
         core_lifecycle = AstrBotCoreLifecycle(self.log_broker, self.db)
-
-        try:
-            await core_lifecycle.initialize()
-        except Exception as e:
-            logger.critical(traceback.format_exc())
-            logger.critical(f"😭 初始化 AstrBot 失败：{e} !!!")
-            return
-
-        core_task = core_lifecycle.start()
-
-        webui_dir = self.webui_dir
-
-        self.dashboard_server = AstrBotDashboard(
-            core_lifecycle,
-            self.db,
-            core_lifecycle.dashboard_shutdown_event,
-            webui_dir,
-        )
-
-        coro = self.dashboard_server.run()
-        if coro:
-            # 启动核心任务和仪表板服务器
-            task = asyncio.gather(core_task, coro)
-        else:
-            task = core_task
         try:
-            await task  # 整个AstrBot在这里运行
+            try:
+                await core_lifecycle.initialize()
+            except Exception as e:
+                logger.critical(traceback.format_exc())
+                logger.critical(f"😭 初始化 AstrBot 失败：{e} !!!")
+                return
+
+            core_task = core_lifecycle.start()
+
+            webui_dir = self.webui_dir
+
+            self.dashboard_server = AstrBotDashboard(
+                core_lifecycle,
+                self.db,
+                core_lifecycle.dashboard_shutdown_event,
+                webui_dir,
+            )
+
+            coro = self.dashboard_server.run()
+            if coro:
+                # 启动核心任务和仪表板服务器
+                task = asyncio.gather(core_task, coro)
+            else:
+                task = core_task
+            try:
+                await task  # 整个AstrBot在这里运行
+            except asyncio.CancelledError:
+                logger.info("🌈 正在关闭 AstrBot...")
         except asyncio.CancelledError:
             logger.info("🌈 正在关闭 AstrBot...")
         finally:

tests/unit/test_core_lifecycle.py

@@ -261,7 +261,7 @@ async def test_subagent_orchestrator_error_is_logged(
 
 class TestAstrBotCoreLifecycleSandboxRestore:
     @pytest.mark.asyncio
-    async def test_initialize_restores_persistent_sandboxes(
+    async def test_initialize_restores_persistent_sandboxes_in_background(
         self, mock_log_broker, mock_db, mock_astrbot_config
     ):
         lifecycle = AstrBotCoreLifecycle(mock_log_broker, mock_db)
@@ -286,7 +286,13 @@ async def test_initialize_restores_persistent_sandboxes(
         mock_pipeline_scheduler = MagicMock(initialize=AsyncMock())
         mock_astrbot_updator = MagicMock()
         mock_event_bus = MagicMock()
-        restore_persistent = AsyncMock()
+        restore_started = asyncio.Event()
+        restore_finished = asyncio.Event()
+
+        async def restore_persistent(_context, **_kwargs):
+            restore_started.set()
+            await restore_finished.wait()
+            return 1, 0
 
         with (
             patch("astrbot.core.core_lifecycle.astrbot_config", mock_astrbot_config),
@@ -357,9 +363,14 @@ async def test_initialize_restores_persistent_sandboxes(
                 restore_persistent,
             ),
         ):
-            await lifecycle.initialize()
-
-        restore_persistent.assert_awaited_once_with(mock_star_context)
+            init_task = asyncio.create_task(lifecycle.initialize())
+            await asyncio.wait_for(restore_started.wait(), timeout=1)
+            await asyncio.wait_for(init_task, timeout=1)
+            assert not restore_finished.is_set()
+            restore_finished.set()
+            await asyncio.wait_for(asyncio.sleep(0), timeout=1)
+
+        assert restore_started.is_set()
 
 
 class TestAstrBotCoreLifecycleDefaultChatProviderWarning:

tests/unit/test_sandbox_computer_client.py

@@ -588,6 +588,7 @@ async def terminate_plugin(plugin):
     lifecycle.platform_manager.terminate = platform_terminate
     lifecycle.kb_manager.terminate = kb_terminate
     lifecycle.plugin_manager._terminate_plugin = terminate_plugin
+    lifecycle._persistent_restore_task = None
 
     await AstrBotCoreLifecycle.stop(lifecycle)
 

tests/unit/test_sandbox_manager.py

@@ -333,6 +333,24 @@ async def test_destroy_sandbox_cancels_deferred_boot_task(tmp_path):
     assert sandbox["sandbox_id"] not in manager.pending_boot_tasks
 
 
+@pytest.mark.asyncio
+async def test_destroy_persistent_sandbox_removes_record(tmp_path):
+    manager, provider = _manager(tmp_path)
+    created = await manager.create_sandbox(None, "session-a", "generic", "Named")
+    manager.update_sandbox_config(
+        created["sandbox_id"],
+        idle_timeout=None,
+        expires_at=None,
+        retention_policy="persistent",
+    )
+
+    destroyed = await manager.destroy_sandbox("session-a", created["sandbox_id"])
+
+    assert destroyed["sandbox_id"] == created["sandbox_id"]
+    assert manager.registry.get_sandbox(created["sandbox_id"]) is None
+    assert provider.destroyed[0][1] == created["sandbox_id"]
+
+
 @pytest.mark.asyncio
 async def test_create_sandbox_sets_current_sandbox_after_lease(tmp_path):
     manager, _provider = _manager(tmp_path)
@@ -673,7 +691,7 @@ async def test_manager_does_not_revive_destroyed_persistent_sandbox(tmp_path):
 
     await manager.destroy_sandbox("session-a", created["sandbox_id"])
 
-    with pytest.raises(RuntimeError, match="has been destroyed"):
+    with pytest.raises(RuntimeError, match="not found"):
         await manager.get_observer_booter_by_id(
             created["sandbox_id"],
             "dashboard",
@@ -923,6 +941,39 @@ async def test_manager_reconcile_on_startup_keeps_valid_persistent_records(
     assert len(provider.created) == 0
 
 
+@pytest.mark.asyncio
+async def test_manager_restore_persistent_sandboxes_times_out_and_deletes_record(
+    tmp_path,
+):
+    provider = FailingReconnectProvider()
+    manager, _provider = _manager(tmp_path, provider)
+    restore_started = asyncio.Event()
+
+    async def slow_create_booter(context, session_id, sandbox_id, config):
+        restore_started.set()
+        await asyncio.sleep(1)
+        return await FakeProvider().create_booter(context, session_id, sandbox_id, config)
+
+    provider.create_booter = slow_create_booter
+    manager.registry.upsert_sandbox(
+        sandbox_id="generic-1",
+        sandbox_name="Persistent",
+        provider="generic",
+        managed=True,
+        created_by_astrbot=True,
+        owner_user_id="session-a",
+        owner_session_id="session-a",
+        connect_info={"name": "Persistent"},
+        status="running",
+        retention_policy="persistent",
+    )
+
+    await manager.restore_persistent_sandboxes(object(), per_sandbox_timeout=0.01)
+
+    assert restore_started.is_set()
+    assert manager.registry.get_sandbox("generic-1") is None
+
+
 def test_manager_reconcile_on_startup_drops_temporary_records(tmp_path):
     manager, _provider = _manager(tmp_path)
     manager.registry.upsert_sandbox(