From d6c2de6ab62c5c11f264bf759ff11d2286e5da6a Mon Sep 17 00:00:00 2001
From: sharziki <sharvil.saxena@gmail.com>
Date: Thu, 28 May 2026 15:01:59 -0400
Subject: [PATCH] fix(proxy): propagate end_user_max_budget in
 update_valid_token_with_end_user_params

`update_valid_token_with_end_user_params()` updated every end-user field
from the DB-derived `end_user_params` dict except `end_user_max_budget`.
When the same API key was reused with different end-users, the cached
`UserAPIKeyAuth` retained the stale budget from the previous end-user,
causing incorrect budget enforcement.

Add the missing conditional update for `end_user_max_budget`, matching
the existing pattern for the other end-user fields.

Closes #29142
---
 litellm/proxy/auth/user_api_key_auth.py       |  2 +
 .../auth/test_custom_auth_end_user_budget.py  | 45 +++++++++++++++++++
 2 files changed, 47 insertions(+)

diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py
index 03278633928..6e366de7599 100644
--- a/litellm/proxy/auth/user_api_key_auth.py
+++ b/litellm/proxy/auth/user_api_key_auth.py
@@ -416,6 +416,8 @@ def update_valid_token_with_end_user_params(
         valid_token.end_user_model_max_budget = end_user_params[
             "end_user_model_max_budget"
         ]
+    if end_user_params.get("end_user_max_budget") is not None:
+        valid_token.end_user_max_budget = end_user_params["end_user_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 4084fa4f3aa..3de87a114a1 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
@@ -143,3 +143,48 @@ def test_update_valid_token_db_values_override_custom_auth_when_set():
     # DB values should win
     assert result.end_user_tpm_limit == 500
     assert result.end_user_model_max_budget == db_budget
+
+
+def test_update_valid_token_updates_end_user_max_budget():
+    """
+    Regression for #29142: update_valid_token_with_end_user_params must
+    propagate end_user_max_budget from the DB-derived end_user_params so
+    that a cached token does not retain a stale budget from a previous
+    end-user.
+    """
+    valid_token = UserAPIKeyAuth(
+        token="test_token",
+        end_user_id="user_old",
+        end_user_max_budget=50.0,
+    )
+
+    end_user_params = {
+        "end_user_id": "user_new",
+        "end_user_max_budget": 5.0,
+    }
+
+    result = update_valid_token_with_end_user_params(valid_token, end_user_params)
+
+    assert result.end_user_id == "user_new"
+    assert result.end_user_max_budget == 5.0
+
+
+def test_update_valid_token_preserves_end_user_max_budget_when_db_none():
+    """
+    When the DB end_user has no max_budget set, the custom-auth-provided
+    value on the token should not be cleared.
+    """
+    valid_token = UserAPIKeyAuth(
+        token="test_token",
+        end_user_id="user_1",
+        end_user_max_budget=100.0,
+    )
+
+    end_user_params = {
+        "end_user_id": "user_1",
+        # No end_user_max_budget from DB
+    }
+
+    result = update_valid_token_with_end_user_params(valid_token, end_user_params)
+
+    assert result.end_user_max_budget == 100.0