mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-03 02:22:24 +00:00
fix(proxy): evict the member's cached user row on team member add
JWT auth caches the user row before it adds the user to the JWT's team, and admission checks the credential's team against the cached row on whichever worker takes the next request. On a two-worker gateway the credential minted for a newly joined team answered 403 "not in your team memberships" until the management-object TTL ran out, because /team/member_add only evicted the membership spend sentinel. The add now evicts the added members' cached user rows and broadcasts the eviction to the other workers, the way /team/member_delete already did The mint test now also covers a user SCIM deactivated after the cache last saw them active: the database read refuses the mint while the cached row still says active
This commit is contained in:
parent
167edf2769
commit
ce722ab1b3
4 changed files with 102 additions and 5 deletions
|
|
@ -279,11 +279,10 @@ async def load_active_user_by_id(
|
|||
catches every DB failure and re-raises a bare ``ValueError`` (a deleted user and a real outage look
|
||||
identical, the original error surviving only as ``__context__``), so the outage check walks the cause
|
||||
chain, and a missing user falls through to ``no_active_key`` rather than an opaque gateway fault.
|
||||
``source="database"`` reads the row from the database, never the cache, and leaves the fresh row in the
|
||||
cache for the requests the credential makes next: JWT auth caches the user it creates before it adds
|
||||
that user to the JWT's team and adding a member never evicts the cached row, so a credential minted
|
||||
off the cache would refuse the very first exchange as not a member. Every other caller keeps the cache
|
||||
read, so introspection, which a resource server may call per request, stays off the database."""
|
||||
``source="database"`` reads the row from the database, never the cache, so the credential mint refuses
|
||||
a user that a writer deactivated or deleted without evicting the cached row, and it leaves the fresh
|
||||
row in the cache for the requests the credential makes next. Every other caller keeps the cache read,
|
||||
so introspection, which a resource server may call per request, stays off the database."""
|
||||
from litellm.proxy._types import (
|
||||
ProxyException, # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
)
|
||||
|
|
|
|||
|
|
@ -3154,6 +3154,7 @@ async def team_member_add(
|
|||
|
||||
```
|
||||
"""
|
||||
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast
|
||||
from litellm.proxy.proxy_server import (
|
||||
litellm_proxy_admin_name,
|
||||
premium_user,
|
||||
|
|
@ -3248,6 +3249,10 @@ async def team_member_add(
|
|||
litellm_proxy_admin_name=litellm_proxy_admin_name,
|
||||
)
|
||||
|
||||
await evict_and_broadcast(
|
||||
cache_keys=tuple(sorted(user.user_id for user in updated_users)),
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
await _evict_created_membership_caches(
|
||||
user_ids=(tm.user_id for tm in updated_team_memberships),
|
||||
team_id=data.team_id,
|
||||
|
|
|
|||
|
|
@ -152,6 +152,28 @@ async def test_mint_reads_the_users_teams_from_the_database_not_a_stale_cached_r
|
|||
assert _decoded(minted).team_id == "team-a"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mint_refuses_a_user_scim_deactivated_after_the_cache_last_saw_them_active(fetch_teams, monkeypatch):
|
||||
"""SCIM deactivation writes the user row without evicting the cached copy, so a mint off the cache would
|
||||
keep issuing credentials for the management-object TTL. The mint reads the database row, so the
|
||||
deactivated user is refused on the first refresh after the deactivation."""
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
cache = UserApiKeyCache()
|
||||
await cache.async_set_cache(
|
||||
key="deactivated-user", value=_user(user_id="deactivated-user", teams=["team-a"]), model_type=LiteLLM_UserTable
|
||||
)
|
||||
prisma = MagicMock()
|
||||
prisma.db.litellm_usertable.find_unique = AsyncMock(
|
||||
return_value=_user(user_id="deactivated-user", teams=["team-a"], metadata={"scim_active": False})
|
||||
)
|
||||
monkeypatch.setattr(proxy_server, "user_api_key_cache", cache)
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", prisma)
|
||||
|
||||
assert await mint_proxy_credential("deactivated-user", "team-a") == "no_active_key"
|
||||
fetch_teams.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mint_refuses_a_team_the_user_is_not_on(load_user, fetch_teams):
|
||||
assert await mint_proxy_credential("u1", "team-c") == "not_a_member"
|
||||
|
|
|
|||
|
|
@ -13090,6 +13090,77 @@ async def test_team_member_add_audits_a_user_created_from_a_list_payload(monkeyp
|
|||
assert created_user_id not in mock_audit.call_args.kwargs["existing_user_ids"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_member_add_evicts_the_new_members_cached_user_row_on_every_worker(monkeypatch):
|
||||
"""Auth admits a team-bound credential off the teams list of the cached user row. The add wrote the
|
||||
new team to the database row only, so a worker still holding the old row refused the member's
|
||||
credential with 403 until the management-object TTL expired. The add now evicts the row here and
|
||||
broadcasts the eviction to the other workers, the way /team/member_delete already does."""
|
||||
from litellm.proxy._types import TeamMemberAddRequest
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.proxy.management_endpoints.team_endpoints import team_member_add
|
||||
|
||||
team_id = "team-b"
|
||||
user_id = "dev-1"
|
||||
cache = UserApiKeyCache()
|
||||
await cache.async_set_cache(
|
||||
key=user_id, value=LiteLLM_UserTable(user_id=user_id, teams=["team-a"]), model_type=LiteLLM_UserTable
|
||||
)
|
||||
broadcast = AsyncMock()
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", AsyncMock())
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "default_user_id")
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", broadcast
|
||||
)
|
||||
|
||||
updated_team = MagicMock()
|
||||
updated_team.model_dump.return_value = {
|
||||
"team_id": team_id,
|
||||
"members_with_roles": [{"user_id": user_id, "role": "user"}],
|
||||
}
|
||||
|
||||
async def fake_add_team_members_to_team(**kwargs):
|
||||
return updated_team, [LiteLLM_UserTable(user_id=user_id, teams=["team-a", team_id])], []
|
||||
|
||||
with (
|
||||
patch( # test-quality-ok: team_member_add has no injection seam for its prisma-backed helpers
|
||||
"litellm.proxy.management_endpoints.team_endpoints.get_team_object",
|
||||
new_callable=AsyncMock,
|
||||
return_value=LiteLLM_TeamTable(team_id=team_id, members_with_roles=[]),
|
||||
),
|
||||
patch( # test-quality-ok: team_member_add has no injection seam for its prisma-backed helpers
|
||||
"litellm.proxy.management_endpoints.team_endpoints._validate_team_member_add_permissions",
|
||||
new_callable=AsyncMock,
|
||||
),
|
||||
patch( # test-quality-ok: team_member_add has no injection seam for its prisma-backed helpers
|
||||
"litellm.proxy.management_endpoints.team_endpoints._validate_and_populate_member_user_info",
|
||||
new_callable=AsyncMock,
|
||||
),
|
||||
patch( # test-quality-ok: team_member_add has no injection seam for its prisma-backed helpers
|
||||
"litellm.proxy.management_endpoints.team_endpoints._resolve_existing_member_user_ids",
|
||||
new_callable=AsyncMock,
|
||||
return_value=frozenset({user_id}),
|
||||
),
|
||||
patch( # test-quality-ok: team_member_add has no injection seam for its prisma-backed helpers
|
||||
"litellm.proxy.management_endpoints.team_endpoints._add_team_members_to_team",
|
||||
side_effect=fake_add_team_members_to_team,
|
||||
),
|
||||
patch( # test-quality-ok: team_member_add has no injection seam for its prisma-backed helpers
|
||||
"litellm.proxy.management_endpoints.team_endpoints._create_team_member_add_audit_logs",
|
||||
new_callable=AsyncMock,
|
||||
),
|
||||
):
|
||||
await team_member_add(
|
||||
data=TeamMemberAddRequest(team_id=team_id, member=Member(user_id=user_id, role="user")),
|
||||
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-1"),
|
||||
)
|
||||
|
||||
assert await cache.async_get_cache(key=user_id, model_type=LiteLLM_UserTable) is None
|
||||
broadcast.assert_awaited_once_with(cache_key=user_id)
|
||||
|
||||
|
||||
def test_validate_member_user_id_provisioning_caps_the_ids_it_echoes_back():
|
||||
"""A large member list must not echo every id back in the error body."""
|
||||
from litellm.proxy.management_endpoints.team_endpoints import (
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue