scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-13 08:52:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 174

fix(auth): fence native context fills during invalidation

Browse source
This commit is contained in:
Devin AI 2026-11-09 02:59:13 +00:00
parent 28eb490f93
commit 5c610bdce0
11 changed files with 314 additions and 14 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

95
litellm/proxy/auth/auth_checks.py
View file

@ -1487,6 +1487,9 @@ async def get_team_member_default_budget(
Returns:
LiteLLM_BudgetTable if found, None otherwise
"""
fill_generation: Final = (
user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None
)
if prisma_client is None:
return None
@ -1524,6 +1527,7 @@ async def get_team_member_default_budget(
budget_record.model_dump() if isinstance(budget_record, BaseModel) else budget_record.dict()
)
await user_api_key_cache.async_set_cache(
fill_generation=fill_generation,
key=cache_key,
value=budget,
model_type=LiteLLM_BudgetTable,
@ -1637,10 +1641,11 @@ async def _cache_registry_answer(
value: tuple[str, ...] | str,
ttl: float,
user_api_key_cache: UserApiKeyCache,
fill_generation: int | None = None,
) -> None:
"""Best-effort: a cache backend failure must not turn a registry load into a failed request."""
try:
await user_api_key_cache.async_set_cache(key=cache_key, value=value, ttl=ttl)
await user_api_key_cache.async_set_cache(key=cache_key, value=value, ttl=ttl, fill_generation=fill_generation)
except NativeAuthFailure:
raise
except Exception as e: # noqa: BLE001 # best-effort cache write: auth must survive a cache backend error
@ -1653,8 +1658,10 @@ async def _fetch_and_cache_registry(
max_size: int,
fetch_ids: Callable[[], Awaitable[tuple[str, ...]]],
user_api_key_cache: UserApiKeyCache,
native: bool = False,
) -> frozenset[str] | None:
"""The registry as the database has it, cached whole, or ``None`` when it is unusable."""
fill_generation: Final = user_api_key_cache.current_generation() if native else None
try:
registry_ids: Final = await fetch_ids()
except NativeAuthFailure:
@ -1668,6 +1675,7 @@ async def _fetch_and_cache_registry(
e,
)
await _cache_registry_answer(
fill_generation=fill_generation,
cache_key=cache_key,
value=overflow_sentinel,
ttl=REGISTRY_ERROR_NEGATIVE_CACHE_TTL,
@ -1677,6 +1685,7 @@ async def _fetch_and_cache_registry(
if len(registry_ids) > max_size:
await _cache_registry_answer(
fill_generation=fill_generation,
cache_key=cache_key,
value=overflow_sentinel,
ttl=get_management_object_ttl(user_api_key_cache),
@ -1685,6 +1694,7 @@ async def _fetch_and_cache_registry(
return None
await _cache_registry_answer(
fill_generation=fill_generation,
cache_key=cache_key,
value=registry_ids,
ttl=get_management_object_ttl(user_api_key_cache),
@ -1700,6 +1710,7 @@ async def _load_bounded_registry(
load_lock: asyncio.Lock,
fetch_ids: Callable[[], Awaitable[tuple[str, ...]]],
user_api_key_cache: UserApiKeyCache,
native: bool = False,
) -> frozenset[str] | None:
"""
A bounded id set under one cache key, so an id outside it costs no DB read.
@ -1719,6 +1730,7 @@ async def _load_bounded_registry(
return cached_after_wait
return await _fetch_and_cache_registry(
native=native,
cache_key=cache_key,
overflow_sentinel=overflow_sentinel,
max_size=max_size,
@ -1748,6 +1760,7 @@ async def load_end_user_restricted_registry(
return tuple(row.user_id for row in restricted_rows)
return await _load_bounded_registry(
native=native_auth_repository(prisma_client) is not None,
cache_key=end_user_restricted_registry_cache_key(),
overflow_sentinel=END_USER_RESTRICTED_REGISTRY_OVERFLOW_SENTINEL,
max_size=END_USER_RESTRICTED_REGISTRY_MAX_SIZE,
@ -1807,6 +1820,9 @@ async def get_end_user_object(
Returns:
LiteLLM_EndUserTable if found, None otherwise
"""
fill_generation: Final = (
user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None
)
if prisma_client is None:
raise Exception("No db connected")
@ -1863,6 +1879,7 @@ async def get_end_user_object(
)
await user_api_key_cache.async_set_cache(
fill_generation=fill_generation,
key=_key,
value=end_user_row,
model_type=LiteLLM_EndUserTable,
@ -1949,6 +1966,7 @@ async def _load_tag_registry(
return tuple(row.tag_name for row in registry_rows)
return await _load_bounded_registry(
native=native_auth_repository(prisma_client) is not None,
cache_key=tag_registry_cache_key(),
overflow_sentinel=TAG_REGISTRY_OVERFLOW_SENTINEL,
max_size=TAG_REGISTRY_MAX_SIZE,
@ -1978,6 +1996,7 @@ async def load_model_access_group_registry(
return tuple(row.access_group_name for row in registry_rows)
return await _load_bounded_registry(
native=native_auth_repository(prisma_client) is not None,
cache_key=model_access_group_registry_cache_key(),
overflow_sentinel=MODEL_ACCESS_GROUP_REGISTRY_OVERFLOW_SENTINEL,
max_size=MODEL_ACCESS_GROUP_REGISTRY_MAX_SIZE,
@ -2002,6 +2021,9 @@ async def _fetch_uncached_model_access_group_budgets(
``matched_model_access_groups``, which :func:`collect_matched_model_access_groups` already
intersected with the registry, so a name that has no row cannot reach this.
"""
fill_generation: Final = (
user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None
)
if not uncached_groups:
return ()
@ -2018,6 +2040,7 @@ async def _fetch_uncached_model_access_group_budgets(
fetched: Final = tuple((row.access_group_name, _model_access_group_budget(row)) for row in db_rows)
for fetched_name, fetched_obj in fetched:
await user_api_key_cache.async_set_cache(
fill_generation=fill_generation,
key=model_access_group_cache_key(fetched_name),
value=fetched_obj,
model_type=ModelAccessGroupBudget,
@ -2081,6 +2104,9 @@ async def _fetch_uncached_tags(
user_api_key_cache: UserApiKeyCache,
) -> tuple[tuple[str, LiteLLM_TagTable], ...] | Exception:
"""Rows for the tags a cache probe missed; names absent from the registry never reach the DB."""
fill_generation: Final = (
user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None
)
if not uncached_tags:
return ()
@ -2115,6 +2141,7 @@ async def _fetch_uncached_tags(
)
for fetched_name, fetched_obj in fetched:
await user_api_key_cache.async_set_cache(
fill_generation=fill_generation,
key=tag_cache_key(fetched_name),
value=fetched_obj,
model_type=LiteLLM_TagTable,
@ -2263,6 +2290,9 @@ async def _fetch_team_membership_from_db(
*,
use_writer: bool = False,
) -> LiteLLM_TeamMembership | None:
fill_generation: Final = (
user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None
)
_ = parent_otel_span, proxy_logging_obj
native: Final = native_auth_repository(prisma_client)
response: Final = (
@ -2285,12 +2315,14 @@ async def _fetch_team_membership_from_db(
_key: Final = team_membership_reservation_cache_key(user_id=user_id, team_id=team_id)
if membership is None:
await user_api_key_cache.async_set_cache(
fill_generation=fill_generation,
key=_key,
value=NO_TEAM_MEMBERSHIP_SENTINEL,
ttl=get_management_object_ttl(user_api_key_cache),
)
else:
await user_api_key_cache.async_set_cache(
fill_generation=fill_generation,
key=_key,
value=membership,
model_type=LiteLLM_TeamMembership,
@ -2534,6 +2566,9 @@ async def _backfill_null_user_email(
user_row: LiteLLM_UserTable,
user_email: str | None,
) -> LiteLLM_UserTable:
fill_generation: Final = (
user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None
)
if user_email is None or user_row.user_email is not None or prisma_client is None:
return user_row
@ -2553,6 +2588,7 @@ async def _backfill_null_user_email(
email_update: Final = {"user_email": db_row.user_email}
updated_row: Final = user_row.model_copy(update=email_update)
await user_api_key_cache.async_set_cache(
fill_generation=fill_generation,
key=user_row.user_id,
value=updated_row,
model_type=LiteLLM_UserTable,
@ -2580,6 +2616,9 @@ async def get_user_object(
- if not, then raise an error
"""
fill_generation: Final = (
user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None
)
if user_id is None:
return None
@ -2694,6 +2733,7 @@ async def get_user_object(
# save the user object to cache
await user_api_key_cache.async_set_cache(
fill_generation=fill_generation,
key=user_id,
value=_response,
model_type=LiteLLM_UserTable,
@ -2733,6 +2773,7 @@ async def _cache_management_object(
proxy_logging_obj: ProxyLogging | None,
*,
model_type: type[BaseModel],
fill_generation: int | None = None,
):
"""
Persist management objects via ``UserApiKeyCache`` (in-memory + optional Redis).
@ -2740,6 +2781,7 @@ async def _cache_management_object(
``UserApiKeyCache`` serializes with ``model_type`` so Redis and in-memory stay aligned.
"""
await user_api_key_cache.async_set_cache(
fill_generation=fill_generation,
key=key,
value=value,
model_type=model_type,
@ -2753,6 +2795,8 @@ async def cache_team_object(
team_table: LiteLLM_TeamTableCachedObj,
user_api_key_cache: UserApiKeyCache,
proxy_logging_obj: ProxyLogging | None,
*,
fill_generation: int | None = None,
) -> None:
## CACHE REFRESH TIME!
team_table.last_refreshed_at = time.time()
@ -2767,6 +2811,7 @@ async def cache_team_object(
# team_id is the table primary key — guaranteed unique, safe to write.
await _cache_management_object(
fill_generation=fill_generation,
key=key,
value=team_table,
user_api_key_cache=user_api_key_cache,
@ -3196,6 +3241,9 @@ async def _get_team_object_from_user_api_key_cache(
team_id_upsert: bool | None = None,
use_writer: bool = False,
) -> LiteLLM_TeamTableCachedObj:
fill_generation: Final = (
user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None
)
db_access_time_key: Final = key
should_check_db: Final = _should_check_db(
key=db_access_time_key,
@ -3243,6 +3291,7 @@ async def _get_team_object_from_user_api_key_cache(
# save the team object to cache
await cache_team_object(
fill_generation=fill_generation,
team_id=team_id,
team_table=_response,
user_api_key_cache=user_api_key_cache,
@ -3346,9 +3395,12 @@ async def cache_access_object(
access_group_table: LiteLLM_AccessGroupTable,
user_api_key_cache: UserApiKeyCache,
proxy_logging_obj: ProxyLogging | None = None,
*,
fill_generation: int | None = None,
) -> None:
key: Final = f"access_group_id:{access_group_id}"
await user_api_key_cache.async_set_cache(
fill_generation=fill_generation,
key=key,
value=access_group_table,
model_type=LiteLLM_AccessGroupTable,
@ -3396,6 +3448,9 @@ async def get_access_object(
Raises:
- HTTPException: If access group doesn't exist in db or cache (status_code=404)
"""
fill_generation: Final = (
user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None
)
if prisma_client is None:
raise Exception("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys")
@ -3429,6 +3484,7 @@ async def get_access_object(
# Save to cache
await cache_access_object(
fill_generation=fill_generation,
access_group_id=access_group_id,
access_group_table=_response,
user_api_key_cache=user_api_key_cache,
@ -3480,6 +3536,9 @@ async def get_team_object_by_alias(
Raises:
HTTPException: If team doesn't exist or multiple teams have the same alias
"""
fill_generation: Final = (
user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None
)
if prisma_client is None:
raise Exception("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys")
@ -3547,6 +3606,7 @@ async def get_team_object_by_alias(
# Cache the result by both alias and team_id
await user_api_key_cache.async_set_cache(
fill_generation=fill_generation,
key=cache_key,
value=team_obj,
model_type=LiteLLM_TeamTableCachedObj,
@ -3555,6 +3615,7 @@ async def get_team_object_by_alias(
# Also cache by team_id for consistency
team_id_cache_key: Final = f"team_id:{team_obj.team_id}"
await user_api_key_cache.async_set_cache(
fill_generation=fill_generation,
key=team_id_cache_key,
value=team_obj,
model_type=LiteLLM_TeamTableCachedObj,
@ -3600,6 +3661,9 @@ async def get_org_object_by_alias(
Raises:
HTTPException: If organization not found or multiple orgs have the same alias
"""
fill_generation: Final = (
user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None
)
if prisma_client is None:
raise Exception("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys")
@ -3644,6 +3708,7 @@ async def get_org_object_by_alias(
# Cache the result
await user_api_key_cache.async_set_cache(
fill_generation=fill_generation,
key=cache_key,
value=org_obj,
model_type=LiteLLM_OrganizationTable,
@ -3651,6 +3716,7 @@ async def get_org_object_by_alias(
)
# Also cache by org_id for consistency
await user_api_key_cache.async_set_cache(
fill_generation=fill_generation,
key=f"org_id:{org_obj.organization_id}",
value=org_obj,
model_type=LiteLLM_OrganizationTable,
@ -4070,6 +4136,9 @@ async def get_object_permission(
- if valid, return LiteLLM_ObjectPermissionTable object
- if not, then raise an error
"""
fill_generation: Final = (
user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None
)
if prisma_client is None:
raise Exception("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys")
@ -4108,6 +4177,7 @@ async def get_object_permission(
response.model_dump() if isinstance(response, BaseModel) else response.dict()
)
await user_api_key_cache.async_set_cache(
fill_generation=fill_generation,
key=key,
value=_perm_obj,
model_type=LiteLLM_ObjectPermissionTable,
@ -4140,6 +4210,9 @@ async def get_managed_vector_store_rows_by_uuids(
Critical-path DB access must go through this helper to avoid raw Prisma
calls on the hot request path.
"""
fill_generation: Final = (
user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None
)
if not uuids or prisma_client is None:
return []
@ -4183,6 +4256,7 @@ async def get_managed_vector_store_rows_by_uuids(
cached_obj = LiteLLM_ManagedVectorStoresTable.model_validate(row_dict)
key = f"managed_vector_store_id:{cached_obj.vector_store_id}"
await user_api_key_cache.async_set_cache(
fill_generation=fill_generation,
key=key,
value=cached_obj,
model_type=LiteLLM_ManagedVectorStoresTable,
@ -4225,6 +4299,9 @@ async def get_org_object(
proxy_logging_obj: Optional proxy logging object
include_budget_table: If True, includes litellm_budget_table in the query
"""
fill_generation: Final = (
user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None
)
if prisma_client is None:
raise Exception("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys")
if not isinstance(org_id, str):
@ -4262,6 +4339,7 @@ async def get_org_object(
_org_obj: Final = LiteLLM_OrganizationTable.model_validate(response.model_dump())
# Cache the result
await user_api_key_cache.async_set_cache(
fill_generation=fill_generation,
key=cache_key,
value=_org_obj,
model_type=LiteLLM_OrganizationTable,
@ -4277,7 +4355,10 @@ def _last_known_org_cache_key(org_id: str) -> str:
@with_service_target(AUTH_OBJECTS_TARGET)
async def _keep_last_known_org(
org: LiteLLM_OrganizationTable, org_id: str, user_api_key_cache: UserApiKeyCache
org: LiteLLM_OrganizationTable,
org_id: str,
user_api_key_cache: UserApiKeyCache,
fill_generation: int | None = None,
) -> None:
cache_key: Final = _last_known_org_cache_key(org_id)
held_locally: Final = await user_api_key_cache.async_get_cache(
@ -4286,6 +4367,7 @@ async def _keep_last_known_org(
if held_locally is not None:
return
await user_api_key_cache.async_set_cache(
fill_generation=fill_generation,
key=cache_key,
value=org,
model_type=LiteLLM_OrganizationTable,
@ -4301,6 +4383,9 @@ async def get_org_object_for_request(
parent_otel_span: Span | None,
proxy_logging_obj: ProxyLogging | None,
) -> LiteLLM_OrganizationTable | None:
fill_generation: Final = (
user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None
)
try:
org: Final = await get_org_object(
org_id=org_id,
@ -4329,7 +4414,7 @@ async def get_org_object_for_request(
raise
if org is None:
return None
await _keep_last_known_org(org, org_id, user_api_key_cache)
await _keep_last_known_org(org, org_id, user_api_key_cache, fill_generation)
return org
@ -5574,6 +5659,9 @@ async def get_project_object(
Returns LiteLLM_ProjectTableCachedObj or None if not found.
"""
fill_generation: Final = (
user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None
)
if prisma_client is None:
return None
@ -5604,6 +5692,7 @@ async def get_project_object(
# Cache with TTL following _cache_management_object pattern
project_obj.last_refreshed_at = time.time()
await _cache_management_object(
fill_generation=fill_generation,
key=cache_key,
value=project_obj,
user_api_key_cache=user_api_key_cache,

24
litellm/proxy/auth/auth_object_prefetch.py
View file

@ -223,6 +223,7 @@ async def _fill_from_redis(
"""On the request's native pipeline when one is open; a failed read is a miss, like ``async_batch_get_cache``."""
if not entries:
return
generation: Final = shared.current_generation()
before_read()
try:
with service_target(AUTH_OBJECTS_TARGET):
@ -232,8 +233,11 @@ async def _fill_from_redis(
return
for entry, value in ((entry, found.get(entry.cache_key)) for entry in entries):
if value is not None:
shared.set_local(
entry.cache_key, value, cache.in_memory_cache.default_ttl if entry.ttl is None else entry.ttl
shared.set_if_generation(
entry.cache_key,
value,
generation,
cache.in_memory_cache.default_ttl if entry.ttl is None else entry.ttl,
)
@ -284,12 +288,18 @@ async def _fetch_rows(
return _RowValues.validate_python(row) if row is not None else _NO_ROWS
async def _write_back(entries: Sequence[tuple[_CacheEntry, BaseModel]], cache: UserApiKeyCache) -> None:
async def _write_back(
entries: Sequence[tuple[_CacheEntry, BaseModel]], cache: UserApiKeyCache, generation: int | None = None
) -> None:
payloads: Final = tuple(
(entry.cache_key, CacheCodec.serialize(value, model_type=entry.model_type), entry.ttl)
for entry, value in entries
)
shared: Final = cache.shared
if generation is not None:
for cache_key, payload, ttl in payloads:
cache.set_if_generation(cache_key, payload, generation, ttl)
return
if shared is not None:
with service_target(AUTH_OBJECTS_TARGET):
await shared.write_back(payloads, batch=Batch.PRE_CALL)
@ -303,12 +313,13 @@ async def _fill_from_db(
) -> None:
if not entries:
return
native: Final = native_auth_repository(prisma_client) is not None
generation: Final = cache.current_generation() if native else None
model_for: Final[Mapping[_RowKind, type[BaseModel]]] = MappingProxyType(
{entry.row: entry.model_type for entry in entries}
)
rows: Final = await _fetch_rows(refs, frozenset(model_for), prisma_client)
refreshed_at: Final = time.time()
native: Final = native_auth_repository(prisma_client) is not None
objects: Final[Mapping[_RowKind, BaseModel | None]] = MappingProxyType(
{
row: _validate_row(rows.get(row), model_type, row, refreshed_at, native=native)
@ -317,7 +328,7 @@ async def _fill_from_db(
)
writes: Final = tuple((entry, value) for entry in entries if (value := objects[entry.row]) is not None)
if writes:
await _write_back(writes, cache)
await _write_back(writes, cache, generation)
async def prefetch_auth_objects(
@ -363,6 +374,7 @@ async def prefetch_identity_keys(
missing: Final = tuple(key for key in dict.fromkeys(cache_keys) if shared.get_local(key) is None)
if not missing:
return
generation: Final = shared.current_generation()
before_read()
with service_target(AUTH_OBJECTS_TARGET):
found: Final = await shared.read_missing(sorted(missing), batch=Batch.PRE_CALL)
@ -372,4 +384,4 @@ async def prefetch_identity_keys(
return
for key, value in ((key, found.get(key)) for key in missing):
if value is not None:
shared.set_local(key, value, _identity_memory_ttl(value, management_ttl))
shared.set_if_generation(key, value, generation, _identity_memory_ttl(value, management_ttl))

7
litellm/proxy/auth/end_user_context_io.py
View file

@ -31,6 +31,7 @@ async def read_default_budget(budget_id: str, io: EndUserIO) -> LiteLLM_BudgetTa
return cached
try:
native: Final = native_auth_repository(io.prisma)
fill_generation: Final = io.cache.current_generation() if native is not None else None
record: Final = (
await native.one("budget", [budget_id], LiteLLM_BudgetTable)
if native is not None
@ -47,7 +48,11 @@ async def read_default_budget(budget_id: str, io: EndUserIO) -> LiteLLM_BudgetTa
record.model_dump() if isinstance(record, LiteLLM_BudgetTable) else record.dict()
)
await io.cache.async_set_cache(
key=key, value=row, model_type=LiteLLM_BudgetTable, ttl=get_management_object_ttl(io.cache)
fill_generation=fill_generation,
key=key,
value=row,
model_type=LiteLLM_BudgetTable,
ttl=get_management_object_ttl(io.cache),
)
except NativeAuthFailure:
raise

2
litellm/proxy/auth/user_api_key_auth.py
View file

@ -2360,6 +2360,7 @@ class _BuilderEffects:
if self._prisma_client is None:
return None
native: Final = native_auth_repository(self._prisma_client)
fill_generation: Final = self._user_api_key_cache.current_generation() if native is not None else None
row: Final = (
await native.one("member", [user_id, team_id], LiteLLM_TeamMembership)
if native is not None
@ -2372,6 +2373,7 @@ class _BuilderEffects:
return None
membership: Final = LiteLLM_TeamMembership.model_validate(row.model_dump())
await self._user_api_key_cache.async_set_cache(
fill_generation=fill_generation,
key=team_membership_auth_cache_key(team_id=team_id, user_id=user_id),
value=membership,
model_type=LiteLLM_TeamMembership,

15
litellm/proxy/common_utils/user_api_key_cache.py
View file

@ -325,11 +325,24 @@ class UserApiKeyCache(DualCache):
return self.key_object_cache.set_cache(key=key, value=payload, local_only=local_only, **kwargs)
return super().set_cache(key=key, value=payload, local_only=local_only, **kwargs)
async def async_set_cache(self, key: str | None, value: CacheValue, local_only: bool = False, **kwargs: object):
async def async_set_cache(
self,
key: str | None,
value: CacheValue,
local_only: bool = False,
*,
fill_generation: int | None = None,
**kwargs: object,
):
"""Inside a request the Redis SET rides the request's pipeline (memory is written at once); anywhere
else, or with options the pipeline does not carry, it goes to Redis directly as before."""
model_type: Final = cast(type[BaseModel] | None, kwargs.pop("model_type", None))
payload: Final[object] = CacheCodec.serialize(value, model_type=model_type)
if fill_generation is not None:
if key is not None:
assert self._shared is not None
self._shared.set_if_generation(key, payload, fill_generation, _ttl(kwargs))
return None
ttl: Final = kwargs.get("ttl")
pipelined: Final = (
key is not None

6
tests/integration/_support/database_relay.py
View file

@ -85,12 +85,14 @@ class HeldStatementRelay:
*,
after: bytes | None = None,
armed: bool = True,
response: bool = False,
) -> None:
self.port: Final = allocate_port()
self._upstream_host: Final = upstream_host
self._upstream_port: Final = upstream_port
self._trigger: Final = trigger
self._after: Final = after
self._response: Final = response
self._marker_length: Final = max(len(trigger), len(after or b""))
self._server: Final = LoopbackServer(self.port, self._serve)
self._loop: Final = self._server.loop
@ -144,8 +146,8 @@ class HeldStatementRelay:
await close_writer(writer)
await asyncio.gather(
forward(client_reader, server_writer, True),
forward(server_reader, client_writer, False),
forward(client_reader, server_writer, not self._response),
forward(server_reader, client_writer, self._response),
)

112
tests/integration/caching/test_auth_cache_native.py
View file

@ -3,15 +3,22 @@ from __future__ import annotations
import asyncio
import hashlib
import json
import os
import uuid
from collections.abc import Callable
from dataclasses import replace
from pathlib import Path
from types import SimpleNamespace
from typing import Final
import psycopg
import pytest
from redis.asyncio import Redis
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.caching.redis_cache import RedisCache
from litellm.proxy.auth.auth_checks import get_team_object
from litellm.proxy.auth.auth_object_prefetch import AuthObjectRefs, prefetch_auth_objects, prefetch_identity_keys
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
AuthCacheInvalidationSubscriber,
auth_cache_invalidation_channel,
@ -21,6 +28,9 @@ from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.redis.stores.connection import connect
from litellm.redis.stores.spend_counters import SpendCounterStore
from litellm.repositories.database_settings import NativeDatabaseSettings
from litellm.rust_bridge._native import NativeDatabase
from tests.integration._support.database_relay import HeldStatementRelay
from tests.integration._support.redis_process import owned_redis
KEY: Final = hashlib.sha256(b"native-cache-test-key").hexdigest()
@ -123,3 +133,105 @@ async def test_management_eviction_reaches_native_memory_and_shared_redis(tmp_pa
await evict_shared((KEY,), first)
second.delete_local(KEY)
assert await second.async_get_cache(KEY) is None
@pytest.mark.asyncio
@pytest.mark.parametrize("identity", (True, False))
async def test_prefetch_does_not_restore_an_invalidated_redis_reply(tmp_path: Path, identity: bool) -> None:
key: Final = KEY if identity else "team_id:prefetch-race"
value: Final = {"token": KEY, "blocked": False, "marker": "held-redis-prefetch"}
with owned_redis(tmp_path) as server:
relay: Final = HeldStatementRelay(server.host, server.port, b"held-redis-prefetch", response=True)
relay.start()
coordination: Final = RedisCache(host=server.host, port=server.port)
cache: Final = UserApiKeyCache(redis_cache=RedisCache(host=server.host, port=relay.port, socket_timeout=30))
writer: Final = UserApiKeyCache(redis_cache=coordination)
subscriber: Final = AuthCacheInvalidationSubscriber(coordination, cache)
redis: Final = Redis(host=server.host, port=server.port)
subscriber.start()
try:
await until(lambda: subscriber.is_ready)
await writer.async_set_cache(key, value)
fill: Final = asyncio.create_task(
prefetch_identity_keys((key,), cache)
if identity
else prefetch_auth_objects(AuthObjectRefs(team_id="prefetch-race"), cache, None)
)
await until(relay.held.is_set)
generation: Final = cache.current_generation()
await evict_shared((key,), writer)
assert (
await redis.publish(auth_cache_invalidation_channel(coordination), json.dumps({"cache_key": key})) == 1
)
await until(lambda: cache.current_generation() > generation)
assert await redis.get(key) is None
assert cache.get_local(key) is None
relay.release()
await fill
assert cache.get_local(key) is None
assert await cache.async_get_cache(key) is None
finally:
relay.release()
await subscriber.stop()
await redis.aclose()
relay.stop()
@pytest.mark.asyncio
@pytest.mark.parametrize("prefetch", (True, False))
async def test_sqlx_prefetch_does_not_restore_local_or_shared_context_after_invalidation(
tmp_path: Path, prefetch: bool
) -> None:
team_id: Final = f"prefetch-{uuid.uuid4().hex}"
key: Final = f"team_id:{team_id}"
settings: Final = NativeDatabaseSettings.from_url(os.environ["DATABASE_URL"])
relay: Final = HeldStatementRelay(settings.host, settings.port, team_id.encode(), response=True)
relay.start()
database: Final = await NativeDatabase.connect(
replace(settings, host="127.0.0.1", port=relay.port, ssl_mode="disable")
)
with psycopg.connect(os.environ["DATABASE_URL"]) as connection:
connection.execute('INSERT INTO "LiteLLM_TeamTable" (team_id, blocked) VALUES (%s, false)', (team_id,))
with owned_redis(tmp_path) as server:
coordination: Final = RedisCache(host=server.host, port=server.port)
cache: Final = UserApiKeyCache(redis_cache=coordination)
writer: Final = UserApiKeyCache(redis_cache=coordination)
subscriber: Final = AuthCacheInvalidationSubscriber(coordination, cache)
redis: Final = Redis(host=server.host, port=server.port)
subscriber.start()
try:
await until(lambda: subscriber.is_ready)
fill: Final = asyncio.create_task(
prefetch_auth_objects(
AuthObjectRefs(team_id=team_id), cache, SimpleNamespace(native_auth_database=database)
)
if prefetch
else get_team_object(team_id, SimpleNamespace(native_auth_database=database), cache)
)
await until(relay.held.is_set)
with psycopg.connect(os.environ["DATABASE_URL"]) as connection:
connection.execute('UPDATE "LiteLLM_TeamTable" SET blocked=true WHERE team_id=%s', (team_id,))
generation: Final = cache.current_generation()
await evict_shared((key,), writer)
assert (
await redis.publish(auth_cache_invalidation_channel(coordination), json.dumps({"cache_key": key})) == 1
)
await until(lambda: cache.current_generation() > generation)
assert await redis.get(key) is None
assert cache.get_local(key) is None
relay.release()
await fill
assert await redis.get(key) is None
assert cache.get_local(key) is None
await prefetch_auth_objects(
AuthObjectRefs(team_id=team_id), cache, SimpleNamespace(native_auth_database=database)
)
assert cache.get_local(key)["blocked"] is True
finally:
relay.release()
await subscriber.stop()
await redis.aclose()
await database.close()
relay.stop()
with psycopg.connect(os.environ["DATABASE_URL"]) as connection:
connection.execute('DELETE FROM "LiteLLM_TeamTable" WHERE team_id=%s', (team_id,))

5
tests/unit/proxy/auth/test_auth_checks.py
View file

@ -2338,7 +2338,10 @@ async def test_bulk_eviction_clears_this_worker_before_waiting_on_a_wedged_redis
monkeypatch.setattr(pubsub_module, "_pending_publishes", set())
monkeypatch.setattr(pubsub_module, "_BULK_EVICTION_WAIT_SECONDS", 0.5)
client = _WedgedPublishRedisClient()
cache: Final = UserApiKeyCache(key_object_in_memory_cache=InMemoryCache(max_size_in_memory=2_000))
cache: Final = UserApiKeyCache(
in_memory_cache=InMemoryCache(max_size_in_memory=2_000),
key_object_in_memory_cache=InMemoryCache(max_size_in_memory=2_000),
)
hashed_tokens: Final = tuple(
hashlib.sha256(f"token-{index}".encode()).hexdigest()
for index in range(600)

31
tests/unit/proxy/auth/test_auth_object_prefetch.py
View file

@ -3,6 +3,7 @@ per-object getters still enforce on their own when the prefetch cannot help."""
import json
from types import SimpleNamespace
from typing import Final
from unittest.mock import AsyncMock, MagicMock
import pytest
@ -29,6 +30,7 @@ from litellm.proxy.common_utils.user_api_key_cache import (
UserApiKeyCache,
end_user_cache_key,
end_user_restricted_registry_cache_key,
team_membership_reservation_cache_key,
)
from tests.unit.redis.stores.auth_cache_double import FakeAuthCache, shared_cache
@ -102,6 +104,35 @@ async def test_native_prefetch_fills_all_getters_with_one_read_and_preserves_opt
assert native.auth_read.await_count == 1
@pytest.mark.asyncio
async def test_native_owner_getters_do_not_cache_rows_invalidated_during_their_read():
cache: Final = _cache(None)
async def read(kind: str, ids: list[str], writer: bool = False) -> list[object]:
cache.delete_local(
{
"user": USER_ID,
"team": f"team_id:{TEAM_ID}",
"member": team_membership_reservation_cache_key(USER_ID, TEAM_ID),
"org": f"org_id:{ORG_ID}:with_budget",
}[kind]
)
return [{"user": USER_ROW, "team": TEAM_ROW, "member": MEMBERSHIP_ROW, "org": ORG_ROW}[kind]]
native: Final = SimpleNamespace(auth_read=AsyncMock(side_effect=read))
prisma: Final = _prisma()
prisma.native_auth_database = native
user, team, member, org = await _read_all_through_getters(cache, prisma)
assert user is not None and team is not None and member is not None and org is not None
assert cache.get_local(USER_ID) is None
assert cache.get_local(f"team_id:{TEAM_ID}") is None
assert cache.get_local(team_membership_reservation_cache_key(USER_ID, TEAM_ID)) is None
assert cache.get_local(f"org_id:{ORG_ID}:with_budget") is None
assert native.auth_read.await_count == 4
assert prisma.db.mock_calls == []
@pytest.mark.asyncio
@pytest.mark.parametrize("rows", [[{"user_row": {"spend": "bad"}}], {"user_row": USER_ROW}])
async def test_native_prefetch_invalid_projection_is_terminal_without_legacy_query(rows):

19
tests/unit/proxy/common_utils/test_user_api_key_cache.py
View file

@ -41,6 +41,25 @@ def _make_key_obj(token: str = "tok") -> UserAPIKeyAuth:
class TestUserApiKeyCache:
@pytest.mark.asyncio
@pytest.mark.parametrize("invalidate", (True, False))
async def test_native_db_fill_is_local_and_generation_fenced(self, invalidate: bool) -> None:
cache, fake = shared_cache()
generation: Final = cache.current_generation()
if invalidate:
cache.delete_local(HASHED_TOKEN)
await cache.async_set_cache(
HASHED_TOKEN,
_make_key_obj(HASHED_TOKEN),
model_type=UserAPIKeyAuth,
fill_generation=generation,
)
assert fake.redis == {}
assert fake.commands == []
assert cache.get_cache(HASHED_TOKEN, local_only=True, model_type=UserAPIKeyAuth) == (
None if invalidate else _make_key_obj(HASHED_TOKEN)
)
def test_native_fill_fence_and_request_local_projection_without_redis(self):
cache: Final = UserApiKeyCache()
stale: Final = cache.begin_fill(HASHED_TOKEN)

12
tests/unit/redis/stores/auth_cache_double.py
View file

@ -44,6 +44,16 @@ class FakeAuthCache:
commands: list[tuple[object, ...]] = field(default_factory=list)
settings: list[AuthCacheSettings] = field(default_factory=list)
fail: BaseException | None = None
generation: int = 0
def current_generation(self) -> int:
return self.generation
def set_if_generation(self, key: str, value: object, generation: int, ttl_seconds: float | None = None) -> bool:
if generation != self.generation:
return False
self.set_local(key, value, ttl_seconds)
return True
def _round_trip(self, *command: object) -> None:
self.commands.append(command)
@ -147,9 +157,11 @@ class FakeAuthCache:
self.memory[key] = (value, ttl_seconds)
def delete_local(self, key: str) -> None:
self.generation += 1
self.memory.pop(key, None)
def flush_local(self) -> None:
self.generation += 1
self.memory.clear()
def len_local(self) -> LocalSizes: