Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/hooks/key_management_event_hooks.py: 24%
226 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1import asyncio
2import json
3from collections.abc import Sequence
4from datetime import datetime, timezone
5from typing import TYPE_CHECKING, Final
7from pydantic import TypeAdapter
9import litellm
10from litellm._logging import verbose_proxy_logger
11from litellm._uuid import uuid
12from litellm.proxy._types import (
13 GenerateKeyRequest,
14 GenerateKeyResponse,
15 KeyRequest,
16 LiteLLM_AuditLogs,
17 Litellm_EntityType,
18 LiteLLM_VerificationToken,
19 LitellmTableNames,
20 RegenerateKeyRequest,
21 UpdateKeyRequest,
22 UserAPIKeyAuth,
23)
24from litellm.proxy.utils import _hash_token_if_needed
25from litellm.secret_managers.base_secret_manager import BaseSecretManager
27if TYPE_CHECKING: 27 ↛ 28line 27 didn't jump to line 28 because the condition on line 27 was never true
28 from prisma import models as prisma_models
30# NOTE: This is the prefix for all virtual keys stored in AWS Secrets Manager
31LITELLM_PREFIX_STORED_VIRTUAL_KEYS: Final = "litellm/"
34class KeyManagementEventHooks:
35 @staticmethod
36 async def async_key_generated_hook(
37 data: GenerateKeyRequest,
38 response: GenerateKeyResponse,
39 user_api_key_dict: UserAPIKeyAuth,
40 litellm_changed_by: str | None = None,
41 ):
42 """
43 Hook that runs after a successful /key/generate request
45 Handles the following:
46 - Sending Email with Key Details
47 - Storing Audit Logs for key generation
48 - Storing Generated Key in DB
49 """
50 from litellm.proxy.management_helpers.audit_logs import (
51 create_audit_log_for_update,
52 get_audit_log_changed_by,
53 is_audit_logging_enabled,
54 )
55 from litellm.proxy.proxy_server import litellm_proxy_admin_name
57 # Send email notification - non-blocking, independent operation
58 if data.send_invite_email is True:
59 try:
60 await KeyManagementEventHooks._send_key_created_email(response.model_dump(exclude_none=True))
61 except Exception as e:
62 verbose_proxy_logger.warning("Failed to send key created email: %s", e)
64 if is_audit_logging_enabled(): 64 ↛ 65line 64 didn't jump to line 65 because the condition on line 64 was never true
65 _updated_values: Final = response.model_dump_json(exclude_none=True)
66 asyncio.create_task(
67 create_audit_log_for_update(
68 request_data=LiteLLM_AuditLogs(
69 id=str(uuid.uuid4()),
70 updated_at=datetime.now(timezone.utc),
71 changed_by=get_audit_log_changed_by(
72 litellm_changed_by=litellm_changed_by,
73 user_api_key_dict=user_api_key_dict,
74 litellm_proxy_admin_name=litellm_proxy_admin_name,
75 ),
76 changed_by_api_key=user_api_key_dict.api_key,
77 table_name=LitellmTableNames.KEY_TABLE_NAME,
78 object_id=response.token_id or "",
79 action="created",
80 updated_values=_updated_values,
81 before_value=None,
82 )
83 )
84 )
86 # Store the generated key in the secret manager - non-blocking, independent operation
87 try:
88 await KeyManagementEventHooks._store_virtual_key_in_secret_manager(
89 secret_name=data.key_alias or f"virtual-key-{response.token_id}",
90 secret_token=response.key,
91 team_id=data.team_id,
92 )
93 except Exception as e:
94 verbose_proxy_logger.warning("Failed to store virtual key in secret manager: %s", e)
96 @staticmethod
97 async def async_key_updated_hook(
98 data: UpdateKeyRequest,
99 existing_key_row: LiteLLM_VerificationToken,
100 response: object,
101 user_api_key_dict: UserAPIKeyAuth,
102 litellm_changed_by: str | None = None,
103 ):
104 """
105 Post /key/update processing hook
107 Handles the following:
108 - Renaming the key's secret in the secret manager when the alias changes
109 - Storing Audit Logs for key update
110 """
111 from litellm.proxy.management_helpers.audit_logs import (
112 create_audit_log_for_update,
113 get_audit_log_changed_by,
114 is_audit_logging_enabled,
115 )
116 from litellm.proxy.proxy_server import litellm_proxy_admin_name
118 if data.key_alias is not None and data.key_alias != existing_key_row.key_alias:
119 try:
120 await KeyManagementEventHooks._rename_virtual_key_in_secret_manager(
121 current_secret_name=existing_key_row.key_alias or f"virtual-key-{existing_key_row.token}",
122 new_secret_name=data.key_alias,
123 team_id=existing_key_row.team_id,
124 )
125 except Exception as e:
126 verbose_proxy_logger.warning("Failed to rename virtual key in secret manager: %s", e)
128 if is_audit_logging_enabled():
129 updated_fields: Final = {
130 **data.model_dump(exclude_none=True),
131 **({"project_id": data.project_id} if "project_id" in data.model_fields_set else {}),
132 }
133 audit_log: Final = LiteLLM_AuditLogs(
134 id=str(uuid.uuid4()),
135 updated_at=datetime.now(timezone.utc),
136 changed_by=get_audit_log_changed_by(
137 litellm_changed_by=litellm_changed_by,
138 user_api_key_dict=user_api_key_dict,
139 litellm_proxy_admin_name=litellm_proxy_admin_name,
140 ),
141 changed_by_api_key=user_api_key_dict.api_key,
142 table_name=LitellmTableNames.KEY_TABLE_NAME,
143 object_id=_hash_token_if_needed(data.key),
144 action="updated",
145 updated_values=json.dumps(updated_fields, default=str),
146 before_value=json.dumps(existing_key_row.json(exclude_none=True), default=str),
147 )
148 masked_values: Final = TypeAdapter(dict[str, object]).validate_json(str(audit_log.updated_values))
149 request_data: Final = (
150 audit_log.model_copy(update={"updated_values": json.dumps({**masked_values, "project_id": None})})
151 if "project_id" in data.model_fields_set and data.project_id is None
152 else audit_log
153 )
154 asyncio.create_task(create_audit_log_for_update(request_data=request_data))
156 @staticmethod
157 async def async_key_rotated_hook(
158 data: RegenerateKeyRequest | None,
159 existing_key_row: LiteLLM_VerificationToken,
160 response: GenerateKeyResponse,
161 user_api_key_dict: UserAPIKeyAuth,
162 litellm_changed_by: str | None = None,
163 ):
164 from litellm.proxy.management_helpers.audit_logs import (
165 create_audit_log_for_update,
166 get_audit_log_changed_by,
167 is_audit_logging_enabled,
168 )
169 from litellm.proxy.proxy_server import litellm_proxy_admin_name
171 # Store the generated key in the secret manager - non-blocking, independent operation
172 if response.token_id is not None:
173 try:
174 initial_secret_name: Final = existing_key_row.key_alias or f"virtual-key-{existing_key_row.token}"
175 requested_alias: Final = data.key_alias if data is not None else None
176 new_secret_name: Final = response.key_alias or requested_alias or initial_secret_name
177 verbose_proxy_logger.info(
178 "Updating secret in secret manager: secret_name=%s",
179 new_secret_name,
180 )
181 team_id: Final = getattr(existing_key_row, "team_id", None)
182 await KeyManagementEventHooks._rotate_virtual_key_in_secret_manager(
183 current_secret_name=initial_secret_name,
184 new_secret_name=new_secret_name,
185 new_secret_value=response.key,
186 team_id=team_id,
187 )
188 verbose_proxy_logger.info(
189 "Secret updated in secret manager: secret_name=%s",
190 new_secret_name,
191 )
192 except Exception as e:
193 verbose_proxy_logger.warning("Failed to rotate virtual key in secret manager: %s", e)
195 # Send key rotated email if configured - non-blocking, independent operation
196 try:
197 await KeyManagementEventHooks._send_key_rotated_email(
198 response=response.model_dump(exclude_none=True),
199 existing_key_alias=existing_key_row.key_alias,
200 )
201 except Exception as e:
202 verbose_proxy_logger.warning("Failed to send key rotated email: %s", e)
204 # store the audit log
205 if is_audit_logging_enabled() and existing_key_row.token is not None:
206 asyncio.create_task(
207 create_audit_log_for_update(
208 request_data=LiteLLM_AuditLogs(
209 id=str(uuid.uuid4()),
210 updated_at=datetime.now(timezone.utc),
211 changed_by=get_audit_log_changed_by(
212 litellm_changed_by=litellm_changed_by,
213 user_api_key_dict=user_api_key_dict,
214 litellm_proxy_admin_name=litellm_proxy_admin_name,
215 ),
216 changed_by_api_key=user_api_key_dict.token,
217 table_name=LitellmTableNames.KEY_TABLE_NAME,
218 object_id=existing_key_row.token,
219 action="rotated",
220 updated_values=response.model_dump_json(exclude_none=True),
221 before_value=existing_key_row.model_dump_json(exclude_none=True),
222 )
223 )
224 )
226 @staticmethod
227 async def async_key_deleted_hook(
228 data: KeyRequest,
229 keys_being_deleted: list[LiteLLM_VerificationToken],
230 response: dict,
231 user_api_key_dict: UserAPIKeyAuth,
232 litellm_changed_by: str | None = None,
233 ):
234 """
235 Post /key/delete processing hook
237 Handles the following:
238 - Storing Audit Logs for key deletion
239 """
240 KeyManagementEventHooks.create_key_deleted_audit_logs(
241 keys_being_deleted=keys_being_deleted,
242 user_api_key_dict=user_api_key_dict,
243 litellm_changed_by=litellm_changed_by,
244 )
245 await KeyManagementEventHooks._delete_virtual_keys_from_secret_manager(keys_being_deleted=keys_being_deleted)
247 @staticmethod
248 def create_key_deleted_audit_logs(
249 keys_being_deleted: Sequence["LiteLLM_VerificationToken | prisma_models.LiteLLM_VerificationToken"],
250 user_api_key_dict: UserAPIKeyAuth,
251 litellm_changed_by: str | None = None,
252 ) -> None:
253 from litellm.proxy.management_helpers.audit_logs import (
254 create_audit_log_for_update,
255 get_audit_log_changed_by,
256 is_audit_logging_enabled,
257 )
258 from litellm.proxy.proxy_server import litellm_proxy_admin_name
260 if not is_audit_logging_enabled():
261 return
262 for key in keys_being_deleted:
263 key_row = LiteLLM_VerificationToken.model_validate(key, from_attributes=True)
264 if key_row.token is None:
265 continue
266 _key_row = key_row.model_dump_json(exclude_none=True)
268 asyncio.create_task(
269 create_audit_log_for_update(
270 request_data=LiteLLM_AuditLogs(
271 id=str(uuid.uuid4()),
272 updated_at=datetime.now(timezone.utc),
273 changed_by=get_audit_log_changed_by(
274 litellm_changed_by=litellm_changed_by,
275 user_api_key_dict=user_api_key_dict,
276 litellm_proxy_admin_name=litellm_proxy_admin_name,
277 ),
278 changed_by_api_key=user_api_key_dict.token,
279 table_name=LitellmTableNames.KEY_TABLE_NAME,
280 object_id=key_row.token,
281 action="deleted",
282 updated_values="{}",
283 before_value=_key_row,
284 )
285 )
286 )
288 @staticmethod
289 async def _store_virtual_key_in_secret_manager(secret_name: str, secret_token: str, team_id: str | None = None):
290 """
291 Store a virtual key in the secret manager
293 Args:
294 secret_name: Name of the virtual key
295 secret_token: Value of the virtual key (example: sk-1234)
296 """
297 if litellm._key_management_settings is not None: 297 ↛ exitline 297 didn't return from function '_store_virtual_key_in_secret_manager' because the condition on line 297 was always true
298 if litellm._key_management_settings.store_virtual_keys is True: 298 ↛ 299line 298 didn't jump to line 299 because the condition on line 298 was never true
299 from litellm.secret_managers.base_secret_manager import (
300 BaseSecretManager,
301 )
303 # store the key in the secret manager
304 if isinstance(litellm.secret_manager_client, BaseSecretManager):
305 tags: Final = getattr(litellm._key_management_settings, "tags", None)
306 description: Final = getattr(litellm._key_management_settings, "description", None)
307 optional_params: Final = await KeyManagementEventHooks._get_secret_manager_optional_params(team_id)
308 verbose_proxy_logger.debug(
309 "Creating secret with %s and tags=%s and description=%s", secret_name, tags, description
310 )
312 await litellm.secret_manager_client.async_write_secret(
313 secret_name=KeyManagementEventHooks._get_secret_name(secret_name),
314 description=description,
315 secret_value=secret_token,
316 tags=tags,
317 optional_params=optional_params,
318 )
320 @staticmethod
321 async def _rotate_virtual_key_in_secret_manager(
322 current_secret_name: str,
323 new_secret_name: str,
324 new_secret_value: str,
325 team_id: str | None = None,
326 ):
327 """
328 Update a virtual key in the secret manager
330 Args:
331 current_secret_name: Current name of the virtual key
332 new_secret_name: New name of the virtual key
333 new_secret_value: New value of the virtual key (example: sk-1234)
334 team_id: Optional team ID to get team-specific secret manager settings
335 """
336 secret_manager: Final = KeyManagementEventHooks._stored_virtual_key_secret_manager()
337 if secret_manager is None:
338 return
339 optional_params: Final = await KeyManagementEventHooks._get_secret_manager_optional_params(team_id)
340 await secret_manager.async_rotate_secret(
341 current_secret_name=KeyManagementEventHooks._get_secret_name(current_secret_name),
342 new_secret_name=KeyManagementEventHooks._get_secret_name(new_secret_name),
343 new_secret_value=new_secret_value,
344 optional_params=optional_params,
345 )
347 @staticmethod
348 def _stored_virtual_key_secret_manager() -> BaseSecretManager | None:
349 """
350 The secret manager client that stores virtual keys, or None when virtual keys are not stored in one
351 """
352 if litellm._key_management_settings is None or litellm._key_management_settings.store_virtual_keys is not True:
353 return None
354 if not isinstance(litellm.secret_manager_client, BaseSecretManager):
355 return None
356 return litellm.secret_manager_client
358 @staticmethod
359 async def _rename_virtual_key_in_secret_manager(
360 current_secret_name: str,
361 new_secret_name: str,
362 team_id: str | None = None,
363 ) -> None:
364 """
365 Move a virtual key to a new secret name, keeping its current value
367 Args:
368 current_secret_name: Current name of the virtual key
369 new_secret_name: New name of the virtual key
370 team_id: Optional team ID to get team-specific secret manager settings
371 """
372 secret_manager: Final = KeyManagementEventHooks._stored_virtual_key_secret_manager()
373 if secret_manager is None:
374 return
375 optional_params: Final = await KeyManagementEventHooks._get_secret_manager_optional_params(team_id)
376 current_secret_value: Final = await secret_manager.async_read_secret(
377 secret_name=KeyManagementEventHooks._get_secret_name(current_secret_name),
378 optional_params=optional_params,
379 )
380 if current_secret_value is None:
381 verbose_proxy_logger.warning(
382 "Secret %s not found in secret manager, skipping rename to %s", current_secret_name, new_secret_name
383 )
384 return
385 verbose_proxy_logger.info(
386 "Renaming secret in secret manager: current_secret_name=%s new_secret_name=%s",
387 current_secret_name,
388 new_secret_name,
389 )
390 await secret_manager.async_rotate_secret(
391 current_secret_name=KeyManagementEventHooks._get_secret_name(current_secret_name),
392 new_secret_name=KeyManagementEventHooks._get_secret_name(new_secret_name),
393 new_secret_value=current_secret_value,
394 optional_params=optional_params,
395 )
397 @staticmethod
398 def _get_secret_name(secret_name: str) -> str:
399 if litellm._key_management_settings.prefix_for_stored_virtual_keys.endswith("/"):
400 return f"{litellm._key_management_settings.prefix_for_stored_virtual_keys}{secret_name}"
401 else:
402 return f"{litellm._key_management_settings.prefix_for_stored_virtual_keys}/{secret_name}"
404 @staticmethod
405 async def _delete_virtual_keys_from_secret_manager(
406 keys_being_deleted: list[LiteLLM_VerificationToken],
407 ):
408 """
409 Deletes virtual keys from the secret manager
411 Args:
412 keys_being_deleted: List of keys being deleted, this is passed down from the /key/delete operation
413 """
414 if litellm._key_management_settings is not None:
415 if litellm._key_management_settings.store_virtual_keys is True:
416 from litellm.secret_managers.base_secret_manager import (
417 BaseSecretManager,
418 )
420 if isinstance(litellm.secret_manager_client, BaseSecretManager):
421 team_settings_cache: Final[dict[str | None, dict | None]] = {}
422 for key in keys_being_deleted:
423 if key.key_alias is not None:
424 team_id = getattr(key, "team_id", None)
425 if team_id not in team_settings_cache:
426 team_settings_cache[
427 team_id
428 ] = await KeyManagementEventHooks._get_secret_manager_optional_params(team_id)
429 optional_params = team_settings_cache[team_id]
430 await litellm.secret_manager_client.async_delete_secret(
431 secret_name=KeyManagementEventHooks._get_secret_name(key.key_alias),
432 optional_params=optional_params,
433 )
434 else:
435 verbose_proxy_logger.warning(
436 "KeyManagementEventHooks._delete_virtual_key_from_secret_manager: Key alias not found for key %s. Skipping deletion from secret manager.",
437 key.token,
438 )
440 @staticmethod
441 async def _get_secret_manager_optional_params(
442 team_id: str | None,
443 ) -> dict | None:
444 if team_id is None:
445 return None
447 try:
448 from litellm.proxy import proxy_server as proxy_server_module
449 except ImportError:
450 return None
452 prisma_client: Final = getattr(proxy_server_module, "prisma_client", None)
453 user_api_key_cache: Final = getattr(proxy_server_module, "user_api_key_cache", None)
455 if prisma_client is None or user_api_key_cache is None:
456 return None
458 try:
459 from litellm.proxy.auth.auth_checks import get_team_object
461 team_obj: Final = await get_team_object(
462 team_id=team_id,
463 prisma_client=prisma_client,
464 user_api_key_cache=user_api_key_cache,
465 )
466 except Exception as exc: # pragma: no cover - defensive logging
467 verbose_proxy_logger.debug("Unable to load team metadata for team_id=%s: %s", team_id, exc)
468 return None
470 metadata = getattr(team_obj, "metadata", None)
471 if metadata is None:
472 return None
474 if hasattr(metadata, "model_dump"):
475 metadata = metadata.model_dump()
477 if not isinstance(metadata, dict):
478 return None
480 team_settings: Final = metadata.get("secret_manager_settings")
481 if isinstance(team_settings, dict) and team_settings:
482 return dict(team_settings)
484 return None
486 @staticmethod
487 def _is_email_sending_enabled() -> bool:
488 """
489 Check if email sending is enabled via v2 enterprise loggers or v0 alerting config.
491 Returns True only if email is actually configured, preventing any email
492 processing when the user has not opted in.
493 """
494 # Check v2 enterprise email loggers
495 try:
496 from litellm_enterprise.enterprise_callbacks.send_emails.base_email import (
497 BaseEmailLogger,
498 )
500 initialized_email_loggers: Final = litellm.logging_callback_manager.get_custom_loggers_for_type(
501 callback_type=BaseEmailLogger
502 )
503 if len(initialized_email_loggers) > 0: 503 ↛ 504line 503 didn't jump to line 504 because the condition on line 503 was never true
504 return True
505 except ImportError:
506 pass
508 # Check v0 alerting config
509 from litellm.proxy.proxy_server import general_settings
511 if "email" in general_settings.get("alerting", []): 511 ↛ 512line 511 didn't jump to line 512 because the condition on line 511 was never true
512 return True
514 return False
516 @staticmethod
517 async def _send_key_created_email(response: dict):
518 """
519 Send key created email if email sending is enabled.
521 This method is non-blocking - it will return silently if email is not
522 configured, and will log warnings instead of raising exceptions on failure.
523 """
524 # Early exit if email is not enabled
525 if not KeyManagementEventHooks._is_email_sending_enabled(): 525 ↛ 529line 525 didn't jump to line 529 because the condition on line 525 was always true
526 verbose_proxy_logger.debug("Email sending not enabled, skipping key created email")
527 return
529 from litellm.proxy.proxy_server import general_settings, proxy_logging_obj
531 ##########################
532 # v2 integration for emails (enterprise)
533 ##########################
534 try:
535 from litellm_enterprise.enterprise_callbacks.send_emails.base_email import (
536 BaseEmailLogger,
537 )
538 from litellm_enterprise.types.enterprise_callbacks.send_emails import (
539 SendKeyCreatedEmailEvent,
540 )
542 initialized_email_loggers: Final = litellm.logging_callback_manager.get_custom_loggers_for_type(
543 callback_type=BaseEmailLogger
544 )
545 if len(initialized_email_loggers) > 0:
546 event = SendKeyCreatedEmailEvent(
547 virtual_key=response.get("key", ""),
548 event="key_created",
549 event_group=Litellm_EntityType.KEY,
550 event_message="API Key Created",
551 token=response.get("token", ""),
552 spend=response.get("spend", 0.0),
553 max_budget=response.get("max_budget", 0.0),
554 user_id=response.get("user_id", None),
555 team_id=response.get("team_id", "Default Team"),
556 key_alias=response.get("key_alias", None),
557 )
558 for email_logger in initialized_email_loggers:
559 if isinstance(email_logger, BaseEmailLogger):
560 await email_logger.send_key_created_email(
561 send_key_created_email_event=event,
562 )
563 return
564 except ImportError:
565 pass
567 ##########################
568 # v0 integration for emails
569 ##########################
570 if "email" in general_settings.get("alerting", []):
571 from litellm.proxy._types import WebhookEvent
573 event = WebhookEvent(
574 event="key_created",
575 event_group=Litellm_EntityType.KEY,
576 event_message="API Key Created",
577 token=response.get("token", ""),
578 spend=response.get("spend", 0.0),
579 max_budget=response.get("max_budget", 0.0),
580 user_id=response.get("user_id", None),
581 team_id=response.get("team_id", "Default Team"),
582 key_alias=response.get("key_alias", None),
583 )
584 # If user configured email alerting - send an Email letting their end-user know the key was created
585 asyncio.create_task(
586 proxy_logging_obj.slack_alerting_instance.send_key_created_or_user_invited_email(
587 webhook_event=event,
588 )
589 )
591 @staticmethod
592 async def _send_key_rotated_email(response: dict, existing_key_alias: str | None):
593 """
594 Send key rotated email if email sending is enabled.
596 This method is non-blocking - it will return silently if email is not
597 configured, and will log warnings instead of raising exceptions on failure.
598 """
599 # Early exit if email is not enabled
600 if not KeyManagementEventHooks._is_email_sending_enabled():
601 verbose_proxy_logger.debug("Email sending not enabled, skipping key rotated email")
602 return
604 try:
605 from litellm_enterprise.enterprise_callbacks.send_emails.base_email import (
606 BaseEmailLogger,
607 )
608 except ImportError:
609 # Enterprise package not installed - v0 doesn't support key rotated email
610 verbose_proxy_logger.debug("Enterprise package not installed, skipping key rotated email")
611 return
613 try:
614 from litellm_enterprise.types.enterprise_callbacks.send_emails import (
615 SendKeyRotatedEmailEvent,
616 )
617 except ImportError:
618 verbose_proxy_logger.debug("Enterprise types not available, skipping key rotated email")
619 return
621 event: Final = SendKeyRotatedEmailEvent(
622 virtual_key=response.get("key", ""),
623 event="key_rotated",
624 event_group=Litellm_EntityType.KEY,
625 event_message="API Key Rotated",
626 token=response.get("token", ""),
627 spend=response.get("spend", 0.0),
628 max_budget=response.get("max_budget", 0.0),
629 user_id=response.get("user_id", None),
630 team_id=response.get("team_id", "Default Team"),
631 key_alias=response.get("key_alias", existing_key_alias),
632 )
634 ##########################
635 # v2 integration for emails
636 ##########################
637 initialized_email_loggers: Final = litellm.logging_callback_manager.get_custom_loggers_for_type(
638 callback_type=BaseEmailLogger
639 )
640 if len(initialized_email_loggers) > 0:
641 for email_logger in initialized_email_loggers:
642 if isinstance(email_logger, BaseEmailLogger):
643 await email_logger.send_key_rotated_email(
644 send_key_rotated_email_event=event,
645 )