scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-15 23:31:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 213

fix(proxy): deliver key rotation emails and warn on unhonored rotation schedules

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Mubashir Osmani 2026-07-25 05:58:59 +00:00
parent b9b27c2beb
commit 83debc98e2
9 changed files with 309 additions and 36 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py
View file

@ -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}
)

69
litellm/proxy/common_utils/key_rotation_manager.py
View file

@ -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,
)

24
litellm/proxy/proxy_server.py
View file

@ -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)

24
litellm/proxy/utils.py
View file

@ -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

18
tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_base_email.py
View file

@ -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()

17
tests/test_litellm/proxy/common_utils/test_key_rotation_e2e.py
View file

@ -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)

4
tests/test_litellm/proxy/common_utils/test_key_rotation_integration.py
View file

@ -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)

125
tests/test_litellm/proxy/common_utils/test_key_rotation_manager.py
View file

@ -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 == ""

60
tests/test_litellm/proxy/test_proxy_utils.py
View file

@ -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