diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 5c48ea5eb9..f66e08ac84 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -237,12 +237,20 @@ def update_valid_token_with_end_user_params( valid_token: UserAPIKeyAuth, end_user_params: dict ) -> UserAPIKeyAuth: valid_token.end_user_id = end_user_params.get("end_user_id") - valid_token.end_user_tpm_limit = end_user_params.get("end_user_tpm_limit") - valid_token.end_user_rpm_limit = end_user_params.get("end_user_rpm_limit") - valid_token.allowed_model_region = end_user_params.get("allowed_model_region") - valid_token.end_user_model_max_budget = end_user_params.get( - "end_user_model_max_budget" - ) + # Only overwrite token fields when the DB-derived value is not None. + # This prevents DB lookups (where the budget table has no value set) + # from silently clearing values that a custom auth function may have + # already set on the token. + if end_user_params.get("end_user_tpm_limit") is not None: + valid_token.end_user_tpm_limit = end_user_params["end_user_tpm_limit"] + if end_user_params.get("end_user_rpm_limit") is not None: + valid_token.end_user_rpm_limit = end_user_params["end_user_rpm_limit"] + if end_user_params.get("allowed_model_region") is not None: + valid_token.allowed_model_region = end_user_params["allowed_model_region"] + if end_user_params.get("end_user_model_max_budget") is not None: + valid_token.end_user_model_max_budget = end_user_params[ + "end_user_model_max_budget" + ] return valid_token diff --git a/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py b/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py index 246048e4bc..73a9718842 100644 --- a/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py +++ b/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py @@ -1,7 +1,10 @@ import pytest from unittest.mock import AsyncMock, patch import litellm -from litellm.proxy.auth.user_api_key_auth import _run_post_custom_auth_checks +from litellm.proxy.auth.user_api_key_auth import ( + _run_post_custom_auth_checks, + update_valid_token_with_end_user_params, +) from litellm.proxy._types import UserAPIKeyAuth @@ -57,3 +60,60 @@ async def test_custom_auth_run_post_custom_auth_checks_with_end_user_budget_exce parent_otel_span=None, ) mock_budget_check.assert_awaited_once() + + +def test_update_valid_token_does_not_override_custom_auth_values_with_none(): + """ + Greptile feedback: if custom auth sets end_user_model_max_budget on the token, + but the DB end_user has no model_max_budget in their budget table, the DB lookup + should NOT clear the custom-auth-provided value. + """ + custom_auth_budget = {"gpt-4": {"budget_limit": 5.0, "time_period": "1d"}} + valid_token = UserAPIKeyAuth( + token="test_token", + end_user_id="user_1", + end_user_tpm_limit=100, + end_user_rpm_limit=50, + end_user_model_max_budget=custom_auth_budget, + ) + + # Simulate DB lookup that found the end_user but budget table has no limits set + end_user_params = { + "end_user_id": "user_1", + "allowed_model_region": None, + # No tpm_limit, rpm_limit, or model_max_budget from DB + } + + result = update_valid_token_with_end_user_params(valid_token, end_user_params) + + # Custom-auth-provided values should be preserved, not cleared to None + assert result.end_user_tpm_limit == 100 + assert result.end_user_rpm_limit == 50 + assert result.end_user_model_max_budget == custom_auth_budget + assert result.end_user_id == "user_1" + + +def test_update_valid_token_db_values_override_custom_auth_when_set(): + """ + When the DB budget table has explicit values, they should override + whatever the custom auth function set (DB is source of truth). + """ + valid_token = UserAPIKeyAuth( + token="test_token", + end_user_id="user_1", + end_user_tpm_limit=100, + end_user_model_max_budget={"gpt-4": {"budget_limit": 5.0, "time_period": "1d"}}, + ) + + db_budget = {"gpt-4": {"budget_limit": 20.0, "time_period": "1d"}} + end_user_params = { + "end_user_id": "user_1", + "end_user_tpm_limit": 500, + "end_user_model_max_budget": db_budget, + } + + result = update_valid_token_with_end_user_params(valid_token, end_user_params) + + # DB values should win + assert result.end_user_tpm_limit == 500 + assert result.end_user_model_max_budget == db_budget