From f6e5951556f39a1be69e87d2d862961363fbb1ca Mon Sep 17 00:00:00 2001
From: Ishaan Jaffer <ishaanjaffer0324@gmail.com>
Date: Wed, 6 May 2026 15:10:37 -0700
Subject: [PATCH] test(agent_session_endpoints): validate cleanup sweeper for
 expiry/dead-daemon/stuck-runs (LIT-2877 #13)

---
 .../test_cleanup_sweeper.py                   | 121 ++++++++++++++++++
 1 file changed, 121 insertions(+)
 create mode 100644 tests/test_litellm/proxy/agent_session_endpoints/test_cleanup_sweeper.py

diff --git a/tests/test_litellm/proxy/agent_session_endpoints/test_cleanup_sweeper.py b/tests/test_litellm/proxy/agent_session_endpoints/test_cleanup_sweeper.py
new file mode 100644
index 00000000000..b1f8c451abb
--- /dev/null
+++ b/tests/test_litellm/proxy/agent_session_endpoints/test_cleanup_sweeper.py
@@ -0,0 +1,121 @@
+"""
+Validation #13 — cleanup sweeper.
+
+Drives:
+  * Force ``expires_at`` past + status=ready → sweeper marks terminated, calls provider.terminate
+  * Force ``last_heartbeat_at`` > 90s ago → sweeper marks error
+  * Force ``status=running, updated_at`` > idle timeout → sweeper marks run error
+"""
+
+from datetime import datetime, timedelta, timezone
+
+import pytest
+
+from litellm.proxy.agent_session_endpoints.cleanup import run_cleanup_pass
+from litellm.proxy.agent_session_endpoints.constants import (
+    DAEMON_HEARTBEAT_DEAD_AFTER_SECONDS,
+    RUN_IDLE_TIMEOUT_SECONDS,
+    RUN_STATUS_ERROR,
+    SESSION_STATUS_ERROR,
+    SESSION_STATUS_READY,
+    SESSION_STATUS_TERMINATED,
+)
+
+
+@pytest.mark.asyncio
+async def test_sweeper_terminates_expired_sessions(
+    client, noop_provider, fake_prisma_client
+):
+    a = client.post(
+        "/v2/agents",
+        headers={"Authorization": "Bearer k"},
+        json={"name": "t", "model": "gpt-4"},
+    ).json()
+    sess = client.post(
+        "/v2/sessions",
+        headers={"Authorization": "Bearer k"},
+        json={"agent_id": a["id"], "repos": []},
+    ).json()
+    sid = sess["id"]
+
+    # Force expiry into the past.
+    row = fake_prisma_client.db.litellm_agentsession.rows[0]
+    row.expires_at = datetime.now(timezone.utc) - timedelta(minutes=1)
+    row.status = SESSION_STATUS_READY
+
+    summary = await run_cleanup_pass(fake_prisma_client)
+    assert summary["expired_sessions"] == 1
+    assert row.status == SESSION_STATUS_TERMINATED
+
+    # provider.terminate was called.
+    assert any(c["session_id"] == sid for c in noop_provider.terminate_calls)
+
+
+@pytest.mark.asyncio
+async def test_sweeper_marks_dead_daemon_sessions_error(
+    client, noop_provider, fake_prisma_client
+):
+    a = client.post(
+        "/v2/agents",
+        headers={"Authorization": "Bearer k"},
+        json={"name": "t", "model": "gpt-4"},
+    ).json()
+    sess = client.post(
+        "/v2/sessions",
+        headers={"Authorization": "Bearer k"},
+        json={"agent_id": a["id"], "repos": []},
+    ).json()
+
+    row = fake_prisma_client.db.litellm_agentsession.rows[0]
+    row.status = SESSION_STATUS_READY
+    row.last_heartbeat_at = datetime.now(timezone.utc) - timedelta(
+        seconds=DAEMON_HEARTBEAT_DEAD_AFTER_SECONDS + 30
+    )
+
+    summary = await run_cleanup_pass(fake_prisma_client)
+    assert summary["dead_daemon_sessions"] == 1
+    assert row.status == SESSION_STATUS_ERROR
+
+
+@pytest.mark.asyncio
+async def test_sweeper_marks_stuck_runs_error(
+    client, noop_provider, fake_prisma_client
+):
+    a = client.post(
+        "/v2/agents",
+        headers={"Authorization": "Bearer k"},
+        json={"name": "t", "model": "gpt-4"},
+    ).json()
+    sess = client.post(
+        "/v2/sessions",
+        headers={"Authorization": "Bearer k"},
+        json={"agent_id": a["id"], "repos": []},
+    ).json()
+    run = client.post(
+        f"/v2/sessions/{sess['id']}/runs",
+        headers={"Authorization": "Bearer k"},
+        json={"prompt": {"text": "x"}},
+    ).json()
+
+    run_row = fake_prisma_client.db.litellm_agentrun.rows[0]
+    run_row.status = "running"
+    run_row.updated_at = datetime.now(timezone.utc) - timedelta(
+        seconds=RUN_IDLE_TIMEOUT_SECONDS + 60
+    )
+
+    summary = await run_cleanup_pass(fake_prisma_client)
+    assert summary["stuck_runs"] == 1
+    assert run_row.status == RUN_STATUS_ERROR
+
+
+@pytest.mark.asyncio
+async def test_sweeper_pass_with_nothing_to_do(
+    client, noop_provider, fake_prisma_client
+):
+    """Empty-DB pass returns all-zero summary, no exception."""
+    summary = await run_cleanup_pass(fake_prisma_client)
+    assert summary == {
+        "expired_sessions": 0,
+        "dead_daemon_sessions": 0,
+        "stuck_runs": 0,
+    }