diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 54b9b26123e..bb15c698355 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -4133,6 +4133,7 @@ async def team_member_update( litellm_proxy_admin_name, premium_user, prisma_client, + proxy_logging_obj, user_api_key_cache, ) @@ -4248,6 +4249,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 697d9810284..e79c258c87e 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_endpoints.py @@ -17111,6 +17111,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) @@ -17144,6 +17145,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,