scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 207

fix(proxy): email the plaintext key when a virtual key is created

Browse source 
/key/generate remaps token to the hashed token_id before the invite email is built. Send the sk- key in the field both key-created email templates show.

Fixes #39315
This commit is contained in:
Atharva Mashalkar 2026-10-20 14:15:18 +05:30
parent c9d081b72a
commit efc2dc1524
2 changed files with 141 additions and 38 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
litellm/proxy/hooks/key_management_event_hooks.py
View file

@ -531,6 +531,9 @@ class KeyManagementEventHooks:
from litellm.proxy.proxy_server import general_settings, proxy_logging_obj
# generate_key_fn remaps token to the hashed token_id. The plaintext sk- key stays on key.
plaintext_key: Final = response.get("key") or response.get("token") or ""
##########################
# v2 integration for emails (enterprise)
##########################
@ -547,11 +550,11 @@ class KeyManagementEventHooks:
)
if len(initialized_email_loggers) > 0:
event = SendKeyCreatedEmailEvent(
virtual_key=response.get("key", ""),
virtual_key=plaintext_key,
event="key_created",
event_group=Litellm_EntityType.KEY,
event_message="API Key Created",
token=response.get("token", ""),
token=plaintext_key,
spend=response.get("spend", 0.0),
max_budget=response.get("max_budget", 0.0),
user_id=response.get("user_id", None),
@ -577,7 +580,7 @@ class KeyManagementEventHooks:
event="key_created",
event_group=Litellm_EntityType.KEY,
event_message="API Key Created",
token=response.get("token", ""),
token=plaintext_key,
spend=response.get("spend", 0.0),
max_budget=response.get("max_budget", 0.0),
user_id=response.get("user_id", None),

170
tests/unit/proxy/hooks/test_send_invite_email.py
View file

@ -1,16 +1,48 @@
import pytest
from unittest.mock import AsyncMock, patch, MagicMock
from litellm.proxy.hooks.user_management_event_hooks import UserManagementEventHooks
from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks
from litellm.proxy._types import (
NewUserRequest,
NewUserResponse,
GenerateKeyRequest,
GenerateKeyResponse,
UserAPIKeyAuth,
)
import asyncio
import sys
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting
from litellm.proxy._types import (
GenerateKeyRequest,
GenerateKeyResponse,
NewUserRequest,
NewUserResponse,
UserAPIKeyAuth,
)
from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks
from litellm.proxy.hooks.user_management_event_hooks import UserManagementEventHooks
_PLAINTEXT_KEY = "sk-plaintext-key-39315"
_TOKEN_ID = "a1b2c3d4e5f6789012345678abcdef01"
def _key_created_email_response() -> dict:
"""Response after generate_key_fn remaps token onto the hashed token_id."""
return {
"key": _PLAINTEXT_KEY,
"token": _TOKEN_ID,
"token_id": _TOKEN_ID,
"user_id": "user-123",
"team_id": "team-1",
"key_alias": "ci-key",
"spend": 0.0,
"max_budget": 10.0,
}
def _proxy_server_module(slack_alerting: object) -> SimpleNamespace:
return SimpleNamespace(
general_settings={"alerting": ["email"]},
proxy_logging_obj=SimpleNamespace(slack_alerting_instance=slack_alerting),
premium_user=False,
prisma_client=None,
CommonProxyErrors=SimpleNamespace(not_premium_user=SimpleNamespace(value="not premium")),
litellm_proxy_admin_name="admin-user",
)
@pytest.mark.asyncio
@ -23,9 +55,7 @@ async def test_v1_user_creation_no_email_when_send_invite_email_false():
mock_proxy_logging_obj = MagicMock()
mock_proxy_logging_obj.slack_alerting_instance = mock_slack_alerting
with patch(
"litellm.logging_callback_manager.get_custom_loggers_for_type", return_value=[]
):
with patch("litellm.logging_callback_manager.get_custom_loggers_for_type", return_value=[]):
mock_proxy_server = SimpleNamespace(
general_settings={"alerting": ["email"]},
proxy_logging_obj=mock_proxy_logging_obj,
@ -41,9 +71,7 @@ async def test_v1_user_creation_no_email_when_send_invite_email_false():
user_email="test@example.com",
key="sk-test-key",
)
user_api_key_dict = UserAPIKeyAuth(
user_id="admin-user", api_key="admin-key"
)
user_api_key_dict = UserAPIKeyAuth(user_id="admin-user", api_key="admin-key")
await UserManagementEventHooks.async_send_user_invitation_email(
data=data,
response=response,
@ -62,9 +90,7 @@ async def test_v1_user_creation_sends_email_when_send_invite_email_true():
mock_proxy_logging_obj = MagicMock()
mock_proxy_logging_obj.slack_alerting_instance = mock_slack_alerting
with patch(
"litellm.logging_callback_manager.get_custom_loggers_for_type", return_value=[]
):
with patch("litellm.logging_callback_manager.get_custom_loggers_for_type", return_value=[]):
mock_proxy_server = SimpleNamespace(
general_settings={"alerting": ["email"]},
proxy_logging_obj=mock_proxy_logging_obj,
@ -80,9 +106,7 @@ async def test_v1_user_creation_sends_email_when_send_invite_email_true():
user_email="test@example.com",
key="sk-test-key",
)
user_api_key_dict = UserAPIKeyAuth(
user_id="admin-user", api_key="admin-key"
)
user_api_key_dict = UserAPIKeyAuth(user_id="admin-user", api_key="admin-key")
await UserManagementEventHooks.async_send_user_invitation_email(
data=data,
response=response,
@ -225,9 +249,7 @@ async def test_v1_key_generation_sends_email_when_send_invite_email_true():
proxy_logging_obj=mock_proxy_logging_obj,
litellm_proxy_admin_name="admin-user",
)
with patch.dict(
sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}
):
with patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}):
data = GenerateKeyRequest(
user_email="test@example.com",
send_invite_email=True, # Should send key email
@ -236,9 +258,7 @@ async def test_v1_key_generation_sends_email_when_send_invite_email_true():
user_email="test@example.com",
key="sk-test-key",
)
user_api_key_dict = UserAPIKeyAuth(
user_id="admin-user", api_key="admin-key"
)
user_api_key_dict = UserAPIKeyAuth(user_id="admin-user", api_key="admin-key")
await KeyManagementEventHooks.async_key_generated_hook(
data=data,
response=response,
@ -271,9 +291,7 @@ async def test_v1_key_generation_no_email_when_send_invite_email_false():
proxy_logging_obj=mock_proxy_logging_obj,
litellm_proxy_admin_name="admin-user",
)
with patch.dict(
sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}
):
with patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}):
data = GenerateKeyRequest(
user_email="test@example.com",
send_invite_email=False, # Should NOT send key email
@ -282,12 +300,94 @@ async def test_v1_key_generation_no_email_when_send_invite_email_false():
user_email="test@example.com",
key="sk-test-key",
)
user_api_key_dict = UserAPIKeyAuth(
user_id="admin-user", api_key="admin-key"
)
user_api_key_dict = UserAPIKeyAuth(user_id="admin-user", api_key="admin-key")
await KeyManagementEventHooks.async_key_generated_hook(
data=data,
response=response,
user_api_key_dict=user_api_key_dict,
)
mock_send_key_created_email.assert_not_called()
@pytest.mark.asyncio
async def test_v0_key_created_email_contains_plaintext_key_not_token_id(monkeypatch):
"""The v0 invite email prints webhook token, which must be the sk- key."""
for name in ("EMAIL_LOGO_URL", "EMAIL_SUPPORT_CONTACT", "SMTP_SENDER_LOGO"):
monkeypatch.delenv(name, raising=False)
sent_html: list[str] = []
async def capture_send_email(*, receiver_email, subject, html):
sent_html.append(html)
scheduled: list[asyncio.Task] = []
real_create_task = asyncio.create_task
def schedule_task(coro, *args, **kwargs):
task = real_create_task(coro, *args, **kwargs)
scheduled.append(task)
return task
slack_alerting = SlackAlerting(alerting=["email"])
with (
patch("litellm.logging_callback_manager.get_custom_loggers_for_type", return_value=[]),
patch("litellm.proxy.utils.send_email", capture_send_email),
patch(
"litellm.proxy.hooks.key_management_event_hooks.asyncio.create_task",
side_effect=schedule_task,
),
patch.dict(sys.modules, {"litellm.proxy.proxy_server": _proxy_server_module(slack_alerting)}),
):
await KeyManagementEventHooks._send_key_created_email(_key_created_email_response())
assert scheduled, "key created email was not scheduled"
await asyncio.gather(*scheduled)
assert sent_html, "key created email was not sent"
assert f"Key: <pre>{_PLAINTEXT_KEY}</pre>" in sent_html[0]
assert _TOKEN_ID not in sent_html[0]
@pytest.mark.asyncio
async def test_enterprise_key_created_email_contains_plaintext_key_not_token_id(monkeypatch):
"""The enterprise invite email prints virtual_key, which must be the sk- key."""
pytest.importorskip("litellm_enterprise")
monkeypatch.setenv("EMAIL_INCLUDE_API_KEY", "true")
for name in ("EMAIL_LOGO_URL", "EMAIL_SUPPORT_CONTACT", "SMTP_SENDER_LOGO", "EMAIL_SIGNATURE"):
monkeypatch.delenv(name, raising=False)
from litellm_enterprise.enterprise_callbacks.send_emails.base_email import (
BaseEmailLogger,
)
class RecordingEmailLogger(BaseEmailLogger):
def __init__(self):
super().__init__()
self.html_bodies: list[str] = []
self.events: list = []
async def send_key_created_email(self, send_key_created_email_event):
self.events.append(send_key_created_email_event)
event = send_key_created_email_event
if event.user_email is None:
event = event.model_copy(update={"user_email": "user@example.com"})
await super().send_key_created_email(send_key_created_email_event=event)
async def send_email(self, from_email, to_email, subject, html_body):
self.html_bodies.append(html_body)
recording_logger = RecordingEmailLogger()
with (
patch(
"litellm.logging_callback_manager.get_custom_loggers_for_type",
return_value=[recording_logger],
),
patch.dict(sys.modules, {"litellm.proxy.proxy_server": _proxy_server_module(None)}),
):
await KeyManagementEventHooks._send_key_created_email(_key_created_email_response())
assert recording_logger.events, "enterprise key created email was not sent"
assert recording_logger.events[0].virtual_key == _PLAINTEXT_KEY
assert recording_logger.events[0].token == _PLAINTEXT_KEY
assert recording_logger.html_bodies, "enterprise key created email was not sent"
html = recording_logger.html_bodies[0]
assert f'<div class="key">{_PLAINTEXT_KEY}</div>' in html
assert _TOKEN_ID not in html