"""Tests for the idle-session reaper and the global create-slot reservation. Covers Parts B (hard global cap), C (idle reaper + activity stamps + safe teardown + submit/reap race), and E (per-user concurrent cap interacts with reaping) of the session-limit fix. """ from __future__ import annotations import asyncio import sys from datetime import datetime, timedelta from pathlib import Path from types import SimpleNamespace from typing import Any import pytest _BACKEND_DIR = Path(__file__).resolve().parent.parent.parent / "backend" if str(_BACKEND_DIR) not in sys.path: sys.path.insert(0, str(_BACKEND_DIR)) import session_manager as sm # noqa: E402 from agent.core.session_persistence import NoopSessionStore # noqa: E402 from session_manager import ( # noqa: E402 AgentSession, Operation, SessionCapacityError, SessionManager, ) from agent.core.session import OpType # noqa: E402 def test_reaper_idle_default_is_fifteen_minutes(): assert sm.REAPER_IDLE_MINUTES == 15 assert sm.REAPER_IDLE == timedelta(minutes=15) class RecordingStore(NoopSessionStore): """Captures every save_snapshot call so tests can assert persistence.""" enabled = True def __init__(self) -> None: self.snapshots: list[dict[str, Any]] = [] async def save_snapshot(self, **kwargs: Any) -> None: self.snapshots.append(kwargs) def snapshots_for(self, session_id: str) -> list[dict[str, Any]]: return [s for s in self.snapshots if s.get("session_id") == session_id] class FakeSession: """Minimal Session stand-in supporting both persistence and _run_session.""" def __init__( self, *, hf_token: str | None = "token", user_plan: str | None = None, ) -> None: self.hf_token = hf_token self.user_plan = user_plan self.context_manager = SimpleNamespace(items=[]) self.pending_approval: Any = None self.turn_count = 0 self.config = SimpleNamespace(model_name="test-model", save_sessions=False) self.notification_destinations: list[str] = [] self.auto_approval_enabled = False self.auto_approval_cost_cap_usd = None self.auto_approval_estimated_spend_usd = 0.0 self.is_running = True async def send_event(self, event: Any) -> None: return None class FakeToolRouter: async def __aenter__(self) -> "FakeToolRouter": return self async def __aexit__(self, *exc: Any) -> bool: return False def _manager() -> SessionManager: manager = object.__new__(SessionManager) manager.config = SimpleNamespace(model_name="test-model") manager.sessions = {} manager._lock = asyncio.Lock() manager.persistence_store = RecordingStore() manager.messaging_gateway = SimpleNamespace() manager._pending_creates = 0 manager._reaper_task = None return manager def _make_agent_session( session_id: str, *, user_id: str = "owner", last_active_at: datetime | None = None, is_processing: bool = False, pending_approval: Any = None, ) -> AgentSession: session = FakeSession() session.pending_approval = pending_approval agent_session = AgentSession( session_id=session_id, session=session, # type: ignore[arg-type] tool_router=FakeToolRouter(), # type: ignore[arg-type] submission_queue=asyncio.Queue(), user_id=user_id, hf_token="token", ) agent_session.is_processing = is_processing if last_active_at is not None: agent_session.last_active_at = last_active_at return agent_session async def _start_real_run_session( manager: SessionManager, session_id: str, *, user_id: str = "owner", last_active_at: datetime | None = None, ) -> AgentSession: """Insert a session and start the REAL _run_session task for it.""" agent_session = _make_agent_session( session_id, user_id=user_id, last_active_at=last_active_at ) event_queue: asyncio.Queue = asyncio.Queue() await manager._start_agent_session( agent_session=agent_session, event_queue=event_queue, tool_router=agent_session.tool_router, ) await asyncio.sleep(0) # let the run loop reach its queue wait return agent_session def _install_fake_create(manager: SessionManager) -> asyncio.Event: """Replace blocking constructors + run loop with fakes for create tests.""" stop = asyncio.Event() def fake_create_session_sync(**kwargs: Any): return object(), FakeSession( hf_token=kwargs.get("hf_token"), user_plan=kwargs.get("user_plan"), ) async def fake_run_session(*_: Any) -> None: await stop.wait() manager._create_session_sync = fake_create_session_sync # type: ignore[method-assign] manager._run_session = fake_run_session # type: ignore[method-assign] manager._start_cpu_sandbox_preload = lambda _agent_session: None # type: ignore[method-assign] return stop async def _cancel_tasks(manager: SessionManager) -> None: tasks = [ agent_session.task for agent_session in manager.sessions.values() if agent_session.task and not agent_session.task.done() ] for task in tasks: task.cancel() if tasks: await asyncio.gather(*tasks, return_exceptions=True) # ── Reaper happy path ─────────────────────────────────────────────────── @pytest.mark.asyncio async def test_reaper_evicts_idle_session_as_resumable(): manager = _manager() cleaned: list[Any] = [] async def fake_cleanup(session: Any) -> None: cleaned.append(session) manager._cleanup_sandbox = fake_cleanup # type: ignore[method-assign] agent_session = await _start_real_run_session( manager, "stale", last_active_at=datetime.utcnow() - timedelta(hours=3), ) session = agent_session.session await manager._reap_idle_sessions() # Evicted from the live pool, sandbox torn down by the task's finally. assert "stale" not in manager.sessions assert cleaned == [session] # Persisted resumable (status="active", runtime_state="idle"), never "ended". store = manager.persistence_store snapshots = store.snapshots_for("stale") assert snapshots, "reaper should persist a snapshot" assert all(s["status"] == "active" for s in snapshots) assert snapshots[-1]["runtime_state"] == "idle" assert not any(s["status"] == "ended" for s in snapshots) # ── Spared sessions ───────────────────────────────────────────────────── @pytest.mark.asyncio @pytest.mark.parametrize( "kwargs", [ # Fresh: touched just now. {"last_active_at": None}, # Currently processing a turn. { "last_active_at": datetime.utcnow() - timedelta(hours=5), "is_processing": True, }, # Awaiting tool approval ("approve later", not idle). { "last_active_at": datetime.utcnow() - timedelta(hours=5), "pending_approval": {"tool_calls": [object()]}, }, # Dev sessions are never reaped. {"last_active_at": datetime.utcnow() - timedelta(hours=5), "user_id": "dev"}, ], ids=["fresh", "processing", "pending_approval", "dev"], ) async def test_reaper_spares(kwargs): manager = _manager() agent_session = _make_agent_session("spared", **kwargs) manager.sessions["spared"] = agent_session await manager._reap_idle_sessions() assert "spared" in manager.sessions assert agent_session.is_reaping is False # ── Submit / reap race ────────────────────────────────────────────────── @pytest.mark.asyncio async def test_reap_aborts_when_message_enqueued_first(): """A message enqueued before teardown makes the queue non-empty, so the reaper's re-check aborts — the message is never lost.""" manager = _manager() agent_session = _make_agent_session( "racing", last_active_at=datetime.utcnow() - timedelta(hours=3) ) manager.sessions["racing"] = agent_session agent_session.submission_queue.put_nowait(object()) cutoff = datetime.utcnow() - sm.REAPER_IDLE reaped = await manager._reap_one("racing", cutoff) assert reaped is False assert "racing" in manager.sessions assert agent_session.is_reaping is False assert agent_session.submission_queue.qsize() == 1 @pytest.mark.asyncio async def test_submit_rejected_while_reaping(): """submit() refuses a session being reaped instead of silently enqueuing onto a dying runtime (the caller then reloads a fresh one).""" manager = _manager() agent_session = _make_agent_session("reaping") agent_session.is_reaping = True manager.sessions["reaping"] = agent_session ok = await manager.submit("reaping", Operation(op_type=OpType.USER_INPUT, data={})) assert ok is False assert agent_session.submission_queue.empty() # ── Turn-finish timestamp ─────────────────────────────────────────────── @pytest.mark.asyncio async def test_turn_finish_restamps_activity(monkeypatch): """A turn that runs longer than the idle window must not be reaped the instant it finishes — the turn-finish stamp resets the idle clock.""" manager = _manager() async def fake_cleanup(session: Any) -> None: return None manager._cleanup_sandbox = fake_cleanup # type: ignore[method-assign] agent_session = await _start_real_run_session( manager, "longturn", last_active_at=datetime.utcnow() ) async def fake_process(session: Any, submission: Any) -> bool: # Simulate a turn that has been running far longer than the idle # window before it completes. agent_session.last_active_at = datetime.utcnow() - timedelta(hours=3) return True monkeypatch.setattr(sm, "process_submission", fake_process) agent_session.submission_queue.put_nowait( sm.Submission(id="s1", operation=Operation(op_type=OpType.USER_INPUT, data={})) ) # Wait for the turn-finish stamp to land. for _ in range(200): await asyncio.sleep(0.01) if datetime.utcnow() - agent_session.last_active_at < timedelta(minutes=1): break assert datetime.utcnow() - agent_session.last_active_at < timedelta(minutes=1) await _cancel_tasks(manager) # ── Global reservation race (Part B) ──────────────────────────────────── @pytest.mark.asyncio async def test_concurrent_creates_cannot_exceed_global_cap(monkeypatch): manager = _manager() stop = _install_fake_create(manager) monkeypatch.setattr(sm, "MAX_SESSIONS", 3) try: results = await asyncio.gather( *[manager.create_session(user_id="owner") for _ in range(10)], return_exceptions=True, ) created = [r for r in results if isinstance(r, str)] errors = [r for r in results if isinstance(r, SessionCapacityError)] assert len(created) == 3 assert len(errors) == 7 assert all(e.error_type == "global" for e in errors) assert len(manager.sessions) == 3 assert manager._pending_creates == 0 finally: stop.set() await _cancel_tasks(manager) @pytest.mark.asyncio async def test_failed_build_releases_reservation(monkeypatch): manager = _manager() _install_fake_create(manager) monkeypatch.setattr(sm, "MAX_SESSIONS", 5) def boom(**_: Any): raise RuntimeError("build failed") manager._create_session_sync = boom # type: ignore[method-assign] with pytest.raises(RuntimeError): await manager.create_session(user_id="owner") # The reservation must be released so a failed create can't shrink the pool. assert manager._pending_creates == 0 assert manager.sessions == {} # ── Per-user concurrent cap interacts with reclaimed slots (Part E) ────── @pytest.mark.asyncio async def test_per_user_cap_frees_up_after_slot_reclaimed(): manager = _manager() stop = _install_fake_create(manager) try: for i in range(sm.MAX_SESSIONS_PER_USER): manager.sessions[f"owner-{i}"] = _make_agent_session( f"owner-{i}", user_id="owner" ) # At the concurrent cap → rejected. with pytest.raises(SessionCapacityError) as exc: await manager.create_session(user_id="owner") assert exc.value.error_type == "per_user" message = str(exc.value) assert f"maximum of {sm.MAX_SESSIONS_PER_USER} live sessions" in message assert "Close an existing session" in message assert f"wait {sm.REAPER_IDLE_MINUTES:g} minutes" in message assert "after your last activity" in message assert "idle session to be released" in message # Reclaiming a slot (the reaper evicts an idle session) frees capacity. manager.sessions.pop("owner-0") new_id = await manager.create_session(user_id="owner") assert isinstance(new_id, str) assert manager._count_user_sessions("owner") == sm.MAX_SESSIONS_PER_USER finally: stop.set() await _cancel_tasks(manager) # ── Persistence safety (resumability invariant) ────────────────────────── @pytest.mark.asyncio async def test_reaper_skips_when_persistence_disabled(): """With no usable store a reaped session couldn't be restored, so eviction would destroy non-dev chats outright. The sweep must be a no-op.""" manager = _manager() manager.persistence_store = NoopSessionStore() # enabled = False agent_session = _make_agent_session( "idle", last_active_at=datetime.utcnow() - timedelta(hours=5) ) manager.sessions["idle"] = agent_session await manager._reap_idle_sessions() assert "idle" in manager.sessions assert agent_session.is_reaping is False @pytest.mark.asyncio async def test_reap_aborts_when_snapshot_write_fails(): """If the resumable snapshot can't be written (e.g. a transient Mongo error), abort rather than evict unrecoverable state — leave it live.""" manager = _manager() class FailingStore(RecordingStore): async def save_snapshot(self, **kwargs: Any) -> None: raise RuntimeError("mongo write failed") manager.persistence_store = FailingStore() agent_session = _make_agent_session( "idle", last_active_at=datetime.utcnow() - timedelta(hours=5) ) manager.sessions["idle"] = agent_session cutoff = datetime.utcnow() - sm.REAPER_IDLE reaped = await manager._reap_one("idle", cutoff) assert reaped is False assert "idle" in manager.sessions assert agent_session.is_reaping is False @pytest.mark.asyncio async def test_reap_aborts_when_message_write_fails_silently(): """The real store swallows message bulk_write errors for best-effort callers and only surfaces them when asked to be strict. The reaper must request strict mode, so a silent message-write failure still aborts the reap (not only metadata/connection failures that already make save_snapshot raise).""" manager = _manager() class StrictModeStore(RecordingStore): # Mirrors MongoSessionStore: message-write failure is swallowed unless # the caller passes raise_on_error (which the reaper does). async def save_snapshot( self, *, raise_on_error: bool = False, **kwargs: Any ) -> None: if raise_on_error: raise RuntimeError("message bulk_write failed") manager.persistence_store = StrictModeStore() agent_session = _make_agent_session( "idle", last_active_at=datetime.utcnow() - timedelta(hours=5) ) manager.sessions["idle"] = agent_session cutoff = datetime.utcnow() - sm.REAPER_IDLE reaped = await manager._reap_one("idle", cutoff) assert reaped is False assert "idle" in manager.sessions assert agent_session.is_reaping is False # ── Reap / restore race ────────────────────────────────────────────────── @pytest.mark.asyncio async def test_old_reaped_task_does_not_end_freshly_restored_session(): """If a user reopens after the reaper pops the old wrapper but before the old task finishes sandbox cleanup, the old finally must not mark the fresh wrapper ended.""" manager = _manager() cleanup_started = asyncio.Event() release_cleanup = asyncio.Event() async def slow_cleanup(session: Any) -> None: cleanup_started.set() await release_cleanup.wait() manager._cleanup_sandbox = slow_cleanup # type: ignore[method-assign] old = await _start_real_run_session( manager, "restore-race", last_active_at=datetime.utcnow() - timedelta(hours=3), ) cutoff = datetime.utcnow() - sm.REAPER_IDLE reap_task = asyncio.create_task(manager._reap_one("restore-race", cutoff)) for _ in range(100): await asyncio.sleep(0.01) if "restore-race" not in manager.sessions and cleanup_started.is_set(): break assert "restore-race" not in manager.sessions assert cleanup_started.is_set() fresh = _make_agent_session("restore-race") async with manager._lock: manager.sessions["restore-race"] = fresh release_cleanup.set() assert await reap_task is True if old.task is not None: assert old.task.done() assert manager.sessions["restore-race"] is fresh assert fresh.is_active is True assert all( snapshot.get("status") != "ended" for snapshot in manager.persistence_store.snapshots_for("restore-race") ) # ── Shutdown safety (cancellation must propagate) ──────────────────────── @pytest.mark.asyncio async def test_reaper_teardown_propagates_outer_cancellation(): """Cancelling the reaper while it awaits a slow teardown must propagate, so close() can't hang. Regression for the CancelledError-conflation bug: the old wait_for + bare-except swallowed the reaper's own cancellation.""" manager = _manager() release = asyncio.Event() async def slow_cleanup(session: Any) -> None: await release.wait() # block teardown so _reap_one parks in the wait manager._cleanup_sandbox = slow_cleanup # type: ignore[method-assign] agent_session = await _start_real_run_session( manager, "slow", last_active_at=datetime.utcnow() - timedelta(hours=3) ) cutoff = datetime.utcnow() - sm.REAPER_IDLE reap_task = asyncio.create_task(manager._reap_one("slow", cutoff)) # Let _reap_one persist + pop, then enter the teardown wait (the session # task is stuck in slow_cleanup, so the wait won't complete on its own). for _ in range(100): await asyncio.sleep(0.01) if "slow" not in manager.sessions: break await asyncio.sleep(0.02) # Simulate close() cancelling the reaper mid-teardown. reap_task.cancel() with pytest.raises(asyncio.CancelledError): await reap_task # Unblock the stuck teardown and reap the orphaned session task. release.set() if agent_session.task is not None: await asyncio.gather(agent_session.task, return_exceptions=True)