From 93cad9c2e6a9ba92f25493cee94bd34631950e02 Mon Sep 17 00:00:00 2001 From: Parman Mohammadalizadeh Date: Thu, 29 Oct 2026 10:17:32 +0100 Subject: [PATCH] fix(proxy): evict the member's cached keys when /team/member_update changes their limits A key's cached auth object carries the member's rpm and tpm limit from the membership budget, so keys already in use kept running without the new limit until their cache TTL ran out. Evict the auth objects of the member's keys in the team the way /key/update does Fixes #49041 --- .../management_endpoints/team_endpoints.py | 9 +++ .../test_team_endpoints.py | 67 +++++++++++++++++++ 2 files changed, 76 insertions(+) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 50b700e712d..653880e568e 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -4123,6 +4123,7 @@ async def team_member_update( litellm_proxy_admin_name, premium_user, prisma_client, + proxy_logging_obj, user_api_key_cache, ) @@ -4238,6 +4239,14 @@ async def team_member_update( team_id=data.team_id, user_api_key_cache=user_api_key_cache, ) + member_keys: Final = await _tokens_db(prisma_client).find_many( + where={"user_id": received_user_id, "team_id": data.team_id} + ) + await delete_cache_key_objects( + hashed_tokens=tuple(key.token for key in member_keys), + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) if role_change is not None: members_before_role_update, team_members = role_change diff --git a/tests/unit/proxy/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index c6d12ade2a0..0dc1482b5a7 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_endpoints.py @@ -17106,6 +17106,7 @@ async def test_team_member_update_invalidates_team_member_spend_state_when_budge } mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) @@ -17139,6 +17140,72 @@ async def test_team_member_update_invalidates_team_member_spend_state_when_budge assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:member-1:team-1") == 999.0 +@pytest.mark.asyncio +@pytest.mark.parametrize("limit_update", [{"rpm_limit": 2}, {"tpm_limit": 100}]) +async def test_team_member_update_evicts_cached_auth_of_the_members_existing_keys(monkeypatch, limit_update): + """Regression for #49041: a key's cached auth object carries the member's rpm/tpm limit, so a new + member limit must evict the auth objects of every key the member already uses in this team, while + the member's keys in other teams and other members' keys stay cached""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + key_rows = ( + SimpleNamespace(token="member-1-team-1-key-a", user_id="member-1", team_id="team-1"), + SimpleNamespace(token="member-1-team-1-key-b", user_id="member-1", team_id="team-1"), + SimpleNamespace(token="member-1-team-2-key", user_id="member-1", team_id="team-2"), + SimpleNamespace(token="member-2-team-1-key", user_id="member-2", team_id="team-1"), + ) + + async def find_key_rows(where): + return [row for row in key_rows if row.user_id == where["user_id"] and row.team_id == where["team_id"]] + + mock_prisma_client = MagicMock() + real_cache = UserApiKeyCache() + for row in key_rows: + await real_cache.async_set_cache(key=row.token, value=f"cached-auth-{row.token}") + + team_row = LiteLLM_TeamTable(team_id="team-1", metadata={}, members_with_roles=[]) + team_info_response = { + "team_info": team_row, + "team_memberships": [LiteLLM_TeamMembership(user_id="member-1", team_id="team-1", budget_id=None)], + } + + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(side_effect=find_key_rows) + + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", real_cache) + monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", DualCache()) + + mock_tx = AsyncMock() + mock_prisma_client.tx.return_value.__aenter__ = AsyncMock(return_value=mock_tx) + mock_prisma_client.tx.return_value.__aexit__ = AsyncMock(return_value=None) + + with ( + patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints.team_info", + AsyncMock(return_value=team_info_response), + ), + patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints.upsert_budget_and_membership", + AsyncMock(), + ), + ): + await team_member_update( + data=TeamMemberUpdateRequest(team_id="team-1", user_id="member-1", **limit_update), + http_request=MagicMock(), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user" + ), + ) + + assert await real_cache.async_get_cache(key="member-1-team-1-key-a") is None + assert await real_cache.async_get_cache(key="member-1-team-1-key-b") is None + assert await real_cache.async_get_cache(key="member-1-team-2-key") == "cached-auth-member-1-team-2-key" + assert await real_cache.async_get_cache(key="member-2-team-1-key") == "cached-auth-member-2-team-1-key" + + @pytest.mark.asyncio async def test_team_member_update_skips_invalidation_when_no_budget_fields_sent(monkeypatch): """A role-only update carries an empty budget_patch and touches no budget state,