From 7fff8405747d1764e729ecf64a8d3ea42be3953c Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 21 Jul 2026 18:06:21 -0700 Subject: [PATCH] fix(team): invalidate stale membership cache after bulk member update get_team_membership caches each member's LiteLLM_TeamMembership under team_membership:{user_id}:{team_id} with no TTL, and the legacy spend-check path in user_api_key_auth.py caches the same row under {team_id}_{user_id}. bulk_update_team_members only refreshed the team-row cache after committing budget/allowed_models writes, so a member could keep invoking a newly disallowed model or spending against a stale limit until those entries separately expired. Evict both cache keys for every member the transaction actually wrote. --- .../management_endpoints/team_endpoints.py | 28 ++++++ .../proxy/test_team_member_update.py | 93 +++++++++++++++++-- 2 files changed, 113 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 69b85f4774f..42c6eca67a9 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -180,6 +180,27 @@ async def _refresh_cached_team( ) +async def _invalidate_team_membership_caches( + team_id: str, + user_ids: List[str], + user_api_key_cache: Any, +) -> None: + """ + Invalidate the per-member membership caches after a budget/allowed_models + write to `litellm_teammembership` for `user_ids` in `team_id`. + + `get_team_membership` (auth_checks.py) caches `LiteLLM_TeamMembership` + under `team_membership:{user_id}:{team_id}` with no TTL, and the legacy + spend-check path in `user_api_key_auth.py` caches the same row under + `{team_id}_{user_id}`. Without invalidating both, a member keeps their + prior `allowed_models`/budget until the cache entry is separately evicted, + bypassing restrictions an admin just applied. + """ + for user_id in user_ids: + await user_api_key_cache.async_delete_cache(key="team_membership:{}:{}".format(user_id, team_id)) + await user_api_key_cache.async_delete_cache(key="{}_{}".format(team_id, user_id)) + + async def _verify_team_access( team_obj: LiteLLM_TeamTable, user_api_key_dict: UserAPIKeyAuth, @@ -3175,6 +3196,13 @@ async def bulk_update_team_members( proxy_logging_obj=proxy_logging_obj, ) + if budget_target_user_ids: + await _invalidate_team_membership_caches( + team_id=team_id, + user_ids=budget_target_user_ids, + user_api_key_cache=user_api_key_cache, + ) + successful_updates = [ TeamMemberUpdateResponse( team_id=team_id, diff --git a/tests/test_litellm/proxy/test_team_member_update.py b/tests/test_litellm/proxy/test_team_member_update.py index 09051e45fb5..9d689ceda41 100644 --- a/tests/test_litellm/proxy/test_team_member_update.py +++ b/tests/test_litellm/proxy/test_team_member_update.py @@ -18,6 +18,7 @@ from litellm.proxy._types import ( TeamMemberUpdateRequest, UserAPIKeyAuth, ) +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.management_endpoints.team_endpoints import ( bulk_update_team_members, team_member_update, @@ -217,15 +218,16 @@ class _FakeBulkDb: def _bulk_setup(monkeypatch, team_row, memberships): db = _FakeBulkDb(team_row, memberships) + user_api_key_cache = UserApiKeyCache() monkeypatch.setattr(proxy_server, "prisma_client", types.SimpleNamespace(db=db)) monkeypatch.setattr(proxy_server, "premium_user", False) - monkeypatch.setattr(proxy_server, "user_api_key_cache", object()) + monkeypatch.setattr(proxy_server, "user_api_key_cache", user_api_key_cache) monkeypatch.setattr(proxy_server, "proxy_logging_obj", object()) upsert_mock = AsyncMock() monkeypatch.setattr(team_endpoints, "_upsert_budget_and_membership", upsert_mock) refresh_mock = AsyncMock() monkeypatch.setattr(team_endpoints, "_refresh_cached_team", refresh_mock) - return db, upsert_mock, refresh_mock + return db, upsert_mock, refresh_mock, user_api_key_cache def _bulk_request(): @@ -249,7 +251,7 @@ async def test_bulk_update_delegates_budget_upsert_per_valid_member(monkeypatch) Member(user_id="user-3", role="user"), ], ) - db, upsert_mock, refresh_mock = _bulk_setup( + db, upsert_mock, refresh_mock, _cache = _bulk_setup( monkeypatch, team_row, memberships=[ @@ -282,6 +284,81 @@ async def test_bulk_update_delegates_budget_upsert_per_valid_member(monkeypatch) assert [member.user_id for member in response.successful_updates] == ["user-1", "user-2"] +@pytest.mark.asyncio +async def test_bulk_update_invalidates_membership_cache_for_written_members_only(monkeypatch): + """A budget/allowed_models write must evict the per-member + `team_membership:{user_id}:{team_id}` cache and the legacy + `{team_id}_{user_id}` cache for every member actually written, so the next + request re-reads the new limits instead of enforcing stale ones. Members + the patch didn't touch (no budget_patch, or filtered out as non-members) + must be left alone.""" + team_row = LiteLLM_TeamTable( + team_id="team-1234", + members_with_roles=[ + Member(user_id="user-1", role="user"), + Member(user_id="user-2", role="user"), + ], + ) + _db, _upsert_mock, _refresh_mock, cache = _bulk_setup( + monkeypatch, + team_row, + memberships=[ + LiteLLM_TeamMembership(user_id="user-1", team_id="team-1234", budget_id="bud-1"), + LiteLLM_TeamMembership(user_id="user-2", team_id="team-1234", budget_id="bud-2"), + ], + ) + + stale = "stale-cached-membership" + for user_id in ("user-1", "user-2", "user-3"): + await cache.async_set_cache(key="team_membership:{}:team-1234".format(user_id), value=stale) + await cache.async_set_cache(key="team-1234_{}".format(user_id), value=stale) + + await bulk_update_team_members( + team_id="team-1234", + data=BulkTeamMemberUpdateRequest( + user_ids=["user-1", "user-2"], + update_fields=TeamMemberBulkUpdateFields(allowed_models=["gpt-4"]), + ), + http_request=_bulk_request(), + user_api_key_dict=_ADMIN, + ) + + for user_id in ("user-1", "user-2"): + assert await cache.async_get_cache(key="team_membership:{}:team-1234".format(user_id)) is None + assert await cache.async_get_cache(key="team-1234_{}".format(user_id)) is None + + # user-3 was never part of this batch; its cache must be untouched. + assert await cache.async_get_cache(key="team_membership:user-3:team-1234") == stale + assert await cache.async_get_cache(key="team-1234_user-3") == stale + + +@pytest.mark.asyncio +async def test_bulk_update_role_only_change_does_not_touch_membership_cache(monkeypatch): + """A role-only patch (no budget fields) never writes litellm_teammembership, + so it must not evict the membership cache -- there's nothing stale to fix, + and clearing it would just cause an avoidable DB read on the next request.""" + team_row = LiteLLM_TeamTable( + team_id="team-1234", + members_with_roles=[Member(user_id="user-1", role="admin")], + ) + _db, _upsert_mock, _refresh_mock, cache = _bulk_setup(monkeypatch, team_row, memberships=[]) + + stale = "stale-cached-membership" + await cache.async_set_cache(key="team_membership:user-1:team-1234", value=stale) + + await bulk_update_team_members( + team_id="team-1234", + data=BulkTeamMemberUpdateRequest( + user_ids=["user-1"], + update_fields=TeamMemberBulkUpdateFields(role="user"), + ), + http_request=_bulk_request(), + user_api_key_dict=_ADMIN, + ) + + assert await cache.async_get_cache(key="team_membership:user-1:team-1234") == stale + + @pytest.mark.asyncio async def test_bulk_update_forwards_shared_default_only_for_members_still_on_it(monkeypatch): team_row = LiteLLM_TeamTable( @@ -292,7 +369,7 @@ async def test_bulk_update_forwards_shared_default_only_for_members_still_on_it( ], metadata={"team_member_budget_id": "default-bud"}, ) - _db, upsert_mock, _refresh = _bulk_setup( + _db, upsert_mock, _refresh, _cache = _bulk_setup( monkeypatch, team_row, memberships=[LiteLLM_TeamMembership(user_id="user-1", team_id="team-1234", budget_id="default-bud")], @@ -325,7 +402,7 @@ async def test_bulk_update_role_writes_team_row_once_and_refreshes_cache(monkeyp Member(user_id="user-3", role="user"), ], ) - db, upsert_mock, refresh_mock = _bulk_setup(monkeypatch, team_row, memberships=[]) + db, upsert_mock, refresh_mock, _cache = _bulk_setup(monkeypatch, team_row, memberships=[]) await bulk_update_team_members( team_id="team-1234", @@ -363,7 +440,7 @@ async def test_bulk_update_all_members_in_team_dedups_members(monkeypatch): Member(user_id="user-2", role="user"), ], ) - db, upsert_mock, _refresh = _bulk_setup(monkeypatch, team_row, memberships=[]) + db, upsert_mock, _refresh, _cache = _bulk_setup(monkeypatch, team_row, memberships=[]) response = await bulk_update_team_members( team_id="team-1234", @@ -386,7 +463,7 @@ async def test_bulk_update_reports_non_members_as_failed(monkeypatch): team_id="team-1234", members_with_roles=[Member(user_id="user-1", role="user")], ) - _db, upsert_mock, _refresh = _bulk_setup( + _db, upsert_mock, _refresh, _cache = _bulk_setup( monkeypatch, team_row, memberships=[LiteLLM_TeamMembership(user_id="user-1", team_id="team-1234", budget_id="bud-1")], @@ -415,7 +492,7 @@ async def test_bulk_update_explicit_null_duration_forwarded_as_patch(monkeypatch team_id="team-1234", members_with_roles=[Member(user_id="user-1", role="user")], ) - _db, upsert_mock, _refresh = _bulk_setup( + _db, upsert_mock, _refresh, _cache = _bulk_setup( monkeypatch, team_row, memberships=[LiteLLM_TeamMembership(user_id="user-1", team_id="team-1234", budget_id="bud-1")],