mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-03 02:22:24 +00:00
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
This commit is contained in:
parent
06f6cfc5ae
commit
d6c2de6ab6
2 changed files with 47 additions and 0 deletions
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue