From 83debc98e2b28e68eda2685100a67e25fde62021 Mon Sep 17 00:00:00 2001 From: Mubashir Osmani Date: Sat, 25 Jul 2026 05:58:59 +0000 Subject: [PATCH] fix(proxy): deliver key rotation emails and warn on unhonored rotation schedules Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../send_emails/base_email.py | 4 + .../common_utils/key_rotation_manager.py | 69 ++++++++-- litellm/proxy/proxy_server.py | 24 ++++ litellm/proxy/utils.py | 24 ++++ .../send_emails/test_base_email.py | 18 +++ .../common_utils/test_key_rotation_e2e.py | 17 +-- .../test_key_rotation_integration.py | 4 +- .../common_utils/test_key_rotation_manager.py | 125 ++++++++++++++++-- tests/test_litellm/proxy/test_proxy_utils.py | 60 +++++++++ 9 files changed, 309 insertions(+), 36 deletions(-) diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py b/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py index e7898cac565..0db9ee2defd 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py @@ -812,6 +812,10 @@ class BaseEmailLogger(CustomLogger): ) return None + if user_id is None: + verbose_proxy_logger.debug("No user_id provided. Unable to lookup user email") + return None + user_row = await prisma_client.db.litellm_usertable.find_unique( where={"user_id": user_id} ) diff --git a/litellm/proxy/common_utils/key_rotation_manager.py b/litellm/proxy/common_utils/key_rotation_manager.py index 51f86e514ca..fe843d6e573 100644 --- a/litellm/proxy/common_utils/key_rotation_manager.py +++ b/litellm/proxy/common_utils/key_rotation_manager.py @@ -13,12 +13,12 @@ from litellm.constants import ( LITELLM_KEY_ROTATION_GRACE_PERIOD, LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS, ) +from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.proxy._types import ( GenerateKeyResponse, LiteLLM_VerificationToken, RegenerateKeyRequest, ) -from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks from litellm.proxy.management_endpoints.key_management_endpoints import ( _calculate_key_rotation_time, regenerate_key_fn, @@ -32,6 +32,63 @@ from litellm.repositories.verification_token_repository import ( ) +async def log_rotation_schedule_warnings( + prisma_client: PrismaClient, + *, + rotation_job_enabled: bool, + check_interval_seconds: int, +) -> None: + """ + Warn about auto-rotating keys the scheduler cannot honor. + + Rotation only happens when the background job runs, so a key asking to + rotate every 60s never rotates while the job is off, and rotates at best + once per check interval (24h by default) while it is on. + """ + try: + auto_rotate_keys = await VerificationTokenRepository(prisma_client).table.find_many(where={"auto_rotate": True}) + except Exception as e: + verbose_proxy_logger.debug("Could not check for auto-rotating keys: %s", e) + return + + if not auto_rotate_keys: + return + + if not rotation_job_enabled: + verbose_proxy_logger.warning( + "%s key(s) have auto_rotate=true but the key rotation job is disabled; they will never rotate. " + "Set LITELLM_KEY_ROTATION_ENABLED=true to enable it.", + len(auto_rotate_keys), + ) + return + + shortest_interval = min( + ( + interval + for interval in (_parse_rotation_interval(key.rotation_interval) for key in auto_rotate_keys) + if interval is not None + ), + default=None, + ) + if shortest_interval is not None and shortest_interval < check_interval_seconds: + verbose_proxy_logger.warning( + "Shortest key rotation_interval is %ss but the rotation job only runs every %ss, " + "so rotations lag by up to that long. Lower LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS.", + shortest_interval, + check_interval_seconds, + ) + + +def _parse_rotation_interval(rotation_interval: "str | None") -> "int | None": + if not rotation_interval: + return None + try: + return duration_in_seconds(rotation_interval) + except ValueError: + verbose_proxy_logger.warning("Invalid rotation_interval: %s", rotation_interval) + return None + + class KeyRotationManager: """ Manages automated key rotation based on individual key rotation schedules. @@ -192,13 +249,3 @@ class KeyRotationManager: "key_rotation_at": next_rotation_time, }, ) - - # Call the existing rotation hook for notifications, audit logs, etc. - if isinstance(response, GenerateKeyResponse): - await KeyManagementEventHooks.async_key_rotated_hook( - data=regenerate_request, - existing_key_row=key, - response=response, - user_api_key_dict=system_user, - litellm_changed_by=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME, - ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index a20b557e38b..e5663dad579 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -544,6 +544,7 @@ from litellm.proxy.utils import ( migrate_passwords_to_scrypt_async, model_dump_with_preserved_fields, prefetch_config_params, + register_email_logger_callback, update_spend, ) from litellm.proxy.video_endpoints.endpoints import router as video_router @@ -5168,6 +5169,13 @@ class ProxyConfig: if _alert == "slack": # [OLD] v0 implementation - already handled by update_values above pass + elif _alert == "email": + if not register_email_logger_callback(proxy_logging_obj.email_logging_instance): + verbose_proxy_logger.warning( + "Email alerting is enabled but no email transport is configured. " + "Set SENDGRID_API_KEY, RESEND_API_KEY or SMTP_HOST so key created / " + "rotated and user invite emails can be sent." + ) else: # [NEW] v1 implementation - init as a custom logger if _alert in litellm._known_custom_logger_compatible_callbacks: @@ -8275,6 +8283,7 @@ class ProxyStartupEvent: try: from litellm.proxy.common_utils.key_rotation_manager import ( KeyRotationManager, + log_rotation_schedule_warnings, ) # Get prisma_client and proxy_logging_obj from global scope @@ -8294,12 +8303,27 @@ class ProxyStartupEvent: seconds=LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS, id="key_rotation_job", ) + await log_rotation_schedule_warnings( + prisma_client, + rotation_job_enabled=True, + check_interval_seconds=LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS, + ) else: verbose_proxy_logger.warning("Key rotation enabled but prisma_client not available") except Exception as e: verbose_proxy_logger.warning(f"Failed to setup key rotation job: {e}") else: verbose_proxy_logger.debug("Key rotation disabled (set LITELLM_KEY_ROTATION_ENABLED=true to enable)") + if prisma_client is not None: + from litellm.proxy.common_utils.key_rotation_manager import ( + log_rotation_schedule_warnings, + ) + + await log_rotation_schedule_warnings( + prisma_client, + rotation_job_enabled=False, + check_interval_seconds=LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS, + ) await cls._initialize_expired_ui_session_key_cleanup_background_job(scheduler=scheduler) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 171e17ce650..f9424cc0cf3 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -226,6 +226,30 @@ def _get_email_logger_class(): return BaseEmailLogger +def register_email_logger_callback(email_logger: "object | None") -> bool: + """ + Register the proxy's email logger as a litellm callback. + + Key lifecycle emails (created / rotated) and user invites are dispatched by + looking up ``BaseEmailLogger`` instances in ``logging_callback_manager``. + Admins who configure email through ``general_settings.alerting: ["email"]`` + (what the Admin UI writes) never register such a logger, so those emails are + silently dropped. Registering the already-built instance here closes that gap. + + ``BaseEmailLogger`` itself has no transport, so only its subclasses register. + + Returns True when a sending-capable logger was registered. + """ + if BaseEmailLogger is None or email_logger is None: + return False + + if type(email_logger) is BaseEmailLogger or not isinstance(email_logger, BaseEmailLogger): + return False + + litellm.logging_callback_manager.add_litellm_callback(email_logger) + return True + + class InternalUsageCache: def __init__(self, dual_cache: DualCache): self.dual_cache: DualCache = dual_cache diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_base_email.py b/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_base_email.py index 61303340570..2475121b3b8 100644 --- a/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_base_email.py +++ b/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_base_email.py @@ -1378,3 +1378,21 @@ async def test_budget_alert_release_failure_does_not_propagate(base_email_logger ) mock_cache.async_delete_cache.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_lookup_user_email_from_db_without_user_id(base_email_logger): + """ + Keys with no owner reach the lookup with user_id=None; prisma rejects that with + "`where.user_id`: A value is required but not set", which aborted the whole email. + """ + mock_prisma_client = mock.MagicMock() + mock_prisma_client.db.litellm_usertable.find_unique = mock.AsyncMock() + + with patch( + "litellm.proxy.proxy_server.prisma_client", + mock_prisma_client, + ): + assert await base_email_logger._lookup_user_email_from_db(user_id=None) is None + + mock_prisma_client.db.litellm_usertable.find_unique.assert_not_called() diff --git a/tests/test_litellm/proxy/common_utils/test_key_rotation_e2e.py b/tests/test_litellm/proxy/common_utils/test_key_rotation_e2e.py index d6e1d22fdde..f6c8cf15c12 100644 --- a/tests/test_litellm/proxy/common_utils/test_key_rotation_e2e.py +++ b/tests/test_litellm/proxy/common_utils/test_key_rotation_e2e.py @@ -258,8 +258,8 @@ class TestKeyRotationErrorResilience: @pytest.mark.asyncio async def test_hook_failure_does_not_prevent_db_update(self): """ - If the rotation hook (async_key_rotated_hook) fails, the database - update for rotation_count should still have succeeded (it runs before the hook). + The rotation hook is owned by regenerate_key_fn, so a failing hook must not + surface in the rotation job or block the rotation_count update. """ mock_prisma = AsyncMock() manager = KeyRotationManager(mock_prisma) @@ -282,15 +282,12 @@ class TestKeyRotationErrorResilience: return_value=mock_response, ): with patch( - "litellm.proxy.common_utils.key_rotation_manager.KeyManagementEventHooks.async_key_rotated_hook", + "litellm.proxy.hooks.key_management_event_hooks.KeyManagementEventHooks.async_key_rotated_hook", new_callable=AsyncMock, side_effect=Exception("Hook failed: secret manager down"), ): - # This will raise because the hook fails - with pytest.raises(Exception, match="Hook failed"): - await manager._rotate_key(key) + await manager._rotate_key(key) - # The DB update should have been called BEFORE the hook mock_prisma.db.litellm_verificationtoken.update.assert_called_once() update_data = mock_prisma.db.litellm_verificationtoken.update.call_args[1][ "data" @@ -370,7 +367,7 @@ class TestKeyRotationFullFlow: return_value=mock_response, ): with patch( - "litellm.proxy.common_utils.key_rotation_manager.KeyManagementEventHooks.async_key_rotated_hook", + "litellm.proxy.hooks.key_management_event_hooks.KeyManagementEventHooks.async_key_rotated_hook", new_callable=AsyncMock, ): await manager.process_rotations() @@ -428,7 +425,7 @@ class TestKeyRotationFullFlow: return_value=mock_response, ): with patch( - "litellm.proxy.common_utils.key_rotation_manager.KeyManagementEventHooks.async_key_rotated_hook", + "litellm.proxy.hooks.key_management_event_hooks.KeyManagementEventHooks.async_key_rotated_hook", new_callable=AsyncMock, ): await manager._rotate_key(key) @@ -493,7 +490,7 @@ class TestKeyRotationFullFlow: return_value=mock_response, ): with patch( - "litellm.proxy.common_utils.key_rotation_manager.KeyManagementEventHooks.async_key_rotated_hook", + "litellm.proxy.hooks.key_management_event_hooks.KeyManagementEventHooks.async_key_rotated_hook", new_callable=AsyncMock, ): await manager._rotate_key(key) diff --git a/tests/test_litellm/proxy/common_utils/test_key_rotation_integration.py b/tests/test_litellm/proxy/common_utils/test_key_rotation_integration.py index bdea37f1358..42ce6c64739 100644 --- a/tests/test_litellm/proxy/common_utils/test_key_rotation_integration.py +++ b/tests/test_litellm/proxy/common_utils/test_key_rotation_integration.py @@ -79,7 +79,7 @@ class TestKeyRotationManagerPassesKeyAlias: side_effect=capture_regenerate_key_fn, ): with patch( - "litellm.proxy.common_utils.key_rotation_manager.KeyManagementEventHooks.async_key_rotated_hook", + "litellm.proxy.hooks.key_management_event_hooks.KeyManagementEventHooks.async_key_rotated_hook", new_callable=AsyncMock, ): rotation_manager = KeyRotationManager(mock_prisma) @@ -132,7 +132,7 @@ class TestKeyRotationManagerPassesKeyAlias: side_effect=capture_regenerate_key_fn, ): with patch( - "litellm.proxy.common_utils.key_rotation_manager.KeyManagementEventHooks.async_key_rotated_hook", + "litellm.proxy.hooks.key_management_event_hooks.KeyManagementEventHooks.async_key_rotated_hook", new_callable=AsyncMock, ): rotation_manager = KeyRotationManager(mock_prisma) diff --git a/tests/test_litellm/proxy/common_utils/test_key_rotation_manager.py b/tests/test_litellm/proxy/common_utils/test_key_rotation_manager.py index 18432d106af..8be660d40d9 100644 --- a/tests/test_litellm/proxy/common_utils/test_key_rotation_manager.py +++ b/tests/test_litellm/proxy/common_utils/test_key_rotation_manager.py @@ -2,6 +2,7 @@ Test key rotation manager functionality """ +import logging import os import sys from datetime import datetime, timedelta, timezone @@ -15,7 +16,10 @@ from litellm.proxy._types import ( GenerateKeyResponse, LiteLLM_VerificationToken, ) -from litellm.proxy.common_utils.key_rotation_manager import KeyRotationManager +from litellm.proxy.common_utils.key_rotation_manager import ( + KeyRotationManager, + log_rotation_schedule_warnings, +) class TestKeyRotationManager: @@ -198,11 +202,8 @@ class TestKeyRotationManager: "litellm.proxy.common_utils.key_rotation_manager.regenerate_key_fn", return_value=mock_response, ): - with patch( - "litellm.proxy.common_utils.key_rotation_manager.KeyManagementEventHooks.async_key_rotated_hook" - ): - # Execute - await manager._rotate_key(key_to_rotate) + # Execute + await manager._rotate_key(key_to_rotate) # Verify database update was called with correct data mock_prisma_client.db.litellm_verificationtoken.update.assert_called_once() @@ -278,16 +279,114 @@ class TestKeyRotationManager: ) as mock_regenerate: mock_regenerate.return_value = mock_response with patch( - "litellm.proxy.common_utils.key_rotation_manager.KeyManagementEventHooks.async_key_rotated_hook", - new_callable=AsyncMock, + "litellm.proxy.common_utils.key_rotation_manager.LITELLM_KEY_ROTATION_GRACE_PERIOD", + "48h", ): - with patch( - "litellm.proxy.common_utils.key_rotation_manager.LITELLM_KEY_ROTATION_GRACE_PERIOD", - "48h", - ): - await manager._rotate_key(key_to_rotate) + await manager._rotate_key(key_to_rotate) mock_regenerate.assert_called_once() call_args = mock_regenerate.call_args regenerate_request = call_args[1]["data"] assert regenerate_request.grace_period == "48h" + + @pytest.mark.asyncio + async def test_rotate_key_does_not_double_fire_rotation_hook(self): + """ + regenerate_key_fn already fires async_key_rotated_hook, so the rotation job must not + fire it a second time; otherwise every automated rotation sends two emails and writes + two audit logs. + """ + from unittest.mock import patch + + mock_prisma_client = AsyncMock() + manager = KeyRotationManager(mock_prisma_client) + + key_to_rotate = LiteLLM_VerificationToken( + token="old-token", + auto_rotate=True, + rotation_interval="30s", + key_rotation_at=None, + rotation_count=0, + ) + + with patch( + "litellm.proxy.common_utils.key_rotation_manager.regenerate_key_fn", + new_callable=AsyncMock, + return_value=GenerateKeyResponse( + key="new-api-key", token_id="new-token-id", user_id="test-user" + ), + ): + with patch( + "litellm.proxy.hooks.key_management_event_hooks.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + ) as mock_hook: + await manager._rotate_key(key_to_rotate) + + mock_hook.assert_not_called() + + +class TestRotationScheduleWarnings: + """Auto-rotating keys the background job cannot honor must be surfaced at startup.""" + + @pytest.mark.asyncio + async def test_warns_when_rotation_job_disabled(self, caplog): + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many.return_value = [ + LiteLLM_VerificationToken(token="t1", auto_rotate=True, rotation_interval="60s") + ] + + with caplog.at_level(logging.WARNING): + await log_rotation_schedule_warnings( + mock_prisma_client, + rotation_job_enabled=False, + check_interval_seconds=86400, + ) + + assert "LITELLM_KEY_ROTATION_ENABLED" in caplog.text + + @pytest.mark.asyncio + async def test_warns_when_check_interval_coarser_than_rotation_interval(self, caplog): + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many.return_value = [ + LiteLLM_VerificationToken(token="t1", auto_rotate=True, rotation_interval="24h"), + LiteLLM_VerificationToken(token="t2", auto_rotate=True, rotation_interval="60s"), + ] + + with caplog.at_level(logging.WARNING): + await log_rotation_schedule_warnings( + mock_prisma_client, + rotation_job_enabled=True, + check_interval_seconds=86400, + ) + + assert "LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS" in caplog.text + + @pytest.mark.asyncio + async def test_silent_when_schedule_is_honored(self, caplog): + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many.return_value = [ + LiteLLM_VerificationToken(token="t1", auto_rotate=True, rotation_interval="60s") + ] + + with caplog.at_level(logging.WARNING): + await log_rotation_schedule_warnings( + mock_prisma_client, + rotation_job_enabled=True, + check_interval_seconds=15, + ) + + assert caplog.text == "" + + @pytest.mark.asyncio + async def test_silent_when_no_auto_rotate_keys(self, caplog): + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many.return_value = [] + + with caplog.at_level(logging.WARNING): + await log_rotation_schedule_warnings( + mock_prisma_client, + rotation_job_enabled=False, + check_interval_seconds=86400, + ) + + assert caplog.text == "" diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index 4673807a135..5aeb725429b 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -947,3 +947,63 @@ class TestSendEmailStartTls: assert isinstance(context, ssl.SSLContext) assert context.verify_mode == ssl.CERT_REQUIRED assert context.check_hostname is True + + +def test_register_email_logger_callback_registers_sending_logger(): + """ + Key created / rotated emails are dispatched by looking up BaseEmailLogger instances in + logging_callback_manager. Configuring email via `general_settings.alerting: ["email"]` + (what the Admin UI writes) has to register one, or those emails are silently dropped. + """ + import litellm + from litellm_enterprise.enterprise_callbacks.send_emails.base_email import ( + BaseEmailLogger, + ) + from litellm_enterprise.enterprise_callbacks.send_emails.smtp_email import ( + SMTPEmailLogger, + ) + + from litellm.proxy.utils import register_email_logger_callback + + original_callbacks = list(litellm.callbacks) + try: + smtp_logger = SMTPEmailLogger() + assert register_email_logger_callback(smtp_logger) is True + assert ( + litellm.logging_callback_manager.get_custom_loggers_for_type( + callback_type=BaseEmailLogger + ) + == [smtp_logger] + ) + finally: + litellm.callbacks = original_callbacks + + +@pytest.mark.parametrize("email_logger", [None, "not-a-logger"]) +def test_register_email_logger_callback_ignores_non_senders(email_logger): + import litellm + from litellm.proxy.utils import register_email_logger_callback + + original_callbacks = list(litellm.callbacks) + try: + assert register_email_logger_callback(email_logger) is False + assert litellm.callbacks == original_callbacks + finally: + litellm.callbacks = original_callbacks + + +def test_register_email_logger_callback_ignores_base_logger_without_transport(): + """BaseEmailLogger has no transport; registering it would only produce send failures.""" + import litellm + from litellm_enterprise.enterprise_callbacks.send_emails.base_email import ( + BaseEmailLogger, + ) + + from litellm.proxy.utils import register_email_logger_callback + + original_callbacks = list(litellm.callbacks) + try: + assert register_email_logger_callback(BaseEmailLogger()) is False + assert litellm.callbacks == original_callbacks + finally: + litellm.callbacks = original_callbacks