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:
sharziki 2026-05-28 15:01:59 -04:00
parent 06f6cfc5ae
commit d6c2de6ab6
2 changed files with 47 additions and 0 deletions

View file

@ -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

View file

@ -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