Merge pull request #41707 from BerriAI/litellm_jwt_mapping_cache_evict_on_bulk_key_delete

fix(proxy): evict jwt key mapping cache on user, team, org, and bulk key deletion
This commit is contained in:
ryan-crabbe-berri 2026-09-18 08:41:31 -07:00 committed by GitHub
commit 4f70b88a1f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 548 additions and 39 deletions

View file

@ -3637,6 +3637,22 @@ async def get_jwt_key_mapping_cache_keys_for_token(
return tuple(jwt_key_mapping_cache_key(m.jwt_claim_name, m.jwt_claim_value, m.jwt_issuer) for m in mappings)
class _TokenInFilter(TypedDict):
token: ReadOnly[Mapping[str, Sequence[str]]]
async def get_jwt_key_mapping_cache_keys_for_tokens(
hashed_tokens: Sequence[str],
prisma_client: PrismaClient,
) -> tuple[str, ...]:
"""Cache keys of every JWT claim mapped to any of the given virtual keys."""
if not hashed_tokens:
return ()
token_filter: Final[_TokenInFilter] = {"token": {"in": tuple(hashed_tokens)}}
mappings: Final = await _jwt_key_mapping_table(JWTKeyMappingRepository(prisma_client)).find_many(where=token_filter)
return tuple(jwt_key_mapping_cache_key(m.jwt_claim_name, m.jwt_claim_value, m.jwt_issuer) for m in mappings)
@log_db_metrics
async def get_jwt_key_mapping_object(
jwt_claim_name: str,

View file

@ -23,12 +23,18 @@ from typing import Any, Final, Literal, Protocol, cast, overload
import fastapi
from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
from pydantic import TypeAdapter, ValidationError
from typing_extensions import ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
from litellm.proxy._types import *
from litellm.proxy.auth.auth_checks import get_team_object, get_user_object
from litellm.proxy.auth.auth_checks import (
delete_cache_key_objects,
get_jwt_key_mapping_cache_keys_for_tokens,
get_team_object,
get_user_object,
)
from litellm.proxy.auth.password_policy import validate_password_policy
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast
@ -126,6 +132,10 @@ def _verification_token_table(
return token_table
class _UserIdInFilter(TypedDict):
user_id: ReadOnly[Mapping[str, Sequence[str]]]
def _organization_membership_table(
prisma_client: "PrismaClient | None",
) -> "TableActions[prisma_models.LiteLLM_OrganizationMembership]":
@ -2345,6 +2355,8 @@ async def delete_user(
create_audit_log_for_update,
litellm_proxy_admin_name,
prisma_client,
proxy_logging_obj,
user_api_key_cache,
)
if prisma_client is None:
@ -2471,7 +2483,20 @@ async def delete_user(
# End of Audit logging
## DELETE ASSOCIATED KEYS
await _verification_token_table(prisma_client).delete_many(where={"user_id": {"in": data.user_ids}})
key_filter: Final[_UserIdInFilter] = {"user_id": {"in": data.user_ids}}
keys_to_delete: Final = await _verification_token_table(prisma_client).find_many(where=key_filter)
hashed_tokens_to_delete: Final = tuple(key.token for key in keys_to_delete)
jwt_mapping_cache_keys: Final = await get_jwt_key_mapping_cache_keys_for_tokens(
hashed_tokens=hashed_tokens_to_delete,
prisma_client=prisma_client,
)
await _verification_token_table(prisma_client).delete_many(where=key_filter)
await delete_cache_key_objects(
hashed_tokens=hashed_tokens_to_delete,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
await evict_and_broadcast(cache_keys=jwt_mapping_cache_keys, user_api_key_cache=user_api_key_cache)
## DELETE ASSOCIATED INVITATION LINKS
await _invitation_link_table(prisma_client).delete_many(

View file

@ -26,13 +26,21 @@ from typing import (
import fastapi
from fastapi import APIRouter, Depends, HTTPException, Request, status
from pydantic import TypeAdapter
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
from litellm.proxy._types import *
from litellm.proxy.auth.auth_checks import can_user_call_model, get_user_object
from litellm.proxy.auth.auth_checks import (
can_user_call_model,
delete_cache_key_objects,
get_jwt_key_mapping_cache_keys_for_tokens,
get_user_object,
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.management_endpoints.budget_management_endpoints import (
new_budget,
update_budget,
@ -52,7 +60,7 @@ from litellm.proxy.management_helpers.utils import (
get_new_internal_user_defaults,
management_endpoint_wrapper,
)
from litellm.proxy.utils import PrismaClient
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.repositories.budget_repository import BudgetRepository
from litellm.repositories.object_permission_repository import ObjectPermissionRepository
from litellm.repositories.organization_repository import OrganizationRepository
@ -79,6 +87,7 @@ if TYPE_CHECKING:
)
from prisma.models import LiteLLM_OrganizationTable as PrismaOrganizationTable
from prisma.models import LiteLLM_UserTable as PrismaUserTable
from prisma.models import LiteLLM_VerificationToken as PrismaVerificationToken
async def _enterprise_license_required(
@ -168,9 +177,15 @@ class _TeamTableClient(Protocol):
class _VerificationTokenTableClient(Protocol):
async def find_many(self, where: Mapping[str, object] | None = None) -> "Sequence[PrismaVerificationToken]": ...
async def delete_many(self, where: Mapping[str, object]) -> int: ...
class _OrganizationIdFilter(TypedDict):
organization_id: ReadOnly[str]
class _ObjectPermissionTxClient(Protocol):
async def upsert(
self, where: Mapping[str, object], data: Mapping[str, object]
@ -961,7 +976,7 @@ async def delete_organization(
- organization_ids: List[str] - The organization ids to delete.
"""
from litellm.proxy.proxy_server import prisma_client
from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache
if prisma_client is None:
raise HTTPException(
@ -983,8 +998,12 @@ async def delete_organization(
await _table(OrganizationMembershipRepository(prisma_client)).delete_many(
where={"organization_id": organization_id}
)
# delete all keys in the organization
await _table(VerificationTokenRepository(prisma_client)).delete_many(where={"organization_id": organization_id})
await _delete_organization_keys(
organization_id=organization_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
# delete the organization
deleted_org = await _table(OrganizationRepository(prisma_client)).delete(
where={"organization_id": organization_id},
@ -1000,6 +1019,28 @@ async def delete_organization(
return deleted_orgs
async def _delete_organization_keys(
organization_id: str,
prisma_client: PrismaClient,
user_api_key_cache: UserApiKeyCache,
proxy_logging_obj: ProxyLogging | None,
) -> None:
key_filter: Final[_OrganizationIdFilter] = {"organization_id": organization_id}
keys_to_delete: Final = await _table(VerificationTokenRepository(prisma_client)).find_many(where=key_filter)
hashed_tokens_to_delete: Final = tuple(key.token for key in keys_to_delete)
jwt_mapping_cache_keys: Final = await get_jwt_key_mapping_cache_keys_for_tokens(
hashed_tokens=hashed_tokens_to_delete,
prisma_client=prisma_client,
)
await _table(VerificationTokenRepository(prisma_client)).delete_many(where=key_filter)
await delete_cache_key_objects(
hashed_tokens=hashed_tokens_to_delete,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
await evict_and_broadcast(cache_keys=jwt_mapping_cache_keys, user_api_key_cache=user_api_key_cache)
@router.get(
"/organization/list",
tags=["organization management"],

View file

@ -99,6 +99,7 @@ from litellm.proxy.auth.auth_checks import (
can_org_access_model,
delete_cache_key_objects,
delete_cache_team_object,
get_jwt_key_mapping_cache_keys_for_tokens,
get_org_object,
get_team_membership,
get_team_object,
@ -110,6 +111,7 @@ from litellm.proxy.auth.auth_utils import (
enforce_output_token_estimates_are_admin_only,
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast
from litellm.proxy.common_utils.callback_utils import encrypt_callback_vars
from litellm.proxy.common_utils.json_merge_patch import apply_json_merge_patch
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
@ -3524,7 +3526,6 @@ async def team_member_delete(
}'
```
"""
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast
from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache
if prisma_client is None:
@ -3626,6 +3627,10 @@ async def team_member_delete(
"team_id": data.team_id,
}
)
jwt_mapping_cache_keys: Final = await get_jwt_key_mapping_cache_keys_for_tokens(
hashed_tokens=tuple(key.token for key in keys_to_delete),
prisma_client=prisma_client,
)
if removed_team_members:
await _team_tx_db(tx).update(
@ -3674,6 +3679,7 @@ async def team_member_delete(
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
await evict_and_broadcast(cache_keys=jwt_mapping_cache_keys, user_api_key_cache=user_api_key_cache)
await evict_and_broadcast(cache_keys=tuple(sorted(user_ids_to_delete)), user_api_key_cache=user_api_key_cache)
for user_id in sorted(user_ids_to_delete):
await invalidate_team_member_spend_state(
@ -4264,6 +4270,10 @@ async def delete_team(
)
keys_to_delete: Final = await _tokens_db(prisma_client).find_many(where={"team_id": {"in": data.team_ids}})
jwt_mapping_cache_keys: Final = await get_jwt_key_mapping_cache_keys_for_tokens(
hashed_tokens=tuple(key.token for key in keys_to_delete),
prisma_client=prisma_client,
)
if keys_to_delete:
await _persist_deleted_verification_tokens(
@ -4280,6 +4290,7 @@ async def delete_team(
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
await evict_and_broadcast(cache_keys=jwt_mapping_cache_keys, user_api_key_cache=user_api_key_cache)
## DELETE ASSOCIATED BYOK MODELS
# Runs before the team rows are deleted so a mid-flight failure never leaves

View file

@ -28,7 +28,7 @@ from litellm.proxy._types import (
MemberDeleteRequest,
UserAPIKeyAuth,
)
from litellm.proxy.auth.auth_checks import delete_cache_key_objects
from litellm.proxy.auth.auth_checks import delete_cache_key_objects, get_jwt_key_mapping_cache_keys_for_tokens
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.hooks.user_management_event_hooks import UserManagementEventHooks
@ -94,12 +94,20 @@ class _TeamRemoval:
removed: frozenset[str]
matched: frozenset[int]
deleted_key_tokens: tuple[str, ...]
jwt_mapping_cache_keys: tuple[str, ...]
@dataclass(frozen=True, slots=True)
class _UserBatchDeletion:
removals: Mapping[str, _TeamRemoval]
deleted_key_tokens: tuple[str, ...]
jwt_mapping_cache_keys: tuple[str, ...]
@dataclass(frozen=True, slots=True)
class _DeletedKeys:
tokens: tuple[str, ...]
jwt_mapping_cache_keys: tuple[str, ...]
def _team_not_found(team_id: str) -> ManagementProblem:
@ -237,6 +245,10 @@ async def _remove_members_from_team(
if any(_addresses_member(m, r) for m in removed_members) or any(_addresses_user(u, r) for u in stale_rows)
)
keys: Final = await _token_tx_db(tx).find_many(where=_team_users_filter(team_id, cleanup_ids))
jwt_mapping_cache_keys: Final = await get_jwt_key_mapping_cache_keys_for_tokens(
hashed_tokens=tuple(k.token for k in keys),
prisma_client=prisma_client,
)
if removed_members:
roster_data: Final[_RosterData] = {
@ -265,6 +277,7 @@ async def _remove_members_from_team(
removed=cleanup_ids,
matched=matched,
deleted_key_tokens=tuple(k.token for k in keys),
jwt_mapping_cache_keys=jwt_mapping_cache_keys,
)
@ -322,6 +335,7 @@ async def bulk_remove_team_members(
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
await evict_and_broadcast(cache_keys=removal.jwt_mapping_cache_keys, user_api_key_cache=user_api_key_cache)
_emit_team_members_metric(removal.team)
matched: Final = frozenset(kept_indexes[j] for j in removal.matched)
@ -368,8 +382,12 @@ async def _delete_user_rows(
user_ids: frozenset[str],
user_api_key_dict: UserAPIKeyAuth,
litellm_changed_by: str | None,
) -> tuple[str, ...]:
) -> _DeletedKeys:
keys: Final = await _token_tx_db(tx).find_many(where=_in_filter("user_id", user_ids))
jwt_mapping_cache_keys: Final = await get_jwt_key_mapping_cache_keys_for_tokens(
hashed_tokens=tuple(k.token for k in keys),
prisma_client=prisma_client,
)
if keys:
await _persist_deleted_verification_tokens(
keys=keys, # pyright: ignore[reportArgumentType] # generated row model carries the same columns as LiteLLM_VerificationToken
@ -389,7 +407,7 @@ async def _delete_user_rows(
await _org_membership_tx_db(tx).delete_many(where=_in_filter("user_id", user_ids))
await _membership_tx_db(tx).delete_many(where=_in_filter("user_id", user_ids))
await _user_tx_db(tx).delete_many(where=_in_filter("user_id", user_ids))
return tuple(k.token for k in keys)
return _DeletedKeys(tokens=tuple(k.token for k in keys), jwt_mapping_cache_keys=jwt_mapping_cache_keys)
async def _delete_users_tx(
@ -423,12 +441,14 @@ async def _delete_users_tx(
for tid in team_ids
}
)
deleted_key_tokens: Final = await _delete_user_rows(
deleted_keys: Final = await _delete_user_rows(
prisma_client, tx, frozenset(u.user_id for u in users), user_api_key_dict, litellm_changed_by
)
return _UserBatchDeletion(
removals=removals,
deleted_key_tokens=deleted_key_tokens + tuple(t for r in removals.values() for t in r.deleted_key_tokens),
deleted_key_tokens=deleted_keys.tokens + tuple(t for r in removals.values() for t in r.deleted_key_tokens),
jwt_mapping_cache_keys=deleted_keys.jwt_mapping_cache_keys
+ tuple(k for r in removals.values() for k in r.jwt_mapping_cache_keys),
)
@ -454,6 +474,7 @@ async def _delete_users(
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
await evict_and_broadcast(cache_keys=deletion.jwt_mapping_cache_keys, user_api_key_cache=user_api_key_cache)
await evict_and_broadcast(cache_keys=sorted(user_ids), user_api_key_cache=user_api_key_cache)
for removal in deletion.removals.values():
_emit_team_members_metric(removal.team)
@ -534,7 +555,7 @@ async def bulk_delete_users(
litellm_changed_by,
)
if candidates
else _UserBatchDeletion(removals=MappingProxyType({}), deleted_key_tokens=())
else _UserBatchDeletion(removals=MappingProxyType({}), deleted_key_tokens=(), jwt_mapping_cache_keys=())
)
def result(index: int, user_id: str) -> UserDeleteResult:

View file

@ -0,0 +1,27 @@
"""LiteLLM_JWTKeyMapping test doubles for the bulk key deletion paths."""
from __future__ import annotations
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
@dataclass(frozen=True, slots=True)
class JWTMappingRow:
token: str
jwt_claim_name: str
jwt_claim_value: str
jwt_issuer: str | None = None
class CascadingJWTMappingTable:
"""Mapping rows that LiteLLM_JWTKeyMapping_token_fkey drops when their key row is deleted."""
def __init__(self, rows: Sequence[JWTMappingRow]) -> None:
self.rows: tuple[JWTMappingRow, ...] = tuple(rows)
async def find_many(self, where: Mapping[str, Mapping[str, Sequence[str]]]) -> list[JWTMappingRow]:
return [row for row in self.rows if row.token in where["token"]["in"]]
def cascade(self, deleted_tokens: Sequence[str]) -> None:
self.rows = tuple(row for row in self.rows if row.token not in deleted_tokens)

View file

@ -4,11 +4,10 @@ from types import SimpleNamespace
from typing import Final
import pytest
from fastapi.testclient import TestClient
from fastapi import HTTPException
from fastapi.testclient import TestClient
from pytest_mock import MockerFixture
from litellm.proxy._types import (
LiteLLM_UserTableFiltered,
LitellmUserRoles,
@ -27,6 +26,10 @@ from litellm.proxy.management_endpoints.internal_user_endpoints import (
ui_view_users,
)
from litellm.proxy.proxy_server import app
from tests.test_litellm.proxy.management_endpoints.jwt_key_mapping_doubles import (
CascadingJWTMappingTable,
JWTMappingRow,
)
client = TestClient(app)
@ -2627,6 +2630,9 @@ async def test_delete_user_cleans_up_created_by_invitation_links(mocker):
)
# Mock all delete_many calls
mock_prisma_client.db.litellm_verificationtoken.find_many = mocker.AsyncMock(
return_value=[]
)
mock_prisma_client.db.litellm_verificationtoken.delete_many = mocker.AsyncMock(
return_value=0
)
@ -2676,6 +2682,84 @@ async def test_delete_user_cleans_up_created_by_invitation_links(mocker):
assert condition[field] == {"in": ["admin-creator"]}
@pytest.mark.asyncio
async def test_delete_user_evicts_jwt_key_mapping_cache_of_its_keys(mocker):
"""/user/delete bulk-deletes the user's keys without going through /key/delete, so the
jwt_key_mapping cache entries pointing at those keys must be evicted here too. A surviving
entry keeps resolving the deleted token hash until the mapping cache TTL expires: the deleted
identity is either still served through the stale key cache or 401s on every JWT call, and it is
never re-registered (LIT-5387).
The FK cascade drops the mapping rows with the key rows, so the cache keys have to be read
before the delete: reading them afterwards finds nothing to evict.
"""
from litellm.proxy._types import DeleteUserRequest, UserAPIKeyAuth
from litellm.proxy.auth.auth_checks import jwt_key_mapping_cache_key
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.management_endpoints.internal_user_endpoints import delete_user
global_cache_key: Final = jwt_key_mapping_cache_key("sub", "jwt-user", None)
issuer_cache_key: Final = jwt_key_mapping_cache_key("sub", "jwt-user", "https://issuer.example")
unrelated_cache_key: Final = jwt_key_mapping_cache_key("sub", "other-user", None)
jwt_table: Final = CascadingJWTMappingTable(
[
JWTMappingRow("hashed-jwt-key", "sub", "jwt-user"),
JWTMappingRow("hashed-issuer-key", "sub", "jwt-user", "https://issuer.example"),
JWTMappingRow("hashed-unrelated-key", "sub", "other-user"),
]
)
cache: Final = UserApiKeyCache()
for cache_key, hashed_token in (
(global_cache_key, "hashed-jwt-key"),
(issuer_cache_key, "hashed-issuer-key"),
(unrelated_cache_key, "hashed-unrelated-key"),
):
cache.set_cache(key=cache_key, value=hashed_token)
cache.set_cache(key=hashed_token, value=UserAPIKeyAuth(token=hashed_token))
user_row: Final = mocker.MagicMock()
user_row.user_id = "jwt-user"
user_row.user_email = "jwt-user@example.com"
user_row.teams = []
user_row.model_dump_json.return_value = "{}"
user_row.model_dump.return_value = {"user_id": "jwt-user", "user_email": "jwt-user@example.com", "teams": []}
mock_prisma_client: Final = mocker.MagicMock()
mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(return_value=user_row)
mock_prisma_client.db.litellm_teamtable.find_many = mocker.AsyncMock(return_value=[])
mock_prisma_client.db.litellm_jwtkeymapping = jwt_table
mock_prisma_client.db.litellm_verificationtoken.find_many = mocker.AsyncMock(
return_value=[SimpleNamespace(token="hashed-jwt-key"), SimpleNamespace(token="hashed-issuer-key")]
)
async def cascading_delete_many(where):
jwt_table.cascade(("hashed-jwt-key", "hashed-issuer-key"))
return 2
mock_prisma_client.db.litellm_verificationtoken.delete_many = mocker.AsyncMock(side_effect=cascading_delete_many)
mock_prisma_client.db.litellm_invitationlink.delete_many = mocker.AsyncMock(return_value=0)
mock_prisma_client.db.litellm_organizationmembership.delete_many = mocker.AsyncMock(return_value=0)
mock_prisma_client.db.litellm_teammembership.delete_many = mocker.AsyncMock(return_value=0)
mock_prisma_client.db.litellm_usertable.delete_many = mocker.AsyncMock(return_value=1)
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) # test-quality-ok: substitute the database dependency
mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache
mocker.patch("litellm.proxy.proxy_server.proxy_logging_obj", None) # test-quality-ok: delete_user reads it off proxy_server at call time
await delete_user(
data=DeleteUserRequest(user_ids=["jwt-user"]),
user_api_key_dict=UserAPIKeyAuth(user_id="proxy-admin", user_role=LitellmUserRoles.PROXY_ADMIN),
)
assert cache.get_cache(key=global_cache_key) is None
assert cache.get_cache(key=issuer_cache_key) is None
assert cache.get_cache(key="hashed-jwt-key") is None
assert cache.get_cache(key="hashed-issuer-key") is None
assert cache.get_cache(key=unrelated_cache_key) == "hashed-unrelated-key"
assert cache.get_cache(key="hashed-unrelated-key") is not None
assert [row.token for row in jwt_table.rows] == ["hashed-unrelated-key"]
@pytest.mark.asyncio
async def test_delete_user_rejects_org_admin_deleting_outside_scope(mocker):
"""Regression: an org admin of org-A must not be able to delete a user

View file

@ -5244,7 +5244,10 @@ async def test_delete_verification_tokens_evicts_jwt_key_mapping_cache(monkeypat
virtual_key_mapping_cache_ttl expires, instead of auto-registering again.
"""
jwt_table = _CascadingJWTMappingTable(
[_JWTMappingRow("hashed-token-1", "email", "user@example.com")]
[
_JWTMappingRow("hashed-token-1", "email", "user@example.com"),
_JWTMappingRow("hashed-token-1", "email", "user@example.com", "https://issuer.example"),
]
)
key1 = LiteLLM_VerificationToken(
@ -5302,7 +5305,10 @@ async def test_delete_verification_tokens_evicts_jwt_key_mapping_cache(monkeypat
),
)
assert recording_evict.cache_keys == (jwt_key_mapping_cache_key("email", "user@example.com", None),)
assert recording_evict.cache_keys == (
jwt_key_mapping_cache_key("email", "user@example.com", None),
jwt_key_mapping_cache_key("email", "user@example.com", "https://issuer.example"),
)
@pytest.mark.asyncio

View file

@ -1,7 +1,6 @@
import asyncio
import json
from litellm._uuid import uuid
from types import MappingProxyType
from types import MappingProxyType, SimpleNamespace
from typing import Final, Mapping, Optional, cast
from unittest.mock import AsyncMock, MagicMock, patch
@ -9,6 +8,11 @@ import pytest
from fastapi import HTTPException
from fastapi.testclient import TestClient
from litellm._uuid import uuid
from tests.test_litellm.proxy.management_endpoints.jwt_key_mapping_doubles import (
CascadingJWTMappingTable,
JWTMappingRow,
)
@pytest.mark.asyncio
@ -499,9 +503,10 @@ async def test_organization_info_includes_user_email(monkeypatch):
"""
Test that GET /organization/info returns user_email in members list.
"""
from litellm.proxy._types import LiteLLM_OrganizationMembershipTable
from datetime import datetime
from litellm.proxy._types import LiteLLM_OrganizationMembershipTable
# Simulate a membership row with a nested user object that has user_email
raw_membership = {
"user_id": "user_abc",
@ -573,6 +578,10 @@ async def test_organization_member_add_rejects_unauthorized_caller(patched_org_p
# ``organization_member_add`` catches HTTPException in its
# catch-all and re-wraps as ProxyException with the original status
# code preserved.
from unittest.mock import Mock
from fastapi import Request
from litellm.proxy._types import (
OrganizationMemberAddRequest,
OrgMember,
@ -581,9 +590,6 @@ async def test_organization_member_add_rejects_unauthorized_caller(patched_org_p
from litellm.proxy.management_endpoints.organization_endpoints import (
organization_member_add,
)
from unittest.mock import Mock
from fastapi import Request
data = OrganizationMemberAddRequest(
organization_id="org-victim",
@ -1438,3 +1444,61 @@ def test_organization_routes_reach_their_handler_with_enterprise_license(monkeyp
assert any(
message in response.text for message in (CommonProxyErrors.db_not_connected_error.value, "No db connected")
)
@pytest.mark.asyncio
async def test_delete_organization_evicts_the_cache_of_the_keys_it_deletes(monkeypatch):
"""/organization/delete bulk-deletes the org's keys without going through /key/delete, so the
key objects and the jwt_key_mapping entries (issuer-scoped ones included) pointing at them
must be evicted here, or a deleted key keeps authenticating and a JWT identity keeps resolving
a token hash that no longer exists until the TTLs expire. The FK cascade drops the mapping
rows with the key rows, so the cache keys have to be read before the delete (LIT-5387)."""
from litellm.proxy._types import DeleteOrganizationRequest, LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.auth.auth_checks import jwt_key_mapping_cache_key
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.management_endpoints.organization_endpoints import delete_organization
doomed_cache_keys: Final = (
"hashed-org-key",
jwt_key_mapping_cache_key("sub", "svc-account", None),
jwt_key_mapping_cache_key("sub", "svc-account", "https://issuer.example"),
)
kept_cache_keys: Final = ("hashed-other-key", jwt_key_mapping_cache_key("sub", "other-account", None))
kept_row: Final = JWTMappingRow("hashed-other-key", "sub", "other-account")
jwt_table: Final = CascadingJWTMappingTable(
[
JWTMappingRow("hashed-org-key", "sub", "svc-account"),
JWTMappingRow("hashed-org-key", "sub", "svc-account", "https://issuer.example"),
kept_row,
]
)
cache: Final = UserApiKeyCache()
for cache_key in (*doomed_cache_keys, *kept_cache_keys):
cache.set_cache(key=cache_key, value={"retained": True})
prisma_client: Final = AsyncMock()
prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
return_value=[SimpleNamespace(token="hashed-org-key")]
)
async def cascading_delete_many(where):
jwt_table.cascade(("hashed-org-key",))
return 1
prisma_client.db.litellm_verificationtoken.delete_many = AsyncMock(side_effect=cascading_delete_many)
prisma_client.db.litellm_jwtkeymapping = jwt_table
prisma_client.db.litellm_organizationtable.delete = AsyncMock(return_value=MagicMock())
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True, raising=False)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client)
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache)
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", None)
await delete_organization(
data=DeleteOrganizationRequest(organization_ids=["org-doomed"]),
user_api_key_dict=UserAPIKeyAuth(api_key="sk-admin", user_role=LitellmUserRoles.PROXY_ADMIN),
)
assert all(cache.get_cache(key=cache_key) is None for cache_key in doomed_cache_keys)
assert all(cache.get_cache(key=cache_key) == {"retained": True} for cache_key in kept_cache_keys)
assert jwt_table.rows == (kept_row,)

View file

@ -12,8 +12,6 @@ from fastapi.testclient import TestClient
from pydantic import ValidationError
from litellm._uuid import uuid
from litellm.proxy._types import UserAPIKeyAuth # Import UserAPIKeyAuth
from litellm.proxy._types import (
LiteLLM_BudgetTable,
LiteLLM_BudgetTableFull,
@ -33,14 +31,12 @@ from litellm.proxy._types import (
TeamMemberAddRequest,
TeamMemberUpdateRequest,
UpdateTeamRequest,
UserAPIKeyAuth, # Import UserAPIKeyAuth
)
from litellm.proxy.management_endpoints.team_endpoints import (
user_api_key_auth, # Assuming this dependency is needed
)
from litellm.proxy.management_endpoints.team_endpoints import (
_STRIP_DELETED_TEAM_FROM_USERS_SQL,
GetTeamMemberPermissionsResponse,
UpdateTeamMemberPermissionsRequest,
_STRIP_DELETED_TEAM_FROM_USERS_SQL,
_persist_deleted_team_records,
_save_deleted_team_records,
_transform_teams_to_deleted_records,
@ -56,6 +52,7 @@ from litellm.proxy.management_endpoints.team_endpoints import (
team_member_delete,
team_member_update,
update_team,
user_api_key_auth, # Assuming this dependency is needed
validate_team_org_change,
)
from litellm.proxy.management_helpers.access_group_team_sync import (
@ -71,6 +68,10 @@ from litellm.types.proxy.management_endpoints.team_endpoints import (
BulkTeamMemberAddResponse,
TeamMemberAddResult,
)
from tests.test_litellm.proxy.management_endpoints.jwt_key_mapping_doubles import (
CascadingJWTMappingTable,
JWTMappingRow,
)
# Setup TestClient
client = TestClient(app)
@ -2788,7 +2789,7 @@ async def test_upsert_team_member_budget_table_existing_budget():
"""
from unittest.mock import AsyncMock, MagicMock, patch
from litellm.proxy._types import LitellmUserRoles, LiteLLM_TeamTable, UserAPIKeyAuth
from litellm.proxy._types import LiteLLM_TeamTable, LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.management_endpoints.team_endpoints import (
TeamMemberBudgetHandler,
)
@ -2849,7 +2850,7 @@ async def test_upsert_team_member_budget_table_no_existing_budget():
"""
from unittest.mock import AsyncMock, MagicMock, patch
from litellm.proxy._types import LitellmUserRoles, LiteLLM_TeamTable, UserAPIKeyAuth
from litellm.proxy._types import LiteLLM_TeamTable, LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.management_endpoints.team_endpoints import (
TeamMemberBudgetHandler,
)
@ -6092,9 +6093,9 @@ async def test_new_team_standalone_validates_against_user_models(monkeypatch):
- Team is created WITHOUT organization_id and models=['gpt-4']
- Expected: Should fail with "Model not in allowed user models"
"""
import litellm
from fastapi import Request
import litellm
from litellm.proxy._types import NewTeamRequest, ProxyException, UserAPIKeyAuth
from litellm.proxy.management_endpoints.team_endpoints import new_team
@ -9180,6 +9181,154 @@ async def test_team_member_delete_persists_deleted_keys(monkeypatch):
assert cache.get_cache(key="unrelated-key") == {"retained": True}
def _seed_jwt_mapping_cache(cache, mapping_rows):
from litellm.proxy.auth.auth_checks import jwt_key_mapping_cache_key
cache_keys = tuple(
jwt_key_mapping_cache_key(row.jwt_claim_name, row.jwt_claim_value, row.jwt_issuer) for row in mapping_rows
)
for cache_key, row in zip(cache_keys, mapping_rows):
cache.set_cache(key=cache_key, value=row.token)
return cache_keys
@pytest.mark.asyncio
async def test_team_member_delete_evicts_jwt_key_mapping_cache_of_the_keys_it_deletes(monkeypatch):
"""The member's team keys are deleted in bulk here, not through /key/delete, so the
jwt_key_mapping cache entries pointing at them must be evicted here too, or every JWT call
from that identity resolves the deleted token hash and 401s until the mapping TTL expires.
The FK cascade drops the mapping rows with the key rows, so the cache keys have to be read
before the delete (LIT-5387)."""
from litellm.proxy._types import TeamMemberDeleteRequest
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.management_endpoints.key_management_endpoints import LiteLLM_VerificationToken
doomed_rows: Final = (
JWTMappingRow("hashed-token-1", "sub", "user-123"),
JWTMappingRow("hashed-token-1", "sub", "user-123", "https://issuer.example"),
)
kept_row: Final = JWTMappingRow("hashed-other-key", "sub", "user-999")
jwt_table: Final = CascadingJWTMappingTable([*doomed_rows, kept_row])
team = LiteLLM_TeamTable(
team_id="team-1",
team_alias="test-team",
members_with_roles=[Member(user_id="user-123", role="admin")],
metadata={},
model_max_budget={},
model_spend={},
)
key1 = LiteLLM_VerificationToken(token="hashed-token-1", user_id="user-123", team_id="team-1")
mock_prisma_client = AsyncMock()
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team)
mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(
return_value=[MagicMock(user_id="user-123", teams=["team-1"])]
)
mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[key1])
async def cascading_delete_many(where):
jwt_table.cascade(("hashed-token-1",))
mock_prisma_client.db.litellm_verificationtoken.delete_many = AsyncMock(side_effect=cascading_delete_many)
mock_prisma_client.db.litellm_jwtkeymapping = jwt_table
_wire_member_delete_tx(mock_prisma_client)
cache: Final = UserApiKeyCache()
doomed_cache_keys: Final = _seed_jwt_mapping_cache(cache, doomed_rows)
(kept_cache_key,) = _seed_jwt_mapping_cache(cache, (kept_row,))
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache)
monkeypatch.setattr("litellm.proxy.management_endpoints.team_endpoints._is_user_team_admin", lambda **kwargs: True)
await team_member_delete(
data=TeamMemberDeleteRequest(team_id="team-1", user_id="user-123"),
user_api_key_dict=UserAPIKeyAuth(
user_id="admin-user", api_key="sk-admin", user_role=LitellmUserRoles.PROXY_ADMIN.value
),
)
assert all(cache.get_cache(key=cache_key) is None for cache_key in doomed_cache_keys)
assert cache.get_cache(key=kept_cache_key) == "hashed-other-key"
assert jwt_table.rows == (kept_row,)
@pytest.mark.asyncio
async def test_delete_team_evicts_jwt_key_mapping_cache_of_the_keys_it_deletes(
monkeypatch,
disable_audit_logging_for_mocked_team,
):
"""Same contract as /team/member_delete for the bulk key delete in /team/delete: the
jwt_key_mapping cache entries of the team's keys, issuer-scoped ones included, are gone
after the delete while entries pointing at other keys survive (LIT-5387)."""
from litellm.proxy._types import DeleteTeamRequest, LiteLLM_VerificationToken
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
doomed_rows: Final = (
JWTMappingRow("hashed-doomed-key", "sub", "svc-account"),
JWTMappingRow("hashed-doomed-key", "sub", "svc-account", "https://issuer.example"),
)
kept_row: Final = JWTMappingRow("hashed-unrelated-key", "sub", "svc-account", "https://other-issuer.example")
jwt_table: Final = CascadingJWTMappingTable([*doomed_rows, kept_row])
team = LiteLLM_TeamTable(
team_id="team-doomed",
team_alias="doomed-team",
members_with_roles=[],
metadata={},
model_max_budget={},
model_spend={},
)
mock_prisma_client = AsyncMock()
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team)
async def cascading_delete_data(team_id_list, table_name):
jwt_table.cascade(("hashed-doomed-key",))
return {"deleted_keys": 1}
mock_prisma_client.delete_data = AsyncMock(side_effect=cascading_delete_data)
mock_prisma_client.db.litellm_deletedteamtable.create_many = AsyncMock()
mock_prisma_client.db.litellm_deletedverificationtoken.create_many = AsyncMock()
mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
return_value=[LiteLLM_VerificationToken(token="hashed-doomed-key", team_id="team-doomed")]
)
mock_prisma_client.db.litellm_jwtkeymapping = jwt_table
mock_prisma_client.db.execute_raw = AsyncMock()
mock_prisma_client.db.litellm_teammembership.delete_many = AsyncMock()
mock_tx = AsyncMock()
mock_tx.litellm_proxymodeltable.find_many = AsyncMock(return_value=[])
mock_tx_cm = MagicMock()
mock_tx_cm.__aenter__ = AsyncMock(return_value=mock_tx)
mock_tx_cm.__aexit__ = AsyncMock(return_value=False)
mock_prisma_client.db.tx = MagicMock(return_value=mock_tx_cm)
_wire_team_delete_tx(mock_prisma_client)
cache: Final = UserApiKeyCache()
doomed_cache_keys: Final = _seed_jwt_mapping_cache(cache, doomed_rows)
(kept_cache_key,) = _seed_jwt_mapping_cache(cache, (kept_row,))
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache)
monkeypatch.setattr("litellm.proxy.proxy_server.create_audit_log_for_update", AsyncMock())
monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin")
await delete_team(
data=DeleteTeamRequest(team_ids=["team-doomed"]),
http_request=MagicMock(),
user_api_key_dict=UserAPIKeyAuth(
user_id="admin-user", api_key="sk-admin", user_role=LitellmUserRoles.PROXY_ADMIN.value
),
litellm_changed_by="admin-user",
)
assert all(cache.get_cache(key=cache_key) is None for cache_key in doomed_cache_keys)
assert cache.get_cache(key=kept_cache_key) == "hashed-unrelated-key"
assert jwt_table.rows == (kept_row,)
@pytest.mark.asyncio
async def test_new_team_negative_max_budget():
"""
@ -10401,7 +10550,7 @@ def test_new_team_request_accepts_team_member_budget_duration():
async def test_create_team_member_budget_table_with_duration():
"""Verify that create_team_member_budget_table passes budget_duration
through to the new_budget call when team_member_budget_duration is provided."""
from litellm.proxy._types import NewTeamRequest, UserAPIKeyAuth, LitellmUserRoles
from litellm.proxy._types import LitellmUserRoles, NewTeamRequest, UserAPIKeyAuth
from litellm.proxy.management_endpoints.team_endpoints import (
TeamMemberBudgetHandler,
)
@ -10891,7 +11040,7 @@ async def test_team_member_me_matches_email_only_member(mock_db_client):
@pytest.mark.asyncio
async def test_team_member_me_returns_404_for_non_member(mock_db_client):
"""A user who is not a member of the team gets 404, regardless of role."""
from fastapi import Request, HTTPException
from fastapi import HTTPException, Request
from litellm.proxy.management_endpoints.team_endpoints import team_member_me
@ -10925,7 +11074,7 @@ async def test_team_member_me_returns_404_for_proxy_admin_not_in_team(
Proxy admins get 404 if they are not actually a member of the team.
`me` only resolves for actual team members; admins use /team/info instead.
"""
from fastapi import Request, HTTPException
from fastapi import HTTPException, Request
from litellm.proxy.management_endpoints.team_endpoints import team_member_me
@ -10986,7 +11135,7 @@ async def test_team_member_me_returns_defaults_when_no_membership_row(mock_db_cl
@pytest.mark.asyncio
async def test_team_member_me_rejects_team_key_without_user_id(mock_db_client):
"""A team key with no user_id can't resolve 'me' — must return 400."""
from fastapi import Request, HTTPException
from fastapi import HTTPException, Request
from litellm.proxy.management_endpoints.team_endpoints import team_member_me
@ -11004,7 +11153,7 @@ async def test_team_member_me_rejects_team_key_without_user_id(mock_db_client):
@pytest.mark.asyncio
async def test_team_member_me_returns_404_for_unknown_team(mock_db_client):
"""Unknown team_id returns 404 — propagated from get_team_object."""
from fastapi import Request, HTTPException
from fastapi import HTTPException, Request
from litellm.proxy.management_endpoints.team_endpoints import team_member_me

View file

@ -8,6 +8,7 @@ import pytest
from pydantic import BaseModel, ConfigDict, ValidationError
from litellm.proxy._types import LiteLLM_TeamTable, LitellmUserRoles, Member, UserAPIKeyAuth
from litellm.proxy.auth.auth_checks import jwt_key_mapping_cache_key
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.list_api.common import ManagementProblem
from litellm.proxy.management_helpers.bulk_user_deletion import bulk_delete_users, bulk_remove_team_members
@ -114,6 +115,7 @@ class _Db:
tokens: Sequence[Mapping[str, object]] = (),
invitations: Sequence[Mapping[str, object]] = (),
org_memberships: Sequence[Mapping[str, object]] = (),
jwt_mappings: Sequence[Mapping[str, object]] = (),
) -> None:
self.litellm_usertable = _UserTable(users)
self.litellm_teamtable = _TeamTable(teams)
@ -122,6 +124,7 @@ class _Db:
self.litellm_deletedverificationtoken = _Rows()
self.litellm_invitationlink = _Rows(invitations)
self.litellm_organizationmembership = _Rows(org_memberships)
self.litellm_jwtkeymapping = _Rows(jwt_mappings)
class _Tx:
@ -163,11 +166,12 @@ class _FakePrisma:
tokens: Sequence[Mapping[str, object]] = (),
invitations: Sequence[Mapping[str, object]] = (),
org_memberships: Sequence[Mapping[str, object]] = (),
jwt_mappings: Sequence[Mapping[str, object]] = (),
on_lock: Callable[[str], None] = lambda _: None,
fail_locks: frozenset[str] = frozenset(),
fail_commit: bool = False,
) -> None:
self.db = _Db(users, teams, memberships, tokens, invitations, org_memberships)
self.db = _Db(users, teams, memberships, tokens, invitations, org_memberships, jwt_mappings)
self._on_lock = on_lock
self._fail_locks = fail_locks
self._fail_commit = fail_commit
@ -212,6 +216,17 @@ def _cache_with(*hashed_tokens: str) -> UserApiKeyCache:
return cache
def _jwt_mapping(token: str, claim_value: str, issuer: str | None = None) -> Mapping[str, object]:
return {"token": token, "jwt_claim_name": "sub", "jwt_claim_value": claim_value, "jwt_issuer": issuer}
def _cache_with_jwt_mapping_keys(*cache_keys: str) -> UserApiKeyCache:
cache = UserApiKeyCache()
for key in cache_keys:
cache.set_cache(key=key, value={"cache_key": key})
return cache
async def _delete(
prisma: _FakePrisma,
user_ids: Sequence[str],
@ -449,6 +464,34 @@ async def test_bulk_delete_evicts_deleted_keys_and_users_from_the_auth_cache():
assert cache.get_cache(key="keep-key") is not None
@pytest.mark.asyncio
async def test_bulk_delete_evicts_jwt_key_mappings_of_the_deleted_users_keys():
issuer: Final = "https://issuer.example"
doomed_global: Final = jwt_key_mapping_cache_key("sub", "alice")
doomed_scoped: Final = jwt_key_mapping_cache_key("sub", "alice", issuer)
kept: Final = jwt_key_mapping_cache_key("sub", "bob")
prisma = _FakePrisma(
users=[_user("u1", "t1"), _user("keep", "t1")],
teams=[_team("t1", "u1", "keep")],
tokens=[
{"token": "team-key", "user_id": "u1", "team_id": "t1"},
{"token": "personal-key", "user_id": "u1"},
{"token": "keep-key", "user_id": "keep", "team_id": "t1"},
],
jwt_mappings=[
_jwt_mapping("personal-key", "alice"),
_jwt_mapping("team-key", "alice", issuer=issuer),
_jwt_mapping("keep-key", "bob"),
],
)
cache = _cache_with_jwt_mapping_keys(doomed_global, doomed_scoped, kept)
await _delete(prisma, ["u1"], cache=cache)
assert cache.get_cache(key=doomed_global) is None and cache.get_cache(key=doomed_scoped) is None
assert cache.get_cache(key=kept) is not None
@pytest.mark.asyncio
async def test_bulk_delete_rejects_non_admin_callers_before_touching_the_db():
prisma = _FakePrisma(users=[_user("u1")])
@ -577,6 +620,28 @@ async def test_bulk_member_delete_evicts_the_removed_team_keys_from_the_auth_cac
assert cache.get_cache(key="keep-key") is not None
@pytest.mark.asyncio
async def test_bulk_member_delete_evicts_jwt_key_mappings_of_the_removed_team_keys():
issuer: Final = "https://issuer.example"
doomed: Final = jwt_key_mapping_cache_key("sub", "alice", issuer)
kept: Final = jwt_key_mapping_cache_key("sub", "bob")
prisma = _FakePrisma(
users=[_user("u1", "t1"), _user("keep", "t1")],
teams=[_team("t1", "u1", "keep")],
tokens=[
{"token": "team-key", "user_id": "u1", "team_id": "t1"},
{"token": "keep-key", "user_id": "keep", "team_id": "t1"},
],
jwt_mappings=[_jwt_mapping("team-key", "alice", issuer=issuer), _jwt_mapping("keep-key", "bob")],
)
cache = _cache_with_jwt_mapping_keys(doomed, kept)
await _remove(prisma, "t1", [{"user_id": "u1"}], cache=cache)
assert cache.get_cache(key=doomed) is None
assert cache.get_cache(key=kept) is not None
@pytest.mark.asyncio
async def test_bulk_member_delete_cleans_a_user_whose_teams_array_still_names_the_team():
prisma = _FakePrisma(users=[_user("stale", "t1")], teams=[_team("t1", "other")], memberships=[("t1", "stale")])