backend: reserve restore slots and extract idle cleanup constants

backend/session_manager.py

@@ -124,6 +124,8 @@ class SessionCapacityError(Exception):
 MAX_SESSIONS: int = 200
 MAX_SESSIONS_PER_USER: int = 10
 DEFAULT_YOLO_COST_CAP_USD: float = 5.0
+INACTIVE_SESSION_SWEEP_INTERVAL_SECONDS: int = 600
+INACTIVE_SESSION_IDLE_THRESHOLD: timedelta = timedelta(hours=24)
 SANDBOX_SHUTDOWN_CLEANUP_CONCURRENCY: int = 10
 SANDBOX_SHUTDOWN_CLEANUP_TIMEOUT_S: float = 60.0
 
@@ -208,6 +210,40 @@ class SessionManager:
         logger.info("Session initialized in %.2fs", t1 - t0)
         return tool_router, session
 
+    def _make_reserved_session(
+        self,
+        *,
+        session_id: str,
+        user_id: str,
+        hf_username: str | None,
+        hf_token: str | None,
+        submission_queue: asyncio.Queue,
+    ) -> AgentSession:
+        """Create a placeholder session that reserves capacity under lock."""
+        return AgentSession(
+            session_id=session_id,
+            session=None,
+            tool_router=None,
+            submission_queue=submission_queue,
+            user_id=user_id,
+            hf_username=hf_username,
+            hf_token=hf_token,
+            is_active=True,
+        )
+
+    async def _release_reserved_session_slot(
+        self,
+        session_id: str,
+        reserved_session: AgentSession | None = None,
+    ) -> None:
+        """Remove a reserved placeholder if it is still present."""
+        async with self._lock:
+            current = self.sessions.get(session_id)
+            if current is None:
+                return
+            if current is reserved_session or getattr(current, "session", None) is None:
+                self.sessions.pop(session_id, None)
+
     def _serialize_messages(self, session: Session) -> list[dict[str, Any]]:
         return [msg.model_dump(mode="json") for msg in session.context_manager.items]
 
@@ -524,7 +560,7 @@ class SessionManager:
         """
         try:
             while True:
-                await asyncio.sleep(600)  # 10 minutes
+                await asyncio.sleep(INACTIVE_SESSION_SWEEP_INTERVAL_SECONDS)
                 now = datetime.utcnow()
                 to_unload: list[tuple[str, AgentSession]] = []
                 async with self._lock:
@@ -534,12 +570,16 @@ class SessionManager:
                                 continue
                             if getattr(agent_session, "is_processing", False):
                                 continue
-                            last = getattr(agent_session, "last_access", agent_session.created_at)
-                            if now - last > timedelta(hours=24):
-                                # Mark inactive and remove from registry under lock
+                            last = getattr(
+                                agent_session,
+                                "last_access",
+                                agent_session.created_at,
+                            )
+                            if now - last > INACTIVE_SESSION_IDLE_THRESHOLD:
+                                # Mark inactive, but keep it resident until the
+                                # snapshot has been persisted successfully.
                                 agent_session.is_active = False
                                 to_unload.append((sid, agent_session))
-                                del self.sessions[sid]
                         except Exception as e:
                             logger.debug("Skipping unload check for %s: %s", sid, e)
 
@@ -556,7 +596,24 @@ class SessionManager:
                             sid,
                             e,
                         )
-                    # Cancel the running task if present
+                        # Keep the session in memory so the next cleanup cycle
+                        # can retry persistence and so callers can still inspect
+                        # the session state.
+                        agent_session.is_active = True
+                        continue
+
+                    removed = False
+                    async with self._lock:
+                        current = self.sessions.get(sid)
+                        if current is agent_session:
+                            self.sessions.pop(sid, None)
+                            removed = True
+
+                    if not removed:
+                        # Session was replaced or revived while we were persisting.
+                        continue
+
+                    # Cancel the running task if present once the snapshot is safe.
                     try:
                         if agent_session.task:
                             agent_session.task.cancel()
@@ -631,146 +688,171 @@ class SessionManager:
         preload_sandbox: bool = True,
     ) -> AgentSession | None:
         """Return a live runtime session, lazily restoring it from Mongo."""
-        async with self._lock:
-            existing = self.sessions.get(session_id)
-        if existing:
-            if self._can_access_session(existing, user_id):
-                self._update_hf_identity(
-                    existing,
-                    hf_token=hf_token,
-                    hf_username=hf_username,
-                )
-                self._restart_cpu_preload_if_token_recovered(
-                    existing,
-                    preload_sandbox=preload_sandbox,
-                )
-                existing.last_access = datetime.utcnow()
-                return existing
-            return None
+        submission_queue: asyncio.Queue = asyncio.Queue()
+        event_queue: asyncio.Queue = asyncio.Queue()
+        reserved_session: AgentSession | None = None
+        should_release_reserved_slot = False
 
-        # Check capacity before restoring from persistence
-        async with self._lock:
-            active_count = self.active_session_count
-            if active_count >= MAX_SESSIONS:
-                logger.warning(
-                    "Cannot restore session %s: server at capacity (%d/%d)",
-                    session_id,
-                    active_count,
-                    MAX_SESSIONS,
-                )
-                return None
+        try:
+            async with self._lock:
+                existing = self.sessions.get(session_id)
+                if existing:
+                    if getattr(existing, "session", None) is None:
+                        return None
+                    if self._can_access_session(existing, user_id):
+                        self._update_hf_identity(
+                            existing,
+                            hf_token=hf_token,
+                            hf_username=hf_username,
+                        )
+                        self._restart_cpu_preload_if_token_recovered(
+                            existing,
+                            preload_sandbox=preload_sandbox,
+                        )
+                        existing.last_access = datetime.utcnow()
+                        return existing
+                    return None
 
-        store = self._store()
-        loaded = await store.load_session(session_id)
-        if not loaded:
-            return None
+                active_count = self.active_session_count
+                if active_count >= MAX_SESSIONS:
+                    logger.warning(
+                        "Cannot restore session %s: server at capacity (%d/%d)",
+                        session_id,
+                        active_count,
+                        MAX_SESSIONS,
+                    )
+                    return None
 
-        async with self._lock:
-            existing = self.sessions.get(session_id)
-        if existing:
-            if self._can_access_session(existing, user_id):
-                self._update_hf_identity(
-                    existing,
-                    hf_token=hf_token,
+                reserved_session = self._make_reserved_session(
+                    session_id=session_id,
+                    user_id=user_id,
                     hf_username=hf_username,
+                    hf_token=hf_token,
+                    submission_queue=submission_queue,
                 )
-                self._restart_cpu_preload_if_token_recovered(
-                    existing,
-                    preload_sandbox=preload_sandbox,
-                )
-                existing.last_access = datetime.utcnow()
-                return existing
-            return None
+                self.sessions[session_id] = reserved_session
+                should_release_reserved_slot = True
 
-        meta = loaded.get("metadata") or {}
-        owner = str(meta.get("user_id") or "")
-        if user_id != "dev" and owner != "dev" and owner != user_id:
-            return None
+            store = self._store()
+            loaded = await store.load_session(session_id)
+            if not loaded:
+                return None
 
-        await self._cleanup_persisted_sandbox(
-            session_id,
-            meta,
-            hf_token=hf_token,
-        )
+            async with self._lock:
+                existing = self.sessions.get(session_id)
+                if existing is not reserved_session:
+                    if existing and getattr(existing, "session", None) is not None:
+                        if self._can_access_session(existing, user_id):
+                            self._update_hf_identity(
+                                existing,
+                                hf_token=hf_token,
+                                hf_username=hf_username,
+                            )
+                            self._restart_cpu_preload_if_token_recovered(
+                                existing,
+                                preload_sandbox=preload_sandbox,
+                            )
+                            existing.last_access = datetime.utcnow()
+                            should_release_reserved_slot = False
+                            return existing
+                    return None
+
+            meta = loaded.get("metadata") or {}
+            owner = str(meta.get("user_id") or "")
+            if user_id != "dev" and owner != "dev" and owner != user_id:
+                return None
 
-        from litellm import Message
+            await self._cleanup_persisted_sandbox(
+                session_id,
+                meta,
+                hf_token=hf_token,
+            )
 
-        model = meta.get("model") or self.config.model_name
-        event_queue: asyncio.Queue = asyncio.Queue()
-        submission_queue: asyncio.Queue = asyncio.Queue()
-        tool_router, session = await asyncio.to_thread(
-            self._create_session_sync,
-            session_id=session_id,
-            user_id=owner or user_id,
-            hf_username=hf_username,
-            hf_token=hf_token,
-            model=model,
-            event_queue=event_queue,
-            notification_destinations=meta.get("notification_destinations") or [],
-        )
+            from litellm import Message
 
-        restored_messages: list[Message] = []
-        for raw in loaded.get("messages") or []:
-            if not isinstance(raw, dict) or raw.get("role") == "system":
-                continue
-            try:
-                restored_messages.append(Message.model_validate(raw))
-            except Exception as e:
-                logger.warning("Dropping malformed restored message: %s", e)
-        if restored_messages:
-            # Keep the freshly-rendered system prompt, then attach the durable
-            # non-system context so tools/date/user context stay current.
-            session.context_manager.items = [
-                session.context_manager.items[0],
-                *restored_messages,
-            ]
-
-        self._restore_pending_approval(session, meta.get("pending_approval") or [])
-        session.turn_count = int(meta.get("turn_count") or 0)
-        session.auto_approval_enabled = bool(meta.get("auto_approval_enabled", False))
-        raw_cap = meta.get("auto_approval_cost_cap_usd")
-        session.auto_approval_cost_cap_usd = (
-            float(raw_cap) if isinstance(raw_cap, int | float) else None
-        )
-        session.auto_approval_estimated_spend_usd = float(
-            meta.get("auto_approval_estimated_spend_usd") or 0.0
-        )
+            model = meta.get("model") or self.config.model_name
+            tool_router, session = await asyncio.to_thread(
+                self._create_session_sync,
+                session_id=session_id,
+                user_id=owner or user_id,
+                hf_username=hf_username,
+                hf_token=hf_token,
+                model=model,
+                event_queue=event_queue,
+                notification_destinations=meta.get("notification_destinations") or [],
+            )
 
-        created_at = meta.get("created_at")
-        if not isinstance(created_at, datetime):
-            created_at = datetime.utcnow()
+            restored_messages: list[Message] = []
+            for raw in loaded.get("messages") or []:
+                if not isinstance(raw, dict) or raw.get("role") == "system":
+                    continue
+                try:
+                    restored_messages.append(Message.model_validate(raw))
+                except Exception as e:
+                    logger.warning("Dropping malformed restored message: %s", e)
+            if restored_messages:
+                # Keep the freshly-rendered system prompt, then attach the durable
+                # non-system context so tools/date/user context stay current.
+                session.context_manager.items = [
+                    session.context_manager.items[0],
+                    *restored_messages,
+                ]
+
+            self._restore_pending_approval(session, meta.get("pending_approval") or [])
+            session.turn_count = int(meta.get("turn_count") or 0)
+            session.auto_approval_enabled = bool(
+                meta.get("auto_approval_enabled", False)
+            )
+            raw_cap = meta.get("auto_approval_cost_cap_usd")
+            session.auto_approval_cost_cap_usd = (
+                float(raw_cap) if isinstance(raw_cap, int | float) else None
+            )
+            session.auto_approval_estimated_spend_usd = float(
+                meta.get("auto_approval_estimated_spend_usd") or 0.0
+            )
 
-        agent_session = AgentSession(
-            session_id=session_id,
-            session=session,
-            tool_router=tool_router,
-            submission_queue=submission_queue,
-            user_id=owner or user_id,
-            hf_username=hf_username,
-            hf_token=hf_token,
-            created_at=created_at,
-            is_active=True,
-            is_processing=False,
-            claude_counted=bool(meta.get("claude_counted")),
-            title=meta.get("title"),
-        )
-        started = await self._start_agent_session(
-            agent_session=agent_session,
-            event_queue=event_queue,
-            tool_router=tool_router,
-        )
-        if started is not agent_session:
-            self._update_hf_identity(
-                started,
-                hf_token=hf_token,
+            created_at = meta.get("created_at")
+            if not isinstance(created_at, datetime):
+                created_at = datetime.utcnow()
+
+            agent_session = AgentSession(
+                session_id=session_id,
+                session=session,
+                tool_router=tool_router,
+                submission_queue=submission_queue,
+                user_id=owner or user_id,
                 hf_username=hf_username,
+                hf_token=hf_token,
+                created_at=created_at,
+                is_active=True,
+                is_processing=False,
+                claude_counted=bool(meta.get("claude_counted")),
+                title=meta.get("title"),
             )
-            started.last_access = datetime.utcnow()
-            return started
-        if preload_sandbox:
-            self._start_cpu_sandbox_preload(agent_session)
-        logger.info("Restored session %s for user %s", session_id, owner or user_id)
-        return agent_session
+            started = await self._start_agent_session(
+                agent_session=agent_session,
+                event_queue=event_queue,
+                tool_router=tool_router,
+            )
+            if started is not agent_session:
+                self._update_hf_identity(
+                    started,
+                    hf_token=hf_token,
+                    hf_username=hf_username,
+                )
+                started.last_access = datetime.utcnow()
+                should_release_reserved_slot = False
+                return started
+            if preload_sandbox:
+                self._start_cpu_sandbox_preload(agent_session)
+            logger.info("Restored session %s for user %s", session_id, owner or user_id)
+            should_release_reserved_slot = False
+            return agent_session
+        except Exception:
+            raise
+        finally:
+            if should_release_reserved_slot:
+                await self._release_reserved_session_slot(session_id, reserved_session)
 
     async def create_session(
         self,
@@ -838,41 +920,64 @@ class SessionManager:
             self.sessions[session_id] = placeholder
 
         # Run blocking constructors in a thread to keep the event loop responsive.
-        tool_router, session = await asyncio.to_thread(
-            self._create_session_sync,
-            session_id=session_id,
-            user_id=user_id,
-            hf_username=hf_username,
-            hf_token=hf_token,
-            model=model,
-            event_queue=event_queue,
-        )
+        agent_session: AgentSession | None = None
+        try:
+            tool_router, session = await asyncio.to_thread(
+                self._create_session_sync,
+                session_id=session_id,
+                user_id=user_id,
+                hf_username=hf_username,
+                hf_token=hf_token,
+                model=model,
+                event_queue=event_queue,
+            )
 
-        # Create wrapper with the real session resources and replace the
-        # placeholder in _start_agent_session.
-        agent_session = AgentSession(
-            session_id=session_id,
-            session=session,
-            tool_router=tool_router,
-            submission_queue=submission_queue,
-            user_id=user_id,
-            hf_username=hf_username,
-            hf_token=hf_token,
-        )
+            # Create wrapper with the real session resources and replace the
+            # placeholder in _start_agent_session.
+            agent_session = AgentSession(
+                session_id=session_id,
+                session=session,
+                tool_router=tool_router,
+                submission_queue=submission_queue,
+                user_id=user_id,
+                hf_username=hf_username,
+                hf_token=hf_token,
+            )
 
-        await self._start_agent_session(
-            agent_session=agent_session,
-            event_queue=event_queue,
-            tool_router=tool_router,
-        )
-        await self.persist_session_snapshot(agent_session, runtime_state="idle")
-        self._start_cpu_sandbox_preload(agent_session)
+            await self._start_agent_session(
+                agent_session=agent_session,
+                event_queue=event_queue,
+                tool_router=tool_router,
+            )
+            await self.persist_session_snapshot(agent_session, runtime_state="idle")
+            self._start_cpu_sandbox_preload(agent_session)
 
-        if is_pro is not None and user_id and user_id != "dev":
-            await self._track_pro_status(agent_session, is_pro=is_pro)
+            if is_pro is not None and user_id and user_id != "dev":
+                await self._track_pro_status(agent_session, is_pro=is_pro)
 
-        logger.info(f"Created session {session_id} for user {user_id}")
-        return session_id
+            logger.info(f"Created session {session_id} for user {user_id}")
+            return session_id
+        except Exception:
+            cleanup_task: asyncio.Task | None = None
+            async with self._lock:
+                current = self.sessions.get(session_id)
+                if current and (
+                    current is agent_session
+                    or getattr(current, "session", None) is None
+                ):
+                    self.sessions.pop(session_id, None)
+                    if agent_session is not None:
+                        agent_session.is_active = False
+                        cleanup_task = agent_session.task
+            if cleanup_task and not cleanup_task.done():
+                cleanup_task.cancel()
+                try:
+                    await cleanup_task
+                except asyncio.CancelledError:
+                    pass
+                except Exception:
+                    pass
+            raise
 
     async def _track_pro_status(
         self, agent_session: AgentSession, *, is_pro: bool

tests/unit/test_session_capacity.py

@@ -1,4 +1,5 @@
 import asyncio
+import importlib
 import sys
 from pathlib import Path
 from types import SimpleNamespace
@@ -6,31 +7,57 @@ from types import SimpleNamespace
 import pytest
 import types
 
-# Prevent importing heavy third-party modules when importing the backend module.
-for _mod in ("litellm", "fastmcp", "thefuzz", "huggingface_hub"):
-    if _mod not in sys.modules:
-        m = types.ModuleType(_mod)
-        # fastmcp is imported as `from fastmcp import Client` in some codepaths
-        if _mod == "fastmcp":
-            class _DummyClient:
-                pass
+_BACKEND_DIR = Path(__file__).resolve().parent.parent.parent / "backend"
 
-            setattr(m, "Client", _DummyClient)
-        sys.modules[_mod] = m
 
-_BACKEND_DIR = Path(__file__).resolve().parent.parent.parent / "backend"
-if str(_BACKEND_DIR) not in sys.path:
-    sys.path.insert(0, str(_BACKEND_DIR))
+@pytest.fixture
+def session_manager_module(monkeypatch):
+    """Import backend.session_manager with temporary dependency stubs.
+
+    The stubs are inserted with monkeypatch so they are restored after the test,
+    and the imported session_manager module is removed from sys.modules to avoid
+    leaking the stubbed import state into other tests.
+    """
+    with monkeypatch.context() as m:
+        m.syspath_prepend(str(_BACKEND_DIR))
+
+        litellm_stub = types.ModuleType("litellm")
 
-from session_manager import SessionManager, AgentSession, MAX_SESSIONS
+        class _DummyMessage:
+            @staticmethod
+            def model_validate(raw):
+                return SimpleNamespace(**raw)
+
+        setattr(litellm_stub, "Message", _DummyMessage)
+        m.setitem(sys.modules, "litellm", litellm_stub)
+
+        fastmcp_stub = types.ModuleType("fastmcp")
+
+        class _DummyClient:
+            pass
+
+        setattr(fastmcp_stub, "Client", _DummyClient)
+        m.setitem(sys.modules, "fastmcp", fastmcp_stub)
+
+        m.setitem(sys.modules, "thefuzz", types.ModuleType("thefuzz"))
+        m.setitem(
+            sys.modules,
+            "huggingface_hub",
+            types.ModuleType("huggingface_hub"),
+        )
+
+        module = importlib.import_module("session_manager")
+        yield module
+
+    sys.modules.pop("session_manager", None)
 
 
 @pytest.mark.asyncio
-async def test_restore_denied_when_at_capacity(caplog):
-    manager = SessionManager()
+async def test_restore_denied_when_at_capacity(session_manager_module, caplog):
+    manager = session_manager_module.SessionManager()
     # Fill in-memory sessions up to MAX_SESSIONS
-    for i in range(MAX_SESSIONS):
-        manager.sessions[str(i)] = AgentSession(
+    for i in range(session_manager_module.MAX_SESSIONS):
+        manager.sessions[str(i)] = session_manager_module.AgentSession(
             session_id=str(i),
             session=object(),
             tool_router=None,
@@ -53,8 +80,8 @@ async def test_restore_denied_when_at_capacity(caplog):
 
 
 @pytest.mark.asyncio
-async def test_restore_allowed_under_capacity(monkeypatch):
-    manager = SessionManager()
+async def test_restore_allowed_under_capacity(session_manager_module, monkeypatch):
+    manager = session_manager_module.SessionManager()
     manager.sessions.clear()
 
     class DummyStore:
@@ -87,4 +114,39 @@ async def test_restore_allowed_under_capacity(monkeypatch):
 
     res = await manager.ensure_session_loaded("restored", user_id="test")
     assert res is not None
-    assert res.session is not None
+    assert res.session is not None
+
+
+@pytest.mark.asyncio
+async def test_restore_rolls_back_placeholder_on_load_failure(session_manager_module):
+    manager = session_manager_module.SessionManager()
+
+    class FailingStore:
+        enabled = True
+
+        async def load_session(self, sid):
+            raise RuntimeError("load failed")
+
+    manager.persistence_store = FailingStore()
+
+    with pytest.raises(RuntimeError, match="load failed"):
+        await manager.ensure_session_loaded("restored", user_id="test")
+
+    assert manager.sessions == {}
+
+
+@pytest.mark.asyncio
+async def test_create_session_rolls_back_placeholder_on_failure(
+    session_manager_module, monkeypatch
+):
+    manager = session_manager_module.SessionManager()
+
+    def fake_create_session_sync(**kwargs):
+        raise RuntimeError("boom")
+
+    monkeypatch.setattr(manager, "_create_session_sync", fake_create_session_sync)
+
+    with pytest.raises(RuntimeError, match="boom"):
+        await manager.create_session(user_id="u1")
+
+    assert manager.sessions == {}