From 68d11f445f7cd6a8fcd88a73fe5a330546d433e2 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Wed, 6 May 2026 15:03:24 -0700 Subject: [PATCH] feat(agent_session_endpoints): add cleanup sweeper for expired sessions, dead daemons, stuck runs --- .../proxy/agent_session_endpoints/cleanup.py | 249 ++++++++++++++++++ 1 file changed, 249 insertions(+) create mode 100644 litellm/proxy/agent_session_endpoints/cleanup.py diff --git a/litellm/proxy/agent_session_endpoints/cleanup.py b/litellm/proxy/agent_session_endpoints/cleanup.py new file mode 100644 index 00000000000..f1f2ebf93c4 --- /dev/null +++ b/litellm/proxy/agent_session_endpoints/cleanup.py @@ -0,0 +1,249 @@ +""" +Background sweeper that handles three failure modes: + + 1. Expired sessions — past ``expires_at``, still non-terminal: + run the standard terminate flow. + 2. Dead daemons — non-terminal session whose ``last_heartbeat_at`` is + older than ``DAEMON_HEARTBEAT_DEAD_AFTER_SECONDS``: mark error. + 3. Stuck runs — running runs that haven't been touched within + ``RUN_IDLE_TIMEOUT_SECONDS``: mark error. + +The sweeper is started from ``proxy_server.py``'s startup hooks. It loops +until cancelled, sleeping `CLEANUP_SWEEPER_INTERVAL_SECONDS` between +passes. +""" + +import asyncio +from datetime import datetime, timedelta, timezone +from typing import Optional + +from litellm._logging import verbose_proxy_logger +from litellm.proxy.agent_session_endpoints.constants import ( + CLEANUP_SWEEPER_INTERVAL_SECONDS, + DAEMON_HEARTBEAT_DEAD_AFTER_SECONDS, + EVENT_TYPE_RUN_ERROR, + RUN_ACTIVE_STATUSES, + RUN_IDLE_TIMEOUT_SECONDS, + RUN_STATUS_ERROR, + SESSION_STATUS_ERROR, + SESSION_STATUS_PROVISIONING, + SESSION_STATUS_TERMINATED, + SESSION_TERMINAL_STATUSES, +) + + +def _now() -> datetime: + return datetime.now(timezone.utc) + + +async def _sweep_expired_sessions(prisma_client) -> int: + """Terminate sessions whose ``expires_at`` has passed and are not + already terminal. Returns count terminated.""" + from litellm.proxy.agent_session_endpoints.session_endpoints import ( + _terminate_session_internal, + ) + + rows = await prisma_client.db.litellm_agentsession.find_many( + where={ + "expires_at": {"lt": _now()}, + "status": { + "notIn": list(SESSION_TERMINAL_STATUSES), + }, + }, + take=200, + ) + for row in rows: + try: + await _terminate_session_internal(row.id, reason="expired") + except Exception as exc: + verbose_proxy_logger.exception( + "sweeper: failed to terminate expired session=%s: %s", row.id, exc + ) + return len(rows) + + +async def _sweep_dead_daemons(prisma_client) -> int: + """Mark sessions ``error`` whose daemon hasn't heartbeat in + ``DAEMON_HEARTBEAT_DEAD_AFTER_SECONDS``. + + Skips ``provisioning`` sessions (they haven't registered yet — their + own timeout is governed by `expires_at`). + """ + threshold = _now() - timedelta(seconds=DAEMON_HEARTBEAT_DEAD_AFTER_SECONDS) + rows = await prisma_client.db.litellm_agentsession.find_many( + where={ + "last_heartbeat_at": {"lt": threshold}, + "status": { + "notIn": list(SESSION_TERMINAL_STATUSES) + + [SESSION_STATUS_PROVISIONING], + }, + }, + take=200, + ) + if not rows: + return 0 + + now = _now() + ids = [r.id for r in rows] + await prisma_client.db.litellm_agentsession.update_many( + where={"id": {"in": ids}}, + data={ + "status": SESSION_STATUS_ERROR, + "terminated_at": now, + "updated_at": now, + }, + ) + + # Also flip any active runs in those sessions to error. + runs = await prisma_client.db.litellm_agentrun.find_many( + where={ + "session_id": {"in": ids}, + "status": {"in": list(RUN_ACTIVE_STATUSES)}, + }, + take=500, + ) + for run in runs: + await prisma_client.db.litellm_agentrun.update( + where={"id": run.id}, + data={ + "status": RUN_STATUS_ERROR, + "terminated_at": now, + "updated_at": now, + }, + ) + last_evt = await prisma_client.db.litellm_agentrunevent.find_first( + where={"run_id": run.id}, order={"seq": "desc"} + ) + next_seq = (last_evt.seq + 1) if last_evt else 1 + try: + await prisma_client.db.litellm_agentrunevent.create( + data={ + "run_id": run.id, + "seq": next_seq, + "event_type": EVENT_TYPE_RUN_ERROR, + "payload": {"reason": "daemon_dead"}, + } + ) + except Exception as exc: + verbose_proxy_logger.warning( + "sweeper: skipped run_error event run=%s seq=%s: %s", + run.id, + next_seq, + exc, + ) + + verbose_proxy_logger.info( + "sweeper: marked %d sessions error (dead daemon)", len(ids) + ) + return len(ids) + + +async def _sweep_stuck_runs(prisma_client) -> int: + """Mark runs ``error`` if they've been ``running`` past the idle timeout. + + Sweeper-driven only — clients may legitimately have long-running runs + so the threshold is generous (``RUN_IDLE_TIMEOUT_SECONDS``, 30 min by + default). + """ + threshold = _now() - timedelta(seconds=RUN_IDLE_TIMEOUT_SECONDS) + rows = await prisma_client.db.litellm_agentrun.find_many( + where={ + "status": "running", + "updated_at": {"lt": threshold}, + }, + take=200, + ) + if not rows: + return 0 + + now = _now() + for run in rows: + await prisma_client.db.litellm_agentrun.update( + where={"id": run.id}, + data={ + "status": RUN_STATUS_ERROR, + "terminated_at": now, + "updated_at": now, + }, + ) + last_evt = await prisma_client.db.litellm_agentrunevent.find_first( + where={"run_id": run.id}, order={"seq": "desc"} + ) + next_seq = (last_evt.seq + 1) if last_evt else 1 + try: + await prisma_client.db.litellm_agentrunevent.create( + data={ + "run_id": run.id, + "seq": next_seq, + "event_type": EVENT_TYPE_RUN_ERROR, + "payload": {"reason": "run_idle_timeout"}, + } + ) + except Exception as exc: + verbose_proxy_logger.warning( + "sweeper: skipped run_error event run=%s seq=%s: %s", + run.id, + next_seq, + exc, + ) + return len(rows) + + +async def run_cleanup_pass(prisma_client) -> dict: + """Single pass — exposed for tests and on-demand triggers.""" + expired = await _sweep_expired_sessions(prisma_client) + dead = await _sweep_dead_daemons(prisma_client) + stuck = await _sweep_stuck_runs(prisma_client) + return { + "expired_sessions": expired, + "dead_daemon_sessions": dead, + "stuck_runs": stuck, + } + + +_sweeper_task: Optional[asyncio.Task] = None + + +async def _cleanup_loop() -> None: + """Long-running loop. Cancellation-safe — bails out cleanly on + ``asyncio.CancelledError``.""" + while True: + try: + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is not None: + summary = await run_cleanup_pass(prisma_client) + if any(summary.values()): + verbose_proxy_logger.info( + "agent_session_cleanup_sweeper: %s", summary + ) + except asyncio.CancelledError: + raise + except Exception as exc: + verbose_proxy_logger.exception( + "agent_session_cleanup_sweeper iteration failed: %s", exc + ) + try: + await asyncio.sleep(CLEANUP_SWEEPER_INTERVAL_SECONDS) + except asyncio.CancelledError: + raise + + +def start_cleanup_sweeper() -> None: + """Idempotent: spawn the sweeper task once. Subsequent calls are no-ops. + + Called from ``proxy_server.startup_event`` so it ships with the proxy. + """ + global _sweeper_task + if _sweeper_task is not None and not _sweeper_task.done(): + return + _sweeper_task = asyncio.create_task(_cleanup_loop()) + verbose_proxy_logger.info("agent_session_cleanup_sweeper started") + + +def stop_cleanup_sweeper() -> None: + """Cancel the sweeper task. Safe to call multiple times.""" + global _sweeper_task + if _sweeper_task is not None and not _sweeper_task.done(): + _sweeper_task.cancel() + _sweeper_task = None