diff --git a/litellm/proxy/auth/cli_session_registry.py b/litellm/proxy/auth/cli_session_registry.py index f66d5432d7a..4c6cc6e534e 100644 --- a/litellm/proxy/auth/cli_session_registry.py +++ b/litellm/proxy/auth/cli_session_registry.py @@ -6,7 +6,9 @@ it by decrypting it. This registry is the server-side record that makes a sessio listable and revocable: one row per login, keyed by the sha256 of the session token. A session with no row predates the registry and still authenticates until it expires; -only a row with ``revoked_at`` set is refused. +only a row with ``revoked_at`` set is refused. Expired rows are dead weight, so every +registration drops them first and the table stays bounded by the number of live +sessions rather than by every login ever made. """ from __future__ import annotations @@ -49,6 +51,8 @@ class _CLISessionTable(Protocol): async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> _CLISessionRecord | None: ... + async def delete_many(self, where: Mapping[str, object]) -> int: ... + def _cli_session_table(prisma_client: PrismaClient) -> _CLISessionTable: table: Final[_CLISessionTable] = CLISessionRepository(prisma_client).table @@ -71,9 +75,14 @@ async def record_cli_session( user_id: str, team_id: str | None, ) -> CLISessionResponse: - """Register a freshly minted session. Raises if the row cannot be written, so a - session that could never be revoked is never handed to the CLI.""" - created: Final = await _cli_session_table(prisma_client).create( + """Register a freshly minted session, first dropping rows that can no longer be + revoked. Raises if the row cannot be written, so a session that could never be + revoked is never handed to the CLI.""" + table: Final = _cli_session_table(prisma_client) + await table.delete_many( + where={"expires_at": {"lte": get_utc_datetime()}} # mutable-ok: prisma query filters are dict-shaped + ) + created: Final = await table.create( data={ # mutable-ok: prisma payloads are plain dicts "session_id": cli_session_id(session_token), "user_id": user_id, diff --git a/tests/test_litellm/proxy/auth/test_cli_session_registry.py b/tests/test_litellm/proxy/auth/test_cli_session_registry.py index 844fd0cdbde..d7ba4012dd7 100644 --- a/tests/test_litellm/proxy/auth/test_cli_session_registry.py +++ b/tests/test_litellm/proxy/auth/test_cli_session_registry.py @@ -71,6 +71,13 @@ class FakeCLISessionTable: gt = where["expires_at"]["gt"] return len([r for r in self.rows.values() if r["expires_at"] > gt]) + async def delete_many(self, where): + lte = where["expires_at"]["lte"] + doomed = [sid for sid, r in self.rows.items() if r["expires_at"] <= lte] + for sid in doomed: + del self.rows[sid] + return len(doomed) + class FakeDB: def __init__(self, table): @@ -321,6 +328,29 @@ async def test_db_outage_follows_the_proxy_wide_posture(monkeypatch): await is_cli_session_revoked(session_token=SESSION_TOKEN, prisma_client=prisma, user_api_key_cache=_cache()) +@pytest.mark.asyncio +async def test_recording_drops_expired_sessions_but_keeps_live_ones(): + """An authenticated user can log in over and over. Without a retention path every + login would leave a row behind forever, so registration must clear the rows that + can no longer be revoked while leaving every live session listable.""" + table = FakeCLISessionTable( + { + "dead-1": _session_row(session_id="dead-1", expires_in_hours=-5), + "dead-2": _session_row(session_id="dead-2", expires_in_hours=-1), + "live": _session_row(session_id="live", expires_in_hours=1), + } + ) + + await record_cli_session( + prisma_client=FakePrismaClient(table), + session_token=SESSION_TOKEN, + user_id="u-1", + team_id="t-1", + ) + + assert sorted(table.rows) == sorted(["live", hash_token(SESSION_TOKEN)]) + + @pytest.mark.asyncio async def test_listing_hides_expired_sessions(): table = FakeCLISessionTable( diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index d1832abe0f3..22f1c681fb4 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -2528,6 +2528,7 @@ def _cli_session_registry_prisma() -> MagicMock: the credential out, so the poll needs a DB whose CLI-session table accepts the write.""" prisma = MagicMock() now = datetime.now(timezone.utc) + prisma.db.litellm_clisessiontable.delete_many = AsyncMock(return_value=0) prisma.db.litellm_clisessiontable.create = AsyncMock( return_value=SimpleNamespace( model_dump=lambda: {