mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
Add expired UI session key cleanup job
Made-with: Cursor
This commit is contained in:
Milan
parent
082a8faf46
commit
21cf42f568
4 changed files with 393 additions and 2 deletions
|
|
@ -1396,6 +1396,15 @@ LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS = int(
|
|||
os.getenv("LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS", 600)
|
||||
) # 10 minutes default — caps the deadlock window if a pod crashes mid-rotation
|
||||
UI_SESSION_TOKEN_TEAM_ID = "litellm-dashboard"
|
||||
LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_ENABLED = os.getenv(
|
||||
"LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_ENABLED", "false"
|
||||
)
|
||||
LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_INTERVAL_SECONDS = int(
|
||||
os.getenv("LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_INTERVAL_SECONDS", 86400)
|
||||
) # 24 hours default
|
||||
LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE = int(
|
||||
os.getenv("LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE", 1000)
|
||||
)
|
||||
LITELLM_PROXY_ADMIN_NAME = "default_user_id"
|
||||
|
||||
########################### CLI SSO AUTHENTICATION CONSTANTS ###########################
|
||||
|
|
@ -1425,6 +1434,7 @@ CLOUDZERO_MAX_FETCHED_DATA_RECORDS = int(
|
|||
)
|
||||
SPEND_LOG_CLEANUP_JOB_NAME = "spend_log_cleanup"
|
||||
KEY_ROTATION_JOB_NAME = "litellm_key_rotation_job"
|
||||
EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME = "litellm_expired_ui_session_key_cleanup_job"
|
||||
SPEND_LOG_RUN_LOOPS = int(os.getenv("SPEND_LOG_RUN_LOOPS", 500))
|
||||
SPEND_LOG_CLEANUP_BATCH_SIZE = int(os.getenv("SPEND_LOG_CLEANUP_BATCH_SIZE", 1000))
|
||||
SPEND_LOG_QUEUE_SIZE_THRESHOLD = int(os.getenv("SPEND_LOG_QUEUE_SIZE_THRESHOLD", 100))
|
||||
|
|
|
|||
|
|
@ -0,0 +1,116 @@
|
|||
"""
|
||||
Expired UI session key cleanup manager.
|
||||
|
||||
Deletes expired virtual keys created for LiteLLM dashboard sessions.
|
||||
"""
|
||||
|
||||
from datetime import datetime, timezone
|
||||
from typing import List
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.caching import DualCache
|
||||
from litellm.constants import (
|
||||
EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME,
|
||||
LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE,
|
||||
LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME,
|
||||
UI_SESSION_TOKEN_TEAM_ID,
|
||||
)
|
||||
from litellm.proxy._types import KeyRequest, LiteLLM_VerificationToken, UserAPIKeyAuth
|
||||
from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
delete_verification_tokens,
|
||||
)
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
|
||||
class ExpiredUISessionKeyCleanupManager:
|
||||
"""
|
||||
Cleans up expired UI session keys.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
prisma_client: PrismaClient,
|
||||
user_api_key_cache: DualCache,
|
||||
pod_lock_manager=None,
|
||||
):
|
||||
self.prisma_client = prisma_client
|
||||
self.user_api_key_cache = user_api_key_cache
|
||||
self.pod_lock_manager = pod_lock_manager
|
||||
|
||||
async def cleanup_expired_keys(self) -> int:
|
||||
"""
|
||||
Main entry point for deleting expired UI session keys.
|
||||
Uses PodLockManager to ensure only one pod runs cleanup in multi-pod deployments.
|
||||
"""
|
||||
lock_acquired = False
|
||||
try:
|
||||
if self.pod_lock_manager and self.pod_lock_manager.redis_cache:
|
||||
lock_acquired = (
|
||||
await self.pod_lock_manager.acquire_lock(
|
||||
cronjob_id=EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME,
|
||||
)
|
||||
or False
|
||||
)
|
||||
if not lock_acquired:
|
||||
verbose_proxy_logger.warning(
|
||||
"Expired UI session key cleanup: another pod is already "
|
||||
"running cleanup or Redis lock acquisition failed - "
|
||||
"skipping this cycle."
|
||||
)
|
||||
return 0
|
||||
|
||||
verbose_proxy_logger.info("Starting expired UI session key cleanup...")
|
||||
|
||||
expired_keys = await self._find_expired_ui_session_keys()
|
||||
if not expired_keys:
|
||||
verbose_proxy_logger.debug("No expired UI session keys found")
|
||||
return 0
|
||||
|
||||
tokens = [key.token for key in expired_keys if key.token is not None]
|
||||
if not tokens:
|
||||
return 0
|
||||
|
||||
system_user = UserAPIKeyAuth.get_litellm_internal_jobs_user_api_key_auth()
|
||||
response, keys_being_deleted = await delete_verification_tokens(
|
||||
tokens=tokens,
|
||||
user_api_key_cache=self.user_api_key_cache,
|
||||
user_api_key_dict=system_user,
|
||||
litellm_changed_by=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME,
|
||||
)
|
||||
await KeyManagementEventHooks.async_key_deleted_hook(
|
||||
data=KeyRequest(keys=tokens),
|
||||
keys_being_deleted=keys_being_deleted,
|
||||
response=response or {},
|
||||
user_api_key_dict=system_user,
|
||||
litellm_changed_by=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME,
|
||||
)
|
||||
verbose_proxy_logger.info(
|
||||
"Deleted %s expired UI session key(s)", len(tokens)
|
||||
)
|
||||
return len(tokens)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(f"Expired UI session key cleanup failed: {e}")
|
||||
return 0
|
||||
finally:
|
||||
if (
|
||||
lock_acquired
|
||||
and self.pod_lock_manager
|
||||
and self.pod_lock_manager.redis_cache
|
||||
):
|
||||
await self.pod_lock_manager.release_lock(
|
||||
cronjob_id=EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME,
|
||||
)
|
||||
|
||||
async def _find_expired_ui_session_keys(self) -> List[LiteLLM_VerificationToken]:
|
||||
"""
|
||||
Find expired LiteLLM dashboard session keys.
|
||||
"""
|
||||
now = datetime.now(timezone.utc)
|
||||
return await self.prisma_client.db.litellm_verificationtoken.find_many(
|
||||
where={
|
||||
"team_id": UI_SESSION_TOKEN_TEAM_ID,
|
||||
"expires": {"lt": now},
|
||||
},
|
||||
take=LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE,
|
||||
)
|
||||
|
|
@ -6737,6 +6737,10 @@ class ProxyStartupEvent:
|
|||
Args:
|
||||
scheduler: The scheduler to add the background jobs to
|
||||
"""
|
||||
global prisma_client
|
||||
global proxy_logging_obj
|
||||
global user_api_key_cache
|
||||
|
||||
########################################################
|
||||
# CloudZero Background Job
|
||||
########################################################
|
||||
|
|
@ -6810,8 +6814,6 @@ class ProxyStartupEvent:
|
|||
)
|
||||
|
||||
# Get prisma_client and proxy_logging_obj from global scope
|
||||
global prisma_client
|
||||
global proxy_logging_obj
|
||||
if prisma_client is not None:
|
||||
# Reuse the PodLockManager from db_spend_update_writer
|
||||
pod_lock_manager = (
|
||||
|
|
@ -6841,6 +6843,83 @@ class ProxyStartupEvent:
|
|||
"Key rotation disabled (set LITELLM_KEY_ROTATION_ENABLED=true to enable)"
|
||||
)
|
||||
|
||||
await cls._initialize_expired_ui_session_key_cleanup_background_job(
|
||||
scheduler=scheduler
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def _initialize_expired_ui_session_key_cleanup_background_job(
|
||||
cls, scheduler: AsyncIOScheduler
|
||||
):
|
||||
"""
|
||||
Initialize the expired UI session key cleanup background job.
|
||||
"""
|
||||
global prisma_client
|
||||
global proxy_logging_obj
|
||||
global user_api_key_cache
|
||||
|
||||
########################################################
|
||||
# Expired UI Session Key Cleanup Background Job
|
||||
########################################################
|
||||
from litellm.constants import (
|
||||
EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME,
|
||||
LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_ENABLED,
|
||||
LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_INTERVAL_SECONDS,
|
||||
)
|
||||
|
||||
expired_ui_session_key_cleanup_enabled: Optional[bool] = str_to_bool(
|
||||
LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_ENABLED
|
||||
)
|
||||
verbose_proxy_logger.debug(
|
||||
"expired_ui_session_key_cleanup_enabled: "
|
||||
f"{expired_ui_session_key_cleanup_enabled}"
|
||||
)
|
||||
|
||||
if expired_ui_session_key_cleanup_enabled is True:
|
||||
try:
|
||||
from litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager import (
|
||||
ExpiredUISessionKeyCleanupManager,
|
||||
)
|
||||
|
||||
if prisma_client is not None:
|
||||
pod_lock_manager = (
|
||||
proxy_logging_obj.db_spend_update_writer.pod_lock_manager
|
||||
)
|
||||
expired_ui_session_key_cleanup_manager = (
|
||||
ExpiredUISessionKeyCleanupManager(
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
pod_lock_manager=pod_lock_manager,
|
||||
)
|
||||
)
|
||||
verbose_proxy_logger.debug(
|
||||
"Expired UI session key cleanup background job scheduled "
|
||||
"every "
|
||||
f"{LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_INTERVAL_SECONDS} "
|
||||
"seconds "
|
||||
"(LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_ENABLED=true)"
|
||||
)
|
||||
scheduler.add_job(
|
||||
expired_ui_session_key_cleanup_manager.cleanup_expired_keys,
|
||||
"interval",
|
||||
seconds=LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_INTERVAL_SECONDS,
|
||||
id=EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME,
|
||||
)
|
||||
else:
|
||||
verbose_proxy_logger.warning(
|
||||
"Expired UI session key cleanup enabled but prisma_client "
|
||||
"not available"
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.warning(
|
||||
f"Failed to setup expired UI session key cleanup job: {e}"
|
||||
)
|
||||
else:
|
||||
verbose_proxy_logger.debug(
|
||||
"Expired UI session key cleanup disabled (set "
|
||||
"LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_ENABLED=true to enable)"
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def _initialize_slack_alerting_jobs(
|
||||
cls,
|
||||
|
|
|
|||
|
|
@ -0,0 +1,186 @@
|
|||
"""
|
||||
Test expired UI session key cleanup manager functionality.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../../.."))
|
||||
|
||||
from litellm.constants import (
|
||||
EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME,
|
||||
LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE,
|
||||
LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME,
|
||||
UI_SESSION_TOKEN_TEAM_ID,
|
||||
)
|
||||
from litellm.proxy._types import LiteLLM_VerificationToken
|
||||
from litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager import (
|
||||
ExpiredUISessionKeyCleanupManager,
|
||||
)
|
||||
|
||||
|
||||
class TestExpiredUISessionKeyCleanupManager:
|
||||
"""Test the ExpiredUISessionKeyCleanupManager class functionality."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_find_expired_ui_session_keys_filters_dashboard_team_and_expiry(self):
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_cache = MagicMock()
|
||||
manager = ExpiredUISessionKeyCleanupManager(
|
||||
prisma_client=mock_prisma_client,
|
||||
user_api_key_cache=mock_cache,
|
||||
)
|
||||
|
||||
now = datetime(2026, 4, 25, 12, 0, 0, tzinfo=timezone.utc)
|
||||
mock_keys = [
|
||||
LiteLLM_VerificationToken(
|
||||
token="expired-dashboard-token",
|
||||
team_id=UI_SESSION_TOKEN_TEAM_ID,
|
||||
expires=now - timedelta(seconds=1),
|
||||
)
|
||||
]
|
||||
mock_prisma_client.db.litellm_verificationtoken.find_many.return_value = (
|
||||
mock_keys
|
||||
)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.datetime"
|
||||
) as mock_datetime:
|
||||
mock_datetime.now.return_value = now
|
||||
mock_datetime.side_effect = lambda *args, **kwargs: datetime(
|
||||
*args, **kwargs
|
||||
)
|
||||
|
||||
keys = await manager._find_expired_ui_session_keys()
|
||||
|
||||
mock_prisma_client.db.litellm_verificationtoken.find_many.assert_called_once_with(
|
||||
where={
|
||||
"team_id": UI_SESSION_TOKEN_TEAM_ID,
|
||||
"expires": {"lt": now},
|
||||
},
|
||||
take=LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE,
|
||||
)
|
||||
assert keys == mock_keys
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cleanup_expired_keys_uses_existing_delete_path(self):
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_cache = MagicMock()
|
||||
manager = ExpiredUISessionKeyCleanupManager(
|
||||
prisma_client=mock_prisma_client,
|
||||
user_api_key_cache=mock_cache,
|
||||
)
|
||||
expired_key = LiteLLM_VerificationToken(
|
||||
token="expired-dashboard-token",
|
||||
team_id=UI_SESSION_TOKEN_TEAM_ID,
|
||||
expires=datetime.now(timezone.utc) - timedelta(seconds=1),
|
||||
)
|
||||
manager._find_expired_ui_session_keys = AsyncMock(return_value=[expired_key])
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.delete_verification_tokens",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_delete_verification_tokens:
|
||||
mock_delete_verification_tokens.return_value = (
|
||||
{"deleted_keys": ["expired-dashboard-token"], "failed_tokens": []},
|
||||
[expired_key],
|
||||
)
|
||||
with patch(
|
||||
"litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.KeyManagementEventHooks.async_key_deleted_hook",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_key_deleted_hook:
|
||||
deleted_count = await manager.cleanup_expired_keys()
|
||||
|
||||
assert deleted_count == 1
|
||||
mock_delete_verification_tokens.assert_called_once()
|
||||
call_kwargs = mock_delete_verification_tokens.call_args.kwargs
|
||||
assert call_kwargs["tokens"] == ["expired-dashboard-token"]
|
||||
assert call_kwargs["user_api_key_cache"] == mock_cache
|
||||
assert (
|
||||
call_kwargs["litellm_changed_by"]
|
||||
== LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME
|
||||
)
|
||||
assert call_kwargs["user_api_key_dict"].user_id == "system"
|
||||
mock_key_deleted_hook.assert_called_once()
|
||||
hook_kwargs = mock_key_deleted_hook.call_args.kwargs
|
||||
assert hook_kwargs["data"].keys == ["expired-dashboard-token"]
|
||||
assert hook_kwargs["keys_being_deleted"] == [expired_key]
|
||||
assert hook_kwargs["response"] == {
|
||||
"deleted_keys": ["expired-dashboard-token"],
|
||||
"failed_tokens": [],
|
||||
}
|
||||
assert (
|
||||
hook_kwargs["litellm_changed_by"]
|
||||
== LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cleanup_expired_keys_noops_when_no_keys_found(self):
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_cache = MagicMock()
|
||||
manager = ExpiredUISessionKeyCleanupManager(
|
||||
prisma_client=mock_prisma_client,
|
||||
user_api_key_cache=mock_cache,
|
||||
)
|
||||
manager._find_expired_ui_session_keys = AsyncMock(return_value=[])
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.delete_verification_tokens",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_delete_verification_tokens:
|
||||
deleted_count = await manager.cleanup_expired_keys()
|
||||
|
||||
assert deleted_count == 0
|
||||
mock_delete_verification_tokens.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cleanup_expired_keys_skips_when_lock_held(self):
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_cache = MagicMock()
|
||||
mock_pod_lock_manager = MagicMock()
|
||||
mock_pod_lock_manager.redis_cache = MagicMock()
|
||||
mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=False)
|
||||
mock_pod_lock_manager.release_lock = AsyncMock()
|
||||
|
||||
manager = ExpiredUISessionKeyCleanupManager(
|
||||
prisma_client=mock_prisma_client,
|
||||
user_api_key_cache=mock_cache,
|
||||
pod_lock_manager=mock_pod_lock_manager,
|
||||
)
|
||||
manager._find_expired_ui_session_keys = AsyncMock()
|
||||
|
||||
deleted_count = await manager.cleanup_expired_keys()
|
||||
|
||||
assert deleted_count == 0
|
||||
mock_pod_lock_manager.acquire_lock.assert_called_once_with(
|
||||
cronjob_id=EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME,
|
||||
)
|
||||
manager._find_expired_ui_session_keys.assert_not_called()
|
||||
mock_pod_lock_manager.release_lock.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cleanup_expired_keys_releases_acquired_lock(self):
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_cache = MagicMock()
|
||||
mock_pod_lock_manager = MagicMock()
|
||||
mock_pod_lock_manager.redis_cache = MagicMock()
|
||||
mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True)
|
||||
mock_pod_lock_manager.release_lock = AsyncMock()
|
||||
|
||||
manager = ExpiredUISessionKeyCleanupManager(
|
||||
prisma_client=mock_prisma_client,
|
||||
user_api_key_cache=mock_cache,
|
||||
pod_lock_manager=mock_pod_lock_manager,
|
||||
)
|
||||
manager._find_expired_ui_session_keys = AsyncMock(return_value=[])
|
||||
|
||||
deleted_count = await manager.cleanup_expired_keys()
|
||||
|
||||
assert deleted_count == 0
|
||||
mock_pod_lock_manager.release_lock.assert_called_once_with(
|
||||
cronjob_id=EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME,
|
||||
)
|
||||
Loading…
Add table
Reference in a new issue