scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-20 00:11:50 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 75

refactor(auth): remove deep copy from update_valid_token_with_end_user_params

Browse source 
Eliminated the unnecessary deep copy of the valid token in the `update_valid_token_with_end_user_params` function. Updated comments to clarify that the function mutates the token in place, ensuring that request-scoped fields do not affect cached tokens. Removed a related test that verified the original token's immutability, as the implementation now relies on in-place updates.
This commit is contained in:
harish-berri 2026-04-29 19:17:24 +00:00
parent ac5b7aa336
commit ab01e18b9e
2 changed files with 2 additions and 32 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
litellm/proxy/auth/user_api_key_auth.py
View file

@ -303,7 +303,6 @@ async def user_api_key_auth_websocket(websocket: WebSocket):
def update_valid_token_with_end_user_params(
valid_token: UserAPIKeyAuth, end_user_params: dict
) -> UserAPIKeyAuth:
valid_token = valid_token.model_copy(deep=True)
valid_token.end_user_id = end_user_params.get("end_user_id")
# Only overwrite token fields when the DB-derived value is not None.
# This prevents DB lookups (where the budget table has no value set)
@ -1221,6 +1220,8 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
raise e
# update end-user params on valid token
# These can change per request - it's important to update them here
# This helper mutates the token in place. get_key_object() returns a
# copied token, so these request-scoped fields cannot poison the cache.
valid_token = update_valid_token_with_end_user_params(
valid_token=valid_token, end_user_params=end_user_params
)

31
tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py
View file

@ -139,34 +139,3 @@ 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_does_not_mutate_original_token():
"""
Request-scoped end-user limits must not mutate the cached UserAPIKeyAuth object.
"""
valid_token = UserAPIKeyAuth(
token="test_token",
end_user_id=None,
end_user_tpm_limit=None,
end_user_rpm_limit=None,
allowed_model_region=None,
)
end_user_params = {
"end_user_id": "attacker-user",
"end_user_tpm_limit": 1,
"end_user_rpm_limit": 1,
"allowed_model_region": "eu",
}
result = update_valid_token_with_end_user_params(valid_token, end_user_params)
assert result is not valid_token
assert result.end_user_id == "attacker-user"
assert result.end_user_tpm_limit == 1
assert result.end_user_rpm_limit == 1
assert result.allowed_model_region == "eu"
assert valid_token.end_user_id is None
assert valid_token.end_user_tpm_limit is None
assert valid_token.end_user_rpm_limit is None
assert valid_token.allowed_model_region is None