diff --git a/litellm/__init__.py b/litellm/__init__.py index 3b67d9e002..0c3d8f9926 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -376,6 +376,7 @@ datadog_params: Optional[Union[DatadogInitParams, Dict]] = None aws_sqs_callback_params: Optional[Dict] = None generic_logger_headers: Optional[Dict] = None default_key_generate_params: Optional[Dict] = None +default_key_max_budget_alert_emails: Optional[Dict[str, list]] = None upperbound_key_generate_params: Optional[LiteLLM_UpperboundKeyGenerateParams] = None key_generation_settings: Optional["StandardKeyGenerationConfig"] = None default_internal_user_params: Optional[Dict] = None diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index fbc24ab17c..19751f6edc 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -2966,7 +2966,7 @@ async def _virtual_key_max_budget_alert_check( owner_email = user_obj.user_email if user_obj else None alert_email_config = (valid_token.metadata or {}).get( "max_budget_alert_emails" - ) + ) or litellm.default_key_max_budget_alert_emails if isinstance(alert_email_config, dict) and alert_email_config: # New path: pass the map through, let the email handler decide what to fire diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index d9ab8e06d0..0f7ffa541a 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -1652,6 +1652,86 @@ async def test_virtual_key_max_budget_alert_check_old_path_below_threshold_no_al assert alert_triggered is False +@pytest.mark.asyncio +async def test_virtual_key_max_budget_alert_check_global_fallback(): + """Test that litellm.default_key_max_budget_alert_emails is used when key metadata has no map""" + alert_triggered = False + captured_call_info = None + + class MockProxyLogging: + async def budget_alerts(self, type, user_info): + nonlocal alert_triggered, captured_call_info + alert_triggered = True + captured_call_info = user_info + + global_config = { + "50": ["global-finance@co.com"], + "75": ["global-finance@co.com", "global-lead@co.com"], + } + valid_token = UserAPIKeyAuth( + token="test-token", + spend=30.0, + max_budget=100.0, + user_id="test-user", + key_alias="test-key", + metadata={}, # no per-key config + ) + + import litellm + original = litellm.default_key_max_budget_alert_emails + try: + litellm.default_key_max_budget_alert_emails = global_config + await _virtual_key_max_budget_alert_check( + valid_token=valid_token, + proxy_logging_obj=MockProxyLogging(), + user_obj=None, + ) + await asyncio.sleep(0.1) + + assert alert_triggered is True + assert captured_call_info.max_budget_alert_emails == global_config + finally: + litellm.default_key_max_budget_alert_emails = original + + +@pytest.mark.asyncio +async def test_virtual_key_max_budget_alert_check_per_key_overrides_global(): + """Test that per-key metadata takes priority over global fallback""" + captured_call_info = None + + class MockProxyLogging: + async def budget_alerts(self, type, user_info): + nonlocal captured_call_info + captured_call_info = user_info + + per_key_config = {"50": ["per-key@co.com"]} + global_config = {"75": ["global@co.com"]} + + valid_token = UserAPIKeyAuth( + token="test-token", + spend=30.0, + max_budget=100.0, + user_id="test-user", + key_alias="test-key", + metadata={"max_budget_alert_emails": per_key_config}, + ) + + import litellm + original = litellm.default_key_max_budget_alert_emails + try: + litellm.default_key_max_budget_alert_emails = global_config + await _virtual_key_max_budget_alert_check( + valid_token=valid_token, + proxy_logging_obj=MockProxyLogging(), + user_obj=None, + ) + await asyncio.sleep(0.1) + + assert captured_call_info.max_budget_alert_emails == per_key_config + finally: + litellm.default_key_max_budget_alert_emails = original + + @pytest.mark.asyncio async def test_get_fuzzy_user_object_case_insensitive_email(): """Test that _get_fuzzy_user_object uses case-insensitive email lookup"""