feat(cua): expire idle sandbox sessions (#8074)

* feat(cua): expire idle sandbox sessions

* fix(cua): simplify idle timeout state

Author: zouyonghe

astrbot/core/computer/computer_client.py

@@ -1,7 +1,10 @@
+import asyncio
 import json
 import os
 import shutil
+import time
 import uuid
+from dataclasses import dataclass
 from pathlib import Path
 
 from astrbot.api import logger
@@ -20,6 +23,70 @@
 _MANAGED_SKILLS_FILE = ".astrbot_managed_skills.json"
 
 
+@dataclass(slots=True)
+class _CUAIdleState:
+    expires_at: float
+    task: asyncio.Task
+
+
+cua_idle_state: dict[str, _CUAIdleState] = {}
+
+
+def _get_cua_idle_timeout(config: dict) -> float:
+    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)
+
+
+def _clear_cua_idle_state(session_id: str) -> None:
+    state = cua_idle_state.pop(session_id, None)
+    if state is not None and not state.task.done():
+        state.task.cancel()
+
+
+def _schedule_cua_idle_cleanup(session_id: str, timeout: float) -> None:
+    _clear_cua_idle_state(session_id)
+    if timeout <= 0:
+        return
+    expires_at = time.monotonic() + timeout
+
+    async def _expire_when_idle() -> None:
+        try:
+            remaining = expires_at - time.monotonic()
+            if remaining > 0:
+                await asyncio.sleep(remaining)
+
+            state = cua_idle_state.get(session_id)
+            if state is None or state.expires_at != expires_at:
+                return
+
+            booter = session_booter.get(session_id)
+            if booter is not None:
+                try:
+                    await booter.shutdown()
+                except Exception as shutdown_err:
+                    logger.warning(
+                        "[Computer] Failed to shutdown idle CUA sandbox for session %s: %s",
+                        session_id,
+                        shutdown_err,
+                    )
+                finally:
+                    session_booter.pop(session_id, None)
+        except asyncio.CancelledError:
+            raise
+        finally:
+            state = cua_idle_state.get(session_id)
+            if state is not None and state.expires_at == expires_at:
+                cua_idle_state.pop(session_id, None)
+
+    task = asyncio.create_task(_expire_when_idle())
+    cua_idle_state[session_id] = _CUAIdleState(expires_at=expires_at, task=task)
+
+
 def _list_local_skill_dirs(skills_root: Path) -> list[Path]:
     skills: list[Path] = []
     for entry in sorted(skills_root.iterdir()):
@@ -486,6 +553,7 @@ async def get_booter(
 
     sandbox_cfg = config.get("provider_settings", {}).get("sandbox", {})
     booter_type = sandbox_cfg.get("booter", "shipyard_neo")
+    cua_idle_timeout = _get_cua_idle_timeout(config) if booter_type == "cua" else 0.0
 
     if session_id in session_booter:
         booter = session_booter[session_id]
@@ -506,6 +574,7 @@ async def get_booter(
                     session_id,
                     shutdown_err,
                 )
+            _clear_cua_idle_state(session_id)
             session_booter.pop(session_id, None)
     if session_id not in session_booter:
         uuid_str = uuid.uuid5(uuid.NAMESPACE_DNS, session_id).hex
@@ -579,9 +648,12 @@ async def get_booter(
                     session_id,
                     shutdown_error,
                 )
+            _clear_cua_idle_state(session_id)
             raise e
 
         session_booter[session_id] = client
+    if booter_type == "cua":
+        _schedule_cua_idle_cleanup(session_id, cua_idle_timeout)
     return session_booter[session_id]
 
 

astrbot/core/config/default.py

@@ -179,6 +179,7 @@
             "cua_image": CUA_DEFAULT_CONFIG["image"],
             "cua_os_type": CUA_DEFAULT_CONFIG["os_type"],
             "cua_ttl": CUA_DEFAULT_CONFIG["ttl"],
+            "cua_idle_timeout": 0,
             "cua_telemetry_enabled": CUA_DEFAULT_CONFIG["telemetry_enabled"],
             "cua_local": CUA_DEFAULT_CONFIG["local"],
             "cua_api_key": CUA_DEFAULT_CONFIG["api_key"],

tests/unit/test_cua_computer_use.py

@@ -20,6 +20,13 @@ def get_config(self, umo: str | None = None):
         return self._config
 
 
+def _clear_cua_session_state(computer_client, session_id: str) -> None:
+    computer_client.session_booter.pop(session_id, None)
+    state = getattr(computer_client, "cua_idle_state", {}).pop(session_id, None)
+    if state is not None and not state.task.done():
+        state.task.cancel()
+
+
 class FakeShell:
     def __init__(self):
         self.commands = []
@@ -267,6 +274,7 @@ def test_cua_default_config_matches_booter_defaults():
     assert sandbox_defaults["cua_image"] == CUA_DEFAULT_CONFIG["image"]
     assert sandbox_defaults["cua_os_type"] == CUA_DEFAULT_CONFIG["os_type"]
     assert sandbox_defaults["cua_ttl"] == CUA_DEFAULT_CONFIG["ttl"]
+    assert sandbox_defaults["cua_idle_timeout"] == 0
     assert (
         sandbox_defaults["cua_telemetry_enabled"]
         == CUA_DEFAULT_CONFIG["telemetry_enabled"]
@@ -367,6 +375,189 @@ async def fail_sync(booter):
     assert "cua-sync-fail" not in computer_client.session_booter
 
 
+@pytest.mark.asyncio
+async def test_cua_idle_timeout_shuts_down_session_proactively(monkeypatch):
+    from astrbot.core.computer import computer_client
+
+    shutdowns = []
+
+    class FakeCuaBooter:
+        def __init__(self, **kwargs):
+            self.kwargs = kwargs
+
+        async def boot(self, session_id: str):
+            self.session_id = session_id
+
+        async def available(self):
+            return True
+
+        async def shutdown(self):
+            shutdowns.append(self.session_id)
+
+    monkeypatch.setattr(
+        computer_client, "_sync_skills_to_sandbox", lambda booter: asyncio.sleep(0)
+    )
+    monkeypatch.setattr(
+        "astrbot.core.computer.booters.cua.CuaBooter",
+        FakeCuaBooter,
+        raising=False,
+    )
+    _clear_cua_session_state(computer_client, "cua-idle-expire")
+
+    ctx = FakeContext(
+        {
+            "provider_settings": {
+                "computer_use_runtime": "sandbox",
+                "sandbox": {
+                    "booter": "cua",
+                    "cua_idle_timeout": 0.1,
+                },
+            }
+        }
+    )
+
+    booter = await computer_client.get_booter(ctx, "cua-idle-expire")
+    await asyncio.sleep(0.2)
+
+    assert shutdowns == [booter.session_id]
+    assert "cua-idle-expire" not in computer_client.session_booter
+
+
+@pytest.mark.asyncio
+async def test_cua_idle_timeout_refreshes_on_reuse(monkeypatch):
+    from astrbot.core.computer import computer_client
+
+    shutdowns = []
+
+    class FakeCuaBooter:
+        def __init__(self, **kwargs):
+            self.kwargs = kwargs
+
+        async def boot(self, session_id: str):
+            self.session_id = session_id
+
+        async def available(self):
+            return True
+
+        async def shutdown(self):
+            shutdowns.append(self.session_id)
+
+    monkeypatch.setattr(
+        computer_client, "_sync_skills_to_sandbox", lambda booter: asyncio.sleep(0)
+    )
+    monkeypatch.setattr(
+        "astrbot.core.computer.booters.cua.CuaBooter",
+        FakeCuaBooter,
+        raising=False,
+    )
+    _clear_cua_session_state(computer_client, "cua-idle-refresh")
+
+    ctx = FakeContext(
+        {
+            "provider_settings": {
+                "computer_use_runtime": "sandbox",
+                "sandbox": {
+                    "booter": "cua",
+                    "cua_idle_timeout": 0.2,
+                },
+            }
+        }
+    )
+
+    booter1 = await computer_client.get_booter(ctx, "cua-idle-refresh")
+    await asyncio.sleep(0.05)
+    booter2 = await computer_client.get_booter(ctx, "cua-idle-refresh")
+    await asyncio.sleep(0.05)
+
+    assert booter2 is booter1
+    assert shutdowns == []
+
+    await asyncio.sleep(0.25)
+
+    assert shutdowns == [booter1.session_id]
+    assert "cua-idle-refresh" not in computer_client.session_booter
+
+
+@pytest.mark.asyncio
+async def test_cua_idle_timeout_zero_disables_proactive_shutdown(monkeypatch):
+    from astrbot.core.computer import computer_client
+
+    shutdowns = []
+
+    class FakeCuaBooter:
+        def __init__(self, **kwargs):
+            self.kwargs = kwargs
+
+        async def boot(self, session_id: str):
+            self.session_id = session_id
+
+        async def available(self):
+            return True
+
+        async def shutdown(self):
+            shutdowns.append(self.session_id)
+
+    monkeypatch.setattr(
+        computer_client, "_sync_skills_to_sandbox", lambda booter: asyncio.sleep(0)
+    )
+    monkeypatch.setattr(
+        "astrbot.core.computer.booters.cua.CuaBooter",
+        FakeCuaBooter,
+        raising=False,
+    )
+    _clear_cua_session_state(computer_client, "cua-idle-disabled")
+
+    ctx = FakeContext(
+        {
+            "provider_settings": {
+                "computer_use_runtime": "sandbox",
+                "sandbox": {
+                    "booter": "cua",
+                    "cua_idle_timeout": 0,
+                },
+            }
+        }
+    )
+
+    await computer_client.get_booter(ctx, "cua-idle-disabled")
+    await asyncio.sleep(0.05)
+
+    assert shutdowns == []
+    assert "cua-idle-disabled" in computer_client.session_booter
+    assert "cua-idle-disabled" not in computer_client.cua_idle_state
+
+
+@pytest.mark.asyncio
+async def test_non_cua_booter_does_not_schedule_idle_cleanup(monkeypatch):
+    from astrbot.core.computer import computer_client
+
+    class FakeShipyardBooter:
+        async def available(self):
+            return True
+
+    _clear_cua_session_state(computer_client, "shipyard-session")
+    computer_client.session_booter["shipyard-session"] = FakeShipyardBooter()
+
+    ctx = FakeContext(
+        {
+            "provider_settings": {
+                "computer_use_runtime": "sandbox",
+                "sandbox": {
+                    "booter": "shipyard",
+                    "shipyard_endpoint": "http://localhost:8080",
+                    "shipyard_access_token": "token",
+                    "cua_idle_timeout": 0.01,
+                },
+            }
+        }
+    )
+
+    booter = await computer_client.get_booter(ctx, "shipyard-session")
+
+    assert isinstance(booter, FakeShipyardBooter)
+    assert "shipyard-session" not in computer_client.cua_idle_state
+
+
 @pytest.mark.asyncio
 async def test_cua_components_map_sdk_results(tmp_path):
     from astrbot.core.computer.booters.cua import (