open-git/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2025-12-06 11:33:26 +08:00
Code Issues Packages Projects Releases Wiki Activity

[Bug fix] Secret Managers Integration - Make email and secret manager operations independent in key management hooks (#17551)

Browse Source 
* TestKeyManagementEventHooksIndependentOperations

* KeyManagementEventHooks - make ops independant
This commit is contained in:
Ishaan Jaff
2025-12-05 15:26:00 -08:00
committed by GitHub
parent a78f40f75a
commit 769f3cc310
2 changed files with 276 additions and 70 deletions
Download Patch File Download Diff File Expand all files Collapse all files

216
litellm/proxy/hooks/key_management_event_hooks.py
View File

@@ -45,9 +45,13 @@ class KeyManagementEventHooks:
)
from litellm.proxy.proxy_server import litellm_proxy_admin_name
await KeyManagementEventHooks._send_key_created_email(
response.model_dump(exclude_none=True)
)
# Send email notification - non-blocking, independent operation
try:
await KeyManagementEventHooks._send_key_created_email(
response.model_dump(exclude_none=True)
)
except Exception as e:
verbose_proxy_logger.warning(f"Failed to send key created email: {e}")
# Enterprise Feature - Audit Logging. Enable with litellm.store_audit_logs = True
if litellm.store_audit_logs is True:
@@ -69,11 +73,17 @@ class KeyManagementEventHooks:
)
)
)
# store the generated key in the secret manager
await KeyManagementEventHooks._store_virtual_key_in_secret_manager(
secret_name=data.key_alias or f"virtual-key-{response.token_id}",
secret_token=response.key,
)
# Store the generated key in the secret manager - non-blocking, independent operation
try:
await KeyManagementEventHooks._store_virtual_key_in_secret_manager(
secret_name=data.key_alias or f"virtual-key-{response.token_id}",
secret_token=response.key,
)
except Exception as e:
verbose_proxy_logger.warning(
f"Failed to store virtual key in secret manager: {e}"
)
@staticmethod
async def async_key_updated_hook(
@@ -132,22 +142,31 @@ class KeyManagementEventHooks:
)
from litellm.proxy.proxy_server import litellm_proxy_admin_name
# store the generated key in the secret manager
# Store the generated key in the secret manager - non-blocking, independent operation
if data is not None and response.token_id is not None:
initial_secret_name = (
existing_key_row.key_alias or f"virtual-key-{existing_key_row.token}"
)
await KeyManagementEventHooks._rotate_virtual_key_in_secret_manager(
current_secret_name=initial_secret_name,
new_secret_name=data.key_alias or f"virtual-key-{response.token_id}",
new_secret_value=response.key,
)
try:
initial_secret_name = (
existing_key_row.key_alias
or f"virtual-key-{existing_key_row.token}"
)
await KeyManagementEventHooks._rotate_virtual_key_in_secret_manager(
current_secret_name=initial_secret_name,
new_secret_name=data.key_alias or f"virtual-key-{response.token_id}",
new_secret_value=response.key,
)
except Exception as e:
verbose_proxy_logger.warning(
f"Failed to rotate virtual key in secret manager: {e}"
)
# send key rotated email if configured
await KeyManagementEventHooks._send_key_rotated_email(
response=response.model_dump(exclude_none=True),
existing_key_alias=existing_key_row.key_alias,
)
# Send key rotated email if configured - non-blocking, independent operation
try:
await KeyManagementEventHooks._send_key_rotated_email(
response=response.model_dump(exclude_none=True),
existing_key_alias=existing_key_row.key_alias,
)
except Exception as e:
verbose_proxy_logger.warning(f"Failed to send key rotated email: {e}")
# store the audit log
if litellm.store_audit_logs is True and existing_key_row.token is not None:
@@ -324,66 +343,109 @@ class KeyManagementEventHooks:
)
@staticmethod
async def _send_key_created_email(response: dict):
def _is_email_sending_enabled() -> bool:
"""
Check if email sending is enabled via v2 enterprise loggers or v0 alerting config.
Returns True only if email is actually configured, preventing any email
processing when the user has not opted in.
"""
# Check v2 enterprise email loggers
try:
from litellm_enterprise.enterprise_callbacks.send_emails.base_email import (
BaseEmailLogger,
)
except ImportError:
raise Exception(
"Trying to use Email Hooks"
+ CommonProxyErrors.missing_enterprise_package.value
initialized_email_loggers = (
litellm.logging_callback_manager.get_custom_loggers_for_type(
callback_type=BaseEmailLogger
)
)
if len(initialized_email_loggers) > 0:
return True
except ImportError:
pass
# Check v0 alerting config
from litellm.proxy.proxy_server import general_settings
if "email" in general_settings.get("alerting", []):
return True
return False
@staticmethod
async def _send_key_created_email(response: dict):
"""
Send key created email if email sending is enabled.
This method is non-blocking - it will return silently if email is not
configured, and will log warnings instead of raising exceptions on failure.
"""
# Early exit if email is not enabled
if not KeyManagementEventHooks._is_email_sending_enabled():
verbose_proxy_logger.debug(
"Email sending not enabled, skipping key created email"
)
return
from litellm.proxy.proxy_server import general_settings, proxy_logging_obj
##########################
# v2 integration for emails (enterprise)
##########################
try:
from litellm_enterprise.enterprise_callbacks.send_emails.base_email import (
BaseEmailLogger,
)
from litellm_enterprise.types.enterprise_callbacks.send_emails import (
SendKeyCreatedEmailEvent,
)
initialized_email_loggers = (
litellm.logging_callback_manager.get_custom_loggers_for_type(
callback_type=BaseEmailLogger
)
)
if len(initialized_email_loggers) > 0:
event = SendKeyCreatedEmailEvent(
virtual_key=response.get("key", ""),
event="key_created",
event_group=Litellm_EntityType.KEY,
event_message="API Key Created",
token=response.get("token", ""),
spend=response.get("spend", 0.0),
max_budget=response.get("max_budget", 0.0),
user_id=response.get("user_id", None),
team_id=response.get("team_id", "Default Team"),
key_alias=response.get("key_alias", None),
)
for email_logger in initialized_email_loggers:
if isinstance(email_logger, BaseEmailLogger):
await email_logger.send_key_created_email(
send_key_created_email_event=event,
)
return
except ImportError:
raise Exception(
"Trying to use Email Hooks"
+ CommonProxyErrors.missing_enterprise_package.value
)
event = SendKeyCreatedEmailEvent(
virtual_key=response.get("key", ""),
event="key_created",
event_group=Litellm_EntityType.KEY,
event_message="API Key Created",
token=response.get("token", ""),
spend=response.get("spend", 0.0),
max_budget=response.get("max_budget", 0.0),
user_id=response.get("user_id", None),
team_id=response.get("team_id", "Default Team"),
key_alias=response.get("key_alias", None),
)
##########################
# v2 integration for emails
##########################
initialized_email_loggers = (
litellm.logging_callback_manager.get_custom_loggers_for_type(
callback_type=BaseEmailLogger
)
)
if len(initialized_email_loggers) > 0:
for email_logger in initialized_email_loggers:
if isinstance(email_logger, BaseEmailLogger):
await email_logger.send_key_created_email(
send_key_created_email_event=event,
)
pass
##########################
# v0 integration for emails
##########################
else:
if "email" not in general_settings.get("alerting", []):
raise ValueError(
"Email alerting not setup on config.yaml. Please set `alerting=['email']. \nDocs: https://docs.litellm.ai/docs/proxy/email`"
)
if "email" in general_settings.get("alerting", []):
from litellm.proxy._types import WebhookEvent
event = WebhookEvent(
event="key_created",
event_group=Litellm_EntityType.KEY,
event_message="API Key Created",
token=response.get("token", ""),
spend=response.get("spend", 0.0),
max_budget=response.get("max_budget", 0.0),
user_id=response.get("user_id", None),
team_id=response.get("team_id", "Default Team"),
key_alias=response.get("key_alias", None),
)
# If user configured email alerting - send an Email letting their end-user know the key was created
asyncio.create_task(
proxy_logging_obj.slack_alerting_instance.send_key_created_or_user_invited_email(
@@ -393,25 +455,39 @@ class KeyManagementEventHooks:
@staticmethod
async def _send_key_rotated_email(response: dict, existing_key_alias: Optional[str]):
"""
Send key rotated email if email sending is enabled.
This method is non-blocking - it will return silently if email is not
configured, and will log warnings instead of raising exceptions on failure.
"""
# Early exit if email is not enabled
if not KeyManagementEventHooks._is_email_sending_enabled():
verbose_proxy_logger.debug(
"Email sending not enabled, skipping key rotated email"
)
return
try:
from litellm_enterprise.enterprise_callbacks.send_emails.base_email import (
BaseEmailLogger,
)
except ImportError:
raise Exception(
"Trying to use Email Hooks"
+ CommonProxyErrors.missing_enterprise_package.value
# Enterprise package not installed - v0 doesn't support key rotated email
verbose_proxy_logger.debug(
"Enterprise package not installed, skipping key rotated email"
)
return
try:
from litellm_enterprise.types.enterprise_callbacks.send_emails import (
SendKeyRotatedEmailEvent,
)
except ImportError:
raise Exception(
"Trying to use Email Hooks"
+ CommonProxyErrors.missing_enterprise_package.value
verbose_proxy_logger.debug(
"Enterprise types not available, skipping key rotated email"
)
return
event = SendKeyRotatedEmailEvent(
virtual_key=response.get("key", ""),

130
tests/test_litellm/proxy/hooks/test_key_management_event_hooks.py Normal file
View File

@@ -0,0 +1,130 @@
"""
Tests for KeyManagementEventHooks.
Validates that email and secret manager operations are independent and non-blocking.
"""
import os
import sys
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
sys.path.insert(0, os.path.abspath("../../../.."))
from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks
class TestKeyManagementEventHooksIndependentOperations:
"""Tests that email and secret manager operations are independent."""
@pytest.mark.asyncio
async def test_email_failure_does_not_block_secret_manager(self):
"""
Test that if email sending fails, secret manager operation still runs.
This validates the independent operation design where one failure
does not block the other operation.
"""
secret_manager_called = {"called": False}
# Mock the email method to raise an exception
async def mock_send_email_raises(*args, **kwargs):
raise Exception("Email service unavailable")
# Mock the secret manager method to track if it was called
async def mock_store_secret(*args, **kwargs):
secret_manager_called["called"] = True
# Create mock objects for the hook parameters
mock_data = MagicMock()
mock_data.key_alias = "test-key-alias"
mock_response = MagicMock()
mock_response.model_dump.return_value = {"key": "sk-test", "token": "test-token"}
mock_response.model_dump_json.return_value = '{"key": "sk-test"}'
mock_response.token_id = "token-123"
mock_response.key = "sk-test-key"
mock_user_api_key_dict = MagicMock()
mock_user_api_key_dict.user_id = "user-123"
mock_user_api_key_dict.api_key = "api-key-123"
with patch.object(
KeyManagementEventHooks,
"_send_key_created_email",
side_effect=mock_send_email_raises,
), patch.object(
KeyManagementEventHooks,
"_store_virtual_key_in_secret_manager",
side_effect=mock_store_secret,
), patch(
"litellm.store_audit_logs", False
), patch(
"litellm.proxy.hooks.key_management_event_hooks.verbose_proxy_logger"
):
# Should not raise even though email fails
await KeyManagementEventHooks.async_key_generated_hook(
data=mock_data,
response=mock_response,
user_api_key_dict=mock_user_api_key_dict,
)
# Secret manager should have been called despite email failure
assert secret_manager_called["called"] is True
@pytest.mark.asyncio
async def test_secret_manager_failure_does_not_block_email(self):
"""
Test that if secret manager fails, email operation still runs.
This validates the independent operation design where one failure
does not block the other operation.
"""
email_called = {"called": False}
# Mock the email method to track if it was called
async def mock_send_email(*args, **kwargs):
email_called["called"] = True
# Mock the secret manager method to raise an exception
async def mock_store_secret_raises(*args, **kwargs):
raise Exception("Secret manager unavailable")
# Create mock objects for the hook parameters
mock_data = MagicMock()
mock_data.key_alias = "test-key-alias"
mock_response = MagicMock()
mock_response.model_dump.return_value = {"key": "sk-test", "token": "test-token"}
mock_response.model_dump_json.return_value = '{"key": "sk-test"}'
mock_response.token_id = "token-123"
mock_response.key = "sk-test-key"
mock_user_api_key_dict = MagicMock()
mock_user_api_key_dict.user_id = "user-123"
mock_user_api_key_dict.api_key = "api-key-123"
with patch.object(
KeyManagementEventHooks,
"_send_key_created_email",
side_effect=mock_send_email,
), patch.object(
KeyManagementEventHooks,
"_store_virtual_key_in_secret_manager",
side_effect=mock_store_secret_raises,
), patch(
"litellm.store_audit_logs", False
), patch(
"litellm.proxy.hooks.key_management_event_hooks.verbose_proxy_logger"
):
# Should not raise even though secret manager fails
await KeyManagementEventHooks.async_key_generated_hook(
data=mock_data,
response=mock_response,
user_api_key_dict=mock_user_api_key_dict,
)
# Email should have been called despite secret manager failure
assert email_called["called"] is True