mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-19 00:01:29 +00:00
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:
commit
4f70b88a1f
11 changed files with 548 additions and 39 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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"],
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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")])
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue