From 8d0fa9465ee611f4593b943fb7f2b927f1f96d39 Mon Sep 17 00:00:00 2001 From: Sricharan Anugandula Date: Wed, 9 Sep 2026 01:04:22 +0530 Subject: [PATCH 1/3] fix(proxy/auth): copy end_user_max_budget in update_valid_token_with_end_user_params --- litellm/proxy/auth/user_api_key_auth.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 93293db24c6..9f288128fce 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -563,6 +563,8 @@ def update_valid_token_with_end_user_params(valid_token: UserAPIKeyAuth, end_use 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_max_budget") is not None: + valid_token.end_user_max_budget = end_user_params["end_user_max_budget"] 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 From e62ada9d63807442955f4c873696eddf8f07b369 Mon Sep 17 00:00:00 2001 From: Sricharan Anugandula Date: Wed, 9 Sep 2026 01:13:41 +0530 Subject: [PATCH 2/3] added test cases --- .../auth/test_custom_auth_end_user_budget.py | 84 +++++++++++++++++++ 1 file changed, 84 insertions(+) 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 cf1f665ad21..baa7e63ffe2 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 @@ -277,3 +277,87 @@ 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_copies_end_user_max_budget(): + # Regression: end_user_max_budget was the only limit field not copied by + # update_valid_token_with_end_user_params, so a budget from + # max_end_user_budget_id was silently dropped and budget reservation + # returned None, letting all concurrent first requests bypass the limit. + valid_token = UserAPIKeyAuth(token="test_token", end_user_id="customer-1") + + end_user_params = { + "end_user_id": "customer-1", + "end_user_max_budget": 0.000000001, + } + + result = update_valid_token_with_end_user_params(valid_token, end_user_params) + + assert result.end_user_max_budget == 0.000000001 + + +def test_update_valid_token_preserves_custom_auth_max_budget_when_db_has_none(): + # If custom auth sets end_user_max_budget and the DB budget table has no + # max_budget, the custom-auth-supplied value must survive the copy. + valid_token = UserAPIKeyAuth( + token="test_token", + end_user_id="customer-1", + end_user_max_budget=50.0, + ) + + end_user_params = { + "end_user_id": "customer-1", + # no end_user_max_budget key: DB found no value + } + + result = update_valid_token_with_end_user_params(valid_token, end_user_params) + + assert result.end_user_max_budget == 50.0 + + +@pytest.mark.asyncio +async def test_end_user_budget_counter_created_from_token_max_budget(): + # Regression: _get_end_user_budget_counter reads valid_token.end_user_max_budget + # first. When that field was always None (due to the copy bug), no counter was + # returned and budget reservation skipped enforcement entirely for new end users. + from litellm.proxy.spend_tracking.budget_reservation import ( + _get_end_user_budget_counter, + ) + + token = UserAPIKeyAuth( + token="test_token", + end_user_id="customer-1", + end_user_max_budget=0.000000001, + ) + + counter = await _get_end_user_budget_counter( + valid_token=token, + end_user_id="customer-1", + end_user_object=None, + ) + + assert counter is not None + assert counter.max_budget == 0.000000001 + assert counter.counter_key == "spend:end_user:customer-1" + + +@pytest.mark.asyncio +async def test_end_user_budget_counter_none_when_max_budget_missing(): + # Without a budget on the token or the end_user_object, no counter should + # be returned: the user is unrestricted and reservation is skipped correctly. + from litellm.proxy.spend_tracking.budget_reservation import ( + _get_end_user_budget_counter, + ) + + token = UserAPIKeyAuth( + token="test_token", + end_user_id="customer-1", + ) + + counter = await _get_end_user_budget_counter( + valid_token=token, + end_user_id="customer-1", + end_user_object=None, + ) + + assert counter is None From 59f67c5832b51b82dce24026d86c3c4c756f8198 Mon Sep 17 00:00:00 2001 From: Sricharan Anugandula Date: Wed, 9 Sep 2026 01:23:53 +0530 Subject: [PATCH 3/3] removed comments --- .../proxy/auth/test_custom_auth_end_user_budget.py | 12 ------------ 1 file changed, 12 deletions(-) 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 baa7e63ffe2..232036ff328 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 @@ -280,10 +280,6 @@ def test_update_valid_token_db_values_override_custom_auth_when_set(): def test_update_valid_token_copies_end_user_max_budget(): - # Regression: end_user_max_budget was the only limit field not copied by - # update_valid_token_with_end_user_params, so a budget from - # max_end_user_budget_id was silently dropped and budget reservation - # returned None, letting all concurrent first requests bypass the limit. valid_token = UserAPIKeyAuth(token="test_token", end_user_id="customer-1") end_user_params = { @@ -297,8 +293,6 @@ def test_update_valid_token_copies_end_user_max_budget(): def test_update_valid_token_preserves_custom_auth_max_budget_when_db_has_none(): - # If custom auth sets end_user_max_budget and the DB budget table has no - # max_budget, the custom-auth-supplied value must survive the copy. valid_token = UserAPIKeyAuth( token="test_token", end_user_id="customer-1", @@ -307,7 +301,6 @@ def test_update_valid_token_preserves_custom_auth_max_budget_when_db_has_none(): end_user_params = { "end_user_id": "customer-1", - # no end_user_max_budget key: DB found no value } result = update_valid_token_with_end_user_params(valid_token, end_user_params) @@ -317,9 +310,6 @@ def test_update_valid_token_preserves_custom_auth_max_budget_when_db_has_none(): @pytest.mark.asyncio async def test_end_user_budget_counter_created_from_token_max_budget(): - # Regression: _get_end_user_budget_counter reads valid_token.end_user_max_budget - # first. When that field was always None (due to the copy bug), no counter was - # returned and budget reservation skipped enforcement entirely for new end users. from litellm.proxy.spend_tracking.budget_reservation import ( _get_end_user_budget_counter, ) @@ -343,8 +333,6 @@ async def test_end_user_budget_counter_created_from_token_max_budget(): @pytest.mark.asyncio async def test_end_user_budget_counter_none_when_max_budget_missing(): - # Without a budget on the token or the end_user_object, no counter should - # be returned: the user is unrestricted and reservation is skipped correctly. from litellm.proxy.spend_tracking.budget_reservation import ( _get_end_user_budget_counter, )