fix: req changes

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
 
 

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