From ab01e18b9e8ab859b3e6fccf5828069cc695dba3 Mon Sep 17 00:00:00 2001 From: harish-berri Date: Wed, 29 Apr 2026 19:17:24 +0000 Subject: [PATCH] refactor(auth): remove deep copy from update_valid_token_with_end_user_params 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. --- litellm/proxy/auth/user_api_key_auth.py | 3 +- .../auth/test_custom_auth_end_user_budget.py | 31 ------------------- 2 files changed, 2 insertions(+), 32 deletions(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index ce93fa2c07a..5a10235ac5a 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -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 ) 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 6bba59ae8b0..3af0cac6fd3 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 @@ -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