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

feat(auth): route effect reads through the native database handle

Browse source
This commit is contained in:
Devin AI 2026-11-09 02:13:03 +00:00
parent ebf40cc8d3
commit dd55f2c292
10 changed files with 655 additions and 65 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -140,6 +140,12 @@ from litellm.redis.stores.spend_counters import (
team_member_counter,
)
from litellm.repositories.budget_repository import BudgetRepository
from litellm.repositories.native_auth_repository import (
EndUserRegistryRow,
ModelAccessGroupRow,
TagRegistryRow,
native_auth_repository,
)
from litellm.repositories.object_permission_repository import ObjectPermissionRepository
from litellm.repositories.organization_repository import OrganizationRepository
from litellm.repositories.prisma_protocols import DatabaseClient, RowT_co
@ -1494,8 +1500,13 @@ async def get_team_member_default_budget(
return cached_budget
try:
budget_record: Final = await _dictable_table(BudgetRepository(prisma_client), "budget").find_unique(
where={"budget_id": budget_id}
native: Final = native_auth_repository(prisma_client)
budget_record: Final = (
await native.one("budget", [budget_id], LiteLLM_BudgetTable)
if native is not None
else await _dictable_table(BudgetRepository(prisma_client), "budget").find_unique(
where={"budget_id": budget_id}
)
)
except NativeAuthFailure:
raise
@ -1509,7 +1520,9 @@ async def get_team_member_default_budget(
verbose_proxy_logger.warning("Team-default member budget not found in database: %s", budget_id)
return None
budget: Final = LiteLLM_BudgetTable.model_validate(budget_record.dict())
budget: Final = LiteLLM_BudgetTable.model_validate(
budget_record.model_dump() if isinstance(budget_record, BaseModel) else budget_record.dict()
)
await user_api_key_cache.async_set_cache(
key=cache_key,
value=budget,
@ -1721,9 +1734,16 @@ async def load_end_user_restricted_registry(
"""The set of end-user ids whose ``LiteLLM_EndUserTable`` row carries a restriction."""
async def fetch_ids() -> tuple[str, ...]:
restricted_rows: Final = await _end_user_table(EndUserRepository(prisma_client)).find_many(
where=_restricted_end_user_where(),
take=END_USER_RESTRICTED_REGISTRY_MAX_SIZE + 1,
native: Final = native_auth_repository(prisma_client)
restricted_rows: Final = (
await native.many(
"end_user_restricted_registry", [str(END_USER_RESTRICTED_REGISTRY_MAX_SIZE + 1)], EndUserRegistryRow
)
if native is not None
else await _end_user_table(EndUserRepository(prisma_client)).find_many(
where=_restricted_end_user_where(),
take=END_USER_RESTRICTED_REGISTRY_MAX_SIZE + 1,
)
)
return tuple(row.user_id for row in restricted_rows)
@ -1820,16 +1840,23 @@ async def get_end_user_object(
# Fetch from database
try:
response: Final = await _dictable_table(EndUserRepository(prisma_client), "end_user").find_unique(
where={"user_id": end_user_id},
include={"litellm_budget_table": True, "object_permission": True},
native: Final = native_auth_repository(prisma_client)
response: Final = (
await native.one("end_user", [end_user_id], LiteLLM_EndUserTable)
if native is not None
else await _dictable_table(EndUserRepository(prisma_client), "end_user").find_unique(
where={"user_id": end_user_id},
include={"litellm_budget_table": True, "object_permission": True},
)
)
if response is None:
raise Exception
end_user_row: Final = await _apply_default_budget_to_end_user(
end_user_obj=LiteLLM_EndUserTable.model_validate(response.dict()),
end_user_obj=LiteLLM_EndUserTable.model_validate(
response.model_dump() if isinstance(response, BaseModel) else response.dict()
),
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
@ -1911,8 +1938,13 @@ async def _load_tag_registry(
"""The set of tag names that have a row in ``LiteLLM_TagTable``."""
async def fetch_ids() -> tuple[str, ...]:
registry_rows: Final = await _tag_table(TagRepository(prisma_client)).find_many(
take=TAG_REGISTRY_MAX_SIZE + 1,
native: Final = native_auth_repository(prisma_client)
registry_rows: Final = (
await native.many("tag_registry", [str(TAG_REGISTRY_MAX_SIZE + 1)], TagRegistryRow)
if native is not None
else await _tag_table(TagRepository(prisma_client)).find_many(
take=TAG_REGISTRY_MAX_SIZE + 1,
)
)
return tuple(row.tag_name for row in registry_rows)
@ -1933,9 +1965,16 @@ async def load_model_access_group_registry(
"""The set of model access group names that have a row in ``LiteLLM_ModelAccessGroupBudgetTable``."""
async def fetch_ids() -> tuple[str, ...]:
registry_rows: Final = await _model_access_group_budget_table(
ModelAccessGroupBudgetRepository(prisma_client)
).find_many(take=MODEL_ACCESS_GROUP_REGISTRY_MAX_SIZE + 1)
native: Final = native_auth_repository(prisma_client)
registry_rows: Final = (
await native.many(
"model_access_group_registry", [str(MODEL_ACCESS_GROUP_REGISTRY_MAX_SIZE + 1)], ModelAccessGroupRow
)
if native is not None
else await _model_access_group_budget_table(ModelAccessGroupBudgetRepository(prisma_client)).find_many(
take=MODEL_ACCESS_GROUP_REGISTRY_MAX_SIZE + 1
)
)
return tuple(row.access_group_name for row in registry_rows)
return await _load_bounded_registry(
@ -1967,11 +2006,14 @@ async def _fetch_uncached_model_access_group_budgets(
return ()
try:
db_rows: Final = await _model_access_group_budget_table(
ModelAccessGroupBudgetRepository(prisma_client)
).find_many(
where={"access_group_name": {"in": list(uncached_groups)}},
include={"litellm_budget_table": True},
native: Final = native_auth_repository(prisma_client)
db_rows: Final = (
await native.many("model_access_group", uncached_groups, ModelAccessGroupRow)
if native is not None
else await _model_access_group_budget_table(ModelAccessGroupBudgetRepository(prisma_client)).find_many(
where={"access_group_name": {"in": list(uncached_groups)}},
include={"litellm_budget_table": True},
)
)
fetched: Final = tuple((row.access_group_name, _model_access_group_budget(row)) for row in db_rows)
for fetched_name, fetched_obj in fetched:
@ -2053,11 +2095,24 @@ async def _fetch_uncached_tags(
return ()
try:
db_tags: Final = await _tag_table(TagRepository(prisma_client)).find_many(
where={"tag_name": {"in": list(tags_to_fetch)}},
include={"litellm_budget_table": True},
native: Final = native_auth_repository(prisma_client)
db_tags: Final = (
await native.many("tag", tags_to_fetch, LiteLLM_TagTable)
if native is not None
else await _tag_table(TagRepository(prisma_client)).find_many(
where={"tag_name": {"in": list(tags_to_fetch)}},
include={"litellm_budget_table": True},
)
)
fetched: Final = tuple(
(
db_tag.tag_name,
LiteLLM_TagTable.model_validate(
db_tag.model_dump() if isinstance(db_tag, BaseModel) else db_tag.dict()
),
)
for db_tag in db_tags
)
fetched: Final = tuple((db_tag.tag_name, LiteLLM_TagTable.model_validate(db_tag.dict())) for db_tag in db_tags)
for fetched_name, fetched_obj in fetched:
await user_api_key_cache.async_set_cache(
key=tag_cache_key(fetched_name),
@ -2209,13 +2264,24 @@ async def _fetch_team_membership_from_db(
use_writer: bool = False,
) -> LiteLLM_TeamMembership | None:
_ = parent_otel_span, proxy_logging_obj
response: Final = await _dictable_table(
TeamMembershipRepository(prisma_client, use_writer=use_writer), "team_membership"
).find_unique(
where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}},
include={"litellm_budget_table": True},
native: Final = native_auth_repository(prisma_client)
response: Final = (
await native.one("member", [user_id, team_id], LiteLLM_TeamMembership, writer=use_writer)
if native is not None
else await _dictable_table(
TeamMembershipRepository(prisma_client, use_writer=use_writer), "team_membership"
).find_unique(
where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}},
include={"litellm_budget_table": True},
)
)
membership: Final = (
None
if response is None
else LiteLLM_TeamMembership.model_validate(
response.model_dump() if isinstance(response, BaseModel) else response.dict()
)
)
membership: Final = None if response is None else LiteLLM_TeamMembership.model_validate(response.dict())
_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(
@ -2427,18 +2493,27 @@ async def _get_fuzzy_user_object(
"""
response = None
native: Final = native_auth_repository(prisma_client)
if sso_user_id is not None:
response = await _user_table(UserRepository(prisma_client)).find_unique(
where={"sso_user_id": sso_user_id},
include={"organization_memberships": True},
response = (
await native.one("user_sso", [sso_user_id], LiteLLM_UserTable)
if native is not None
else await _user_table(UserRepository(prisma_client)).find_unique(
where={"sso_user_id": sso_user_id},
include={"organization_memberships": True},
)
)
if response is None and user_email is not None:
# Use case-insensitive query to handle emails with different casing
# This matches the pattern used in _check_duplicate_user_email
response = await _user_table(UserRepository(prisma_client)).find_first(
where={"user_email": {"equals": user_email, "mode": "insensitive"}},
include={"organization_memberships": True},
response = (
await native.one("user_email", [user_email], LiteLLM_UserTable)
if native is not None
else await _user_table(UserRepository(prisma_client)).find_first(
where={"user_email": {"equals": user_email, "mode": "insensitive"}},
include={"organization_memberships": True},
)
)
if response is not None and sso_user_id is not None: # update sso_user_id
@ -2528,8 +2603,13 @@ async def get_user_object(
)
if should_check_db:
response = await _user_table(UserRepository(prisma_client, use_writer=bool(check_db_only))).find_unique(
where={"user_id": user_id}, include={"organization_memberships": True}
native: Final = native_auth_repository(prisma_client)
response = (
await native.one("user", [user_id], LiteLLM_UserTable, writer=bool(check_db_only))
if native is not None
else await _user_table(UserRepository(prisma_client, use_writer=bool(check_db_only))).find_unique(
where={"user_id": user_id}, include={"organization_memberships": True}
)
)
if response is None:
@ -3074,8 +3154,13 @@ class AccessGroupNotFoundError(HTTPException):
async def _get_team_db_check(
team_id: str, prisma_client: PrismaClient, team_id_upsert: bool | None = None, *, use_writer: bool = False
) -> "_PrismaTeamRow | None":
response = await _team_table(TeamRepository(prisma_client, use_writer=use_writer)).find_unique(
where={"team_id": team_id}, include=_TEAM_GRANT_RELATIONS
native: Final = native_auth_repository(prisma_client)
response = (
await native.one("team", [team_id], LiteLLM_TeamTable, writer=use_writer)
if native is not None
else await _team_table(TeamRepository(prisma_client, use_writer=use_writer)).find_unique(
where={"team_id": team_id}, include=_TEAM_GRANT_RELATIONS
)
)
if response is None and team_id_upsert:
@ -3321,14 +3406,21 @@ async def get_access_object(
# Not in cache - fetch from DB
try:
response: Final = await _dictable_table(
AccessGroupRepository(prisma_client, use_writer=check_db_only), "access_group"
).find_unique(where={"access_group_id": access_group_id})
native: Final = native_auth_repository(prisma_client)
response: Final = (
await native.one("access_group", [access_group_id], LiteLLM_AccessGroupTable, writer=check_db_only)
if native is not None
else await _dictable_table(
AccessGroupRepository(prisma_client, use_writer=check_db_only), "access_group"
).find_unique(where={"access_group_id": access_group_id})
)
if response is None:
raise AccessGroupNotFoundError(access_group_id=access_group_id)
_response: Final = LiteLLM_AccessGroupTable.model_validate(response.dict())
_response: Final = LiteLLM_AccessGroupTable.model_validate(
response.model_dump() if isinstance(response, BaseModel) else response.dict()
)
# Save to cache
await cache_access_object(
@ -3400,8 +3492,13 @@ async def get_team_object_by_alias(
# Query database by team_alias
try:
teams: Final = await _team_table(TeamRepository(prisma_client)).find_many(
where={"team_alias": team_alias}, include=_TEAM_GRANT_RELATIONS
native: Final = native_auth_repository(prisma_client)
teams: Final = (
await native.many("team_alias", [team_alias], LiteLLM_TeamTable)
if native is not None
else await _team_table(TeamRepository(prisma_client)).find_many(
where={"team_alias": team_alias}, include=_TEAM_GRANT_RELATIONS
)
)
if not teams:
@ -3512,8 +3609,13 @@ async def get_org_object_by_alias(
# Query database by organization_alias
try:
orgs = await _model_dump_table(OrganizationRepository(prisma_client), "organization").find_many(
where={"organization_alias": org_alias}
native: Final = native_auth_repository(prisma_client)
orgs = (
await native.many("org_alias", [org_alias], LiteLLM_OrganizationTable)
if native is not None
else await _model_dump_table(OrganizationRepository(prisma_client), "organization").find_many(
where={"organization_alias": org_alias}
)
)
if not orgs:
@ -3735,6 +3837,9 @@ async def _fetch_key_object_from_db_unbounded(
*,
check_db_only: bool = False,
) -> BaseModel | None:
native: Final = native_auth_repository(prisma_client)
if native is not None:
return await native.key(hashed_token, writer=check_db_only)
fetch: Final = partial(prisma_client.get_data, use_writer=True) if check_db_only else prisma_client.get_data
async with db_lookup_gate.current():
try:
@ -3978,16 +4083,25 @@ async def get_object_permission(
# else, check db
try:
response: Final = await _dictable_table(
ObjectPermissionRepository(prisma_client, use_writer=check_db_only), "object_permission"
).find_unique(where={"object_permission_id": object_permission_id})
native: Final = native_auth_repository(prisma_client)
response: Final = (
await native.one(
"object_permission", [object_permission_id], LiteLLM_ObjectPermissionTable, writer=check_db_only
)
if native is not None
else await _dictable_table(
ObjectPermissionRepository(prisma_client, use_writer=check_db_only), "object_permission"
).find_unique(where={"object_permission_id": object_permission_id})
)
if response is None:
if check_db_only:
raise HTTPException(status_code=403, detail="Referenced object permission does not exist")
return None
_perm_obj: Final = LiteLLM_ObjectPermissionTable.model_validate(response.dict())
_perm_obj: Final = LiteLLM_ObjectPermissionTable.model_validate(
response.model_dump() if isinstance(response, BaseModel) else response.dict()
)
await user_api_key_cache.async_set_cache(
key=key,
value=_perm_obj,
@ -4119,8 +4233,11 @@ async def get_org_object(
if include_budget_table:
query_kwargs["include"] = {"litellm_budget_table": True}
response: Final = await _model_dump_table(OrganizationRepository(prisma_client), "organization").find_unique(
**query_kwargs
native: Final = native_auth_repository(prisma_client)
response: Final = (
await native.one("org", [org_id], LiteLLM_OrganizationTable)
if native is not None
else await _model_dump_table(OrganizationRepository(prisma_client), "organization").find_unique(**query_kwargs)
)
if response is None:
@ -5456,9 +5573,14 @@ async def get_project_object(
return deserialized_project
# Fetch from DB
project_row: Final = await _model_dump_table(ProjectRepository(prisma_client), "project").find_unique(
where={"project_id": project_id},
include={"litellm_budget_table": True},
native: Final = native_auth_repository(prisma_client)
project_row: Final = (
await native.one("project", [project_id], LiteLLM_ProjectTableCachedObj)
if native is not None
else await _model_dump_table(ProjectRepository(prisma_client), "project").find_unique(
where={"project_id": project_id},
include={"litellm_budget_table": True},
)
)
if project_row is None:
return None

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

@ -4,6 +4,7 @@ the readers and the fallback, so enforcement never depends on this running."""
from __future__ import annotations
import json
import time
from collections.abc import Callable, Iterator, Mapping, Sequence
from dataclasses import dataclass
@ -12,6 +13,7 @@ from typing import Final, Literal, TypeAlias
from pydantic import BaseModel, TypeAdapter, ValidationError
from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError
from litellm._internal_context import service_target
from litellm._logging import verbose_proxy_logger
from litellm.caching.redis_batch import Batch
@ -32,6 +34,7 @@ from litellm.proxy.common_utils.user_api_key_cache import (
from litellm.proxy.db.db_span import db_span
from litellm.proxy.utils import PrismaClient
from litellm.redis.stores.auth_cache import AuthCacheStore
from litellm.repositories.native_auth_repository import native_auth_repository
_RowKind: TypeAlias = Literal["user_row", "team_row", "membership_row", "organization_row", "project_row"]
@ -235,7 +238,7 @@ async def _fill_from_redis(
def _validate_row(
row_value: object, model_type: type[BaseModel], row: _RowKind, refreshed_at: float
row_value: object, model_type: type[BaseModel], row: _RowKind, refreshed_at: float, *, native: bool = False
) -> BaseModel | None:
if row_value is None:
return None
@ -246,6 +249,8 @@ def _validate_row(
return model_type.model_validate(stamped)
return model_type.model_validate(columns)
except ValidationError as e:
if native:
raise NativeAuthRuntimeError(f"Invalid native auth {row} projection") from e
verbose_proxy_logger.warning("auth prefetch: %s did not validate as %s: %s", row, model_type.__name__, e)
return None
@ -253,6 +258,20 @@ def _validate_row(
async def _fetch_rows(
refs: AuthObjectRefs, kinds: frozenset[_RowKind], prisma_client: PrismaClient
) -> Mapping[str, object]:
native: Final = native_auth_repository(prisma_client)
if native is not None:
refs_json: Final = json.dumps(
[
refs.user_id if "user_row" in kinds else None,
refs.team_id if kinds & _TEAM_BOUND_ROWS else None,
refs.membership_user_id if "membership_row" in kinds else None,
refs.organization_id if "organization_row" in kinds else None,
refs.project_id if "project_row" in kinds else None,
]
)
async with db_span("prefetch_auth_objects", AUTH_OBJECTS_TARGET):
rows: Final = await native.read("prefetch", [refs_json])
return rows[0] if rows else _NO_ROWS
async with db_span("prefetch_auth_objects", AUTH_OBJECTS_TARGET):
row: Final[object] = await prisma_client.db.query_first( # pyright: ignore[reportAny] # prisma types query_first as Any
_SQL,
@ -289,8 +308,12 @@ async def _fill_from_db(
)
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) for row, model_type in model_for.items()}
{
row: _validate_row(rows.get(row), model_type, row, refreshed_at, native=native)
for row, model_type in model_for.items()
}
)
writes: Final = tuple((entry, value) for entry in entries if (value := objects[entry.row]) is not None)
if writes:
@ -313,6 +336,8 @@ async def prefetch_auth_objects(
if prisma_client is None:
return
await _fill_from_db(refs, _missing_in_memory(missing, user_api_key_cache), user_api_key_cache, prisma_client)
except NativeAuthFailure:
raise
except Exception as e: # noqa: BLE001 # warm-up only; the getters enforce and fail closed on their own
verbose_proxy_logger.warning("auth prefetch skipped, falling back to per-object lookups: %s", e)

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

@ -9,6 +9,7 @@ from litellm.proxy._types import LiteLLM_BudgetTable
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, get_management_object_ttl
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.repositories.budget_repository import BudgetRepository
from litellm.repositories.native_auth_repository import native_auth_repository
@dataclass(frozen=True, slots=True)
@ -29,14 +30,22 @@ async def read_default_budget(budget_id: str, io: EndUserIO) -> LiteLLM_BudgetTa
if cached is not None:
return cached
try:
table: Final = auth_checks._dictable_table(BudgetRepository(io.prisma), "budget") # pyright: ignore[reportPrivateUsage] # shared raw repository projection
record: Final = await table.find_unique(where={"budget_id": budget_id})
native: Final = native_auth_repository(io.prisma)
record: Final = (
await native.one("budget", [budget_id], LiteLLM_BudgetTable)
if native is not None
else await auth_checks._dictable_table(BudgetRepository(io.prisma), "budget").find_unique(
where={"budget_id": budget_id}
) # pyright: ignore[reportPrivateUsage] # shared raw repository projection
)
if record is None:
verbose_proxy_logger.warning(
"Default end user budget not found in database: %s", budget_id.replace("\r", "").replace("\n", "")
)
return None
row: Final = LiteLLM_BudgetTable.model_validate(record.dict())
row: Final = LiteLLM_BudgetTable.model_validate(
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)
)

17
litellm/proxy/auth/model_checks.py
View file

@ -14,14 +14,16 @@ from types import MappingProxyType
from typing import Final, TypeVar
import litellm
from litellm._auth.failure import NativeAuthFailure
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
from litellm.litellm_core_utils.litellm_logging import use_custom_pricing_for_model
from litellm.litellm_core_utils.pattern_match_deployments import PatternMatchRouter
from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping
from litellm.proxy._types import SpecialModelNames, UserAPIKeyAuth
from litellm.proxy._types import LiteLLM_ObjectPermissionTable, SpecialModelNames, UserAPIKeyAuth
from litellm.proxy.common_utils.model_listing_utils import CallerAliases, alias_map, dispatched_model
from litellm.proxy.common_utils.token_count_settings import cost_map_entry
from litellm.repositories.native_auth_repository import native_auth_repository
from litellm.repositories.object_permission_repository import ObjectPermissionRepository
from litellm.responses.main import strip_responses_routing_prefix
from litellm.router import Router
@ -132,12 +134,21 @@ async def get_mcp_server_ids(
# Make a direct SQL query to get just the mcp_servers
try:
result: Final = await ObjectPermissionRepository(prisma_client).table.find_unique(
where={"object_permission_id": user_api_key_dict.object_permission_id},
native: Final = native_auth_repository(prisma_client)
result: Final = (
await native.one(
"object_permission", [user_api_key_dict.object_permission_id], LiteLLM_ObjectPermissionTable
)
if native is not None
else await ObjectPermissionRepository(prisma_client).table.find_unique(
where={"object_permission_id": user_api_key_dict.object_permission_id},
)
)
if result and result.mcp_servers:
return result.mcp_servers
return []
except NativeAuthFailure:
raise
except Exception:
return []

33
litellm/proxy/db/spend_counter_reseed.py
View file

@ -30,6 +30,7 @@ from litellm.proxy.spend_tracking.spend_counter_batch import (
record_spend_counter_value,
)
from litellm.redis.stores.spend_counters import END_USER_COUNTER_PREFIX, SPEND_COUNTERS_TARGET, SpendCounterStore
from litellm.repositories.native_auth_repository import BudgetWindowRow, SpendRow, native_auth_repository
from litellm.repositories.organization_repository import OrganizationRepository
from litellm.repositories.project_repository import ProjectRepository
from litellm.repositories.table_repositories import (
@ -134,6 +135,20 @@ class SpendCounterReseed:
# or tag names that merely contain ":window:".
if SpendCounterReseed._is_key_or_team_window_counter(counter_key):
return None
native: Final = native_auth_repository(prisma_client)
if native is not None:
prefix, separator, suffix = counter_key.partition("spend:")
kind, delimiter, entity_id = suffix.partition(":")
if (
prefix
or not separator
or not delimiter
or kind not in {"key", "team", "team_member", "user", "org", "project"}
):
return None
if kind == "team_member" and ":" not in entity_id:
return None
return await native.spend("member" if kind == "team_member" else kind, entity_id)
try:
row: Final = await bounded_db_lookup(
SpendCounterReseed._counter_row(prisma_client, counter_key), name="spend_counter"
@ -181,6 +196,9 @@ class SpendCounterReseed:
async def end_user_from_db(prisma_client: Optional["PrismaClient"], counter_key: str) -> float | None:
if prisma_client is None or not counter_key.startswith(END_USER_COUNTER_PREFIX):
return None
native: Final = native_auth_repository(prisma_client)
if native is not None:
return await native.spend("end_user", counter_key[len(END_USER_COUNTER_PREFIX) :])
where: Final[LiteLLM_EndUserTableWhereUniqueInput] = {"user_id": counter_key[len(END_USER_COUNTER_PREFIX) :]}
try:
row: Final = await bounded_db_lookup(
@ -296,6 +314,15 @@ class SpendCounterReseed:
if row_entity_type is None:
return None
native: Final = native_auth_repository(prisma_client)
if native is not None:
window: Final = await native.one(
"budget_window", [entity_type, entity_id, window_duration], BudgetWindowRow
)
if window is None or _as_utc(window.window_start) < _as_utc(expected_window_start):
return None
return float(window.spend or 0.0)
try:
row: Final = await BudgetWindowSpendRepository(prisma_client).table.find_unique(
where={
@ -367,6 +394,12 @@ class SpendCounterReseed:
group_field: Final = _WINDOW_SPEND_LOG_FIELDS.get(entity_type)
if group_field is None:
return None
native: Final = native_auth_repository(prisma_client)
if native is not None:
aggregate: Final = await native.one(
"window_spend", [entity_type, entity_id, window_start.isoformat()], SpendRow
)
return float(aggregate.spend or 0.0) if aggregate is not None else 0.0
where: Final = {
group_field: entity_id,
"startTime": {"gte": window_start},

143
litellm/repositories/native_auth_repository.py Normal file
View file

@ -0,0 +1,143 @@
import time
from collections.abc import Sequence
from dataclasses import dataclass
from datetime import datetime
from inspect import getattr_static
from typing import Final, Protocol, TypeVar
from pydantic import BaseModel, ConfigDict, TypeAdapter
from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError
from litellm.litellm_core_utils.core_helpers import require_attributes
from litellm.models.budget import LiteLLM_BudgetTable
from litellm.models.team import Member
from litellm.proxy.db.db_lookup_gate import DBLookupDeadlineExceeded, bounded_db_lookup, db_lookup_gate
from litellm.types.proxy.auth.user_api_key_auth import LiteLLM_VerificationTokenView
_Model = TypeVar("_Model", bound=BaseModel)
_ROWS: Final = TypeAdapter(list[dict[str, object]])
_SPEND: Final[TypeAdapter[float | None]] = TypeAdapter(float | None)
_MEMBERS: Final = TypeAdapter(list[dict[str, object]])
class ModelAccessGroupRow(BaseModel):
model_config = ConfigDict(frozen=True)
access_group_name: str
spend: float | None = 0.0
litellm_budget_table: LiteLLM_BudgetTable | None = None
class EndUserRegistryRow(BaseModel):
model_config = ConfigDict(frozen=True)
user_id: str
class TagRegistryRow(BaseModel):
model_config = ConfigDict(frozen=True)
tag_name: str
class SpendRow(BaseModel):
model_config = ConfigDict(frozen=True)
spend: float | None = None
class BudgetWindowRow(SpendRow):
window_start: datetime
class NativeAuthDeadlineExceeded(NativeAuthFailure, DBLookupDeadlineExceeded):
pass
class AuthReadHandle(Protocol):
async def auth_read(self, kind: str, ids: list[str], writer: bool = False) -> object: ...
async def auth_spend(self, kind: str, id: str) -> object: ...
class _NativeAuthClient(Protocol):
@property
def native_auth_database(self) -> AuthReadHandle | None: ...
@dataclass(frozen=True, slots=True)
class NativeAuthRepository:
database: AuthReadHandle
async def key(self, token: str, *, writer: bool = False) -> LiteLLM_VerificationTokenView | None:
rows: Final = await self.read("key", [token], writer=writer)
if not rows:
return None
try:
row: Final = rows[0]
members: Final = _MEMBERS.validate_python(row.get("team_members_with_roles") or [])
matches: Final = tuple(
Member.model_validate(member)
for member in members
if member.get("user_id") is not None and member.get("user_id") == row.get("user_id")
)
projection: Final = LiteLLM_VerificationTokenView.model_validate(
{
**row,
"team_models": row.get("team_models") if row.get("team_models") is not None else [],
"team_blocked": row.get("team_blocked") if row.get("team_blocked") is not None else False,
"team_member": matches[-1] if matches else None,
"last_refreshed_at": time.time(),
}
)
return (
projection.model_copy(update={"expires": projection.expires.isoformat()})
if isinstance(projection.expires, datetime)
else projection
)
except Exception as error: # noqa: BLE001 # preserve terminal native projection errors
raise NativeAuthRuntimeError("Invalid native auth key projection") from error
async def read(self, kind: str, ids: Sequence[str], writer: bool = False) -> list[dict[str, object]]:
try:
value: Final = await bounded_db_lookup(self._read(kind, ids, writer), name=kind)
return _ROWS.validate_python(value)
except NativeAuthFailure:
raise
except DBLookupDeadlineExceeded as error:
raise NativeAuthDeadlineExceeded(error.lookup, error.deadline_seconds) from error
except Exception as error: # noqa: BLE001 # native errors must not enter legacy fallback paths
raise NativeAuthRuntimeError(f"Native auth {kind} read failed") from error
async def _read(self, kind: str, ids: Sequence[str], writer: bool) -> object:
async with db_lookup_gate.current():
return await self.database.auth_read(kind, list(ids), writer=writer)
async def one(self, kind: str, ids: Sequence[str], model: type[_Model], *, writer: bool = False) -> _Model | None:
rows: Final = await self.many(kind, ids, model, writer=writer)
return rows[0] if rows else None
async def many(self, kind: str, ids: Sequence[str], model: type[_Model], *, writer: bool = False) -> list[_Model]:
rows: Final = await self.read(kind, ids, writer=writer)
try:
return [model.model_validate(row) for row in rows]
except Exception as error: # noqa: BLE001 # invalid native projections are terminal
raise NativeAuthRuntimeError(f"Invalid native auth {kind} projection") from error
async def spend(self, kind: str, entity_id: str) -> float | None:
try:
value: Final = await bounded_db_lookup(self._spend(kind, entity_id), name="spend_counter")
return _SPEND.validate_python(value)
except NativeAuthFailure:
raise
except DBLookupDeadlineExceeded as error:
raise NativeAuthDeadlineExceeded(error.lookup, error.deadline_seconds) from error
except Exception as error: # noqa: BLE001 # a failed native read is not an authoritative miss
raise NativeAuthRuntimeError(f"Native auth {kind} spend read failed") from error
async def _spend(self, kind: str, entity_id: str) -> object:
async with db_lookup_gate.current():
return await self.database.auth_spend(kind, entity_id)
def native_auth_repository(client: object) -> NativeAuthRepository | None:
if getattr_static(client, "native_auth_database", None) is None:
return None
database: Final = require_attributes(client, _NativeAuthClient, "native_auth_database").native_auth_database
return NativeAuthRepository(database) if database is not None else None

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

@ -632,6 +632,37 @@ async def test_get_key_object_should_reconnect_once_on_db_connection_error():
)
@pytest.mark.asyncio
async def test_native_key_lookup_populates_cache_without_prisma_io():
native = SimpleNamespace(auth_read=AsyncMock(return_value=[{"token": "native-hash", "team_models": None, "team_blocked": None}]))
prisma = MagicMock(native_auth_database=native)
prisma.get_data = AsyncMock(side_effect=AssertionError("legacy read"))
prisma.attempt_db_reconnect = AsyncMock(side_effect=AssertionError("legacy reconnect"))
cache = UserApiKeyCache()
first = await get_key_object("native-hash", prisma, cache)
second = await get_key_object("native-hash", prisma, cache)
assert first.token == second.token == "native-hash"
assert first.team_models == second.team_models == []
assert native.auth_read.await_count == 1
prisma.get_data.assert_not_awaited()
prisma.attempt_db_reconnect.assert_not_awaited()
@pytest.mark.asyncio
async def test_native_key_transport_error_never_enters_prisma_reconnect():
from litellm._auth.failure import NativeAuthFailure
native = SimpleNamespace(auth_read=AsyncMock(side_effect=httpx.ConnectError("db offline")))
prisma = MagicMock(native_auth_database=native)
prisma.get_data = AsyncMock()
prisma.attempt_db_reconnect = AsyncMock()
with pytest.raises(NativeAuthFailure):
await fetch_key_object_from_db_with_reconnect("hash", prisma, None, None, check_db_only=True)
native.auth_read.assert_awaited_once_with("key", ["hash"], writer=True)
prisma.get_data.assert_not_awaited()
prisma.attempt_db_reconnect.assert_not_awaited()
@pytest.mark.asyncio
async def test_get_key_object_should_raise_if_reconnect_fails_on_db_connection_error():
mock_prisma_client = MagicMock()

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

@ -1,11 +1,14 @@
"""Counts the Redis round trips and DB queries auth object reads cost per cache regime, and checks that the
per-object getters still enforce on their own when the prefetch cannot help."""
import json
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock
import pytest
from fastapi import HTTPException
from litellm._auth.failure import NativeAuthFailure
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.proxy._types import (
LiteLLM_OrganizationTable,
@ -75,6 +78,51 @@ ALL_ROWS = {
}
@pytest.mark.asyncio
async def test_native_prefetch_fills_all_getters_with_one_read_and_preserves_optional_ids():
native = SimpleNamespace(auth_read=AsyncMock(return_value=[ALL_ROWS]))
prisma = _prisma()
prisma.native_auth_database = native
cache = _cache(None)
await prefetch_auth_objects(refs=_refs(), prisma_client=prisma, user_api_key_cache=cache)
user, team, member, org = await _read_all_through_getters(cache, prisma)
assert user is not None and user.user_id == USER_ID
assert team.team_id == TEAM_ID
assert member is not None and member.spend == 3.0
assert org is not None and org.litellm_budget_table.max_budget == 1000.0
assert native.auth_read.await_count == 1
kind, ids = native.auth_read.call_args.args
assert kind == "prefetch"
assert len(ids) == 1 and json.loads(ids[0]) == [USER_ID, TEAM_ID, USER_ID, ORG_ID, None]
assert prisma.db.mock_calls == []
await prefetch_auth_objects(refs=_refs(), prisma_client=prisma, user_api_key_cache=cache)
assert native.auth_read.await_count == 1
@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):
prisma = _prisma()
prisma.native_auth_database = SimpleNamespace(auth_read=AsyncMock(return_value=rows))
with pytest.raises(NativeAuthFailure):
await prefetch_auth_objects(
refs=AuthObjectRefs(user_id=USER_ID), prisma_client=prisma, user_api_key_cache=_cache(None)
)
assert prisma.db.mock_calls == []
@pytest.mark.asyncio
async def test_native_prefetch_transport_failure_is_not_best_effort():
prisma = _prisma()
prisma.native_auth_database = SimpleNamespace(auth_read=AsyncMock(side_effect=RuntimeError("offline")))
with pytest.raises(NativeAuthFailure):
await prefetch_auth_objects(refs=_refs(), prisma_client=prisma, user_api_key_cache=_cache(None))
assert prisma.db.mock_calls == []
def _prisma(rows: dict[str, object] | None = ALL_ROWS) -> MagicMock:
prisma = MagicMock(name="prisma_client")
prisma.db.query_first = AsyncMock(return_value=rows)

64
tests/unit/proxy/db/test_spend_counter_reseed.py
View file

@ -16,6 +16,7 @@ from unittest.mock import AsyncMock
import pytest
from litellm._auth.failure import NativeAuthFailure
from litellm.caching.dual_cache import DualCache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.constants import PROXY_DB_LOOKUP_MAX_CONCURRENCY
@ -27,6 +28,69 @@ from litellm.redis.stores.spend_counters import SpendCounterStore
WINDOW_START = datetime(2026, 8, 1, tzinfo=timezone.utc)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("counter", "kind", "entity_id"),
[
("spend:key:k", "key", "k"),
("spend:user:u", "user", "u"),
("spend:team:t", "team", "t"),
("spend:org:o", "org", "o"),
("spend:project:p", "project", "p"),
("spend:team_member:u:with:colon:t", "member", "u:with:colon:t"),
],
)
async def test_native_reseed_reads_authoritative_zero_without_a_prisma_handle(counter, kind, entity_id):
native = SimpleNamespace(auth_spend=AsyncMock(return_value=0.0))
prisma = SimpleNamespace(native_auth_database=native)
assert await SpendCounterReseed.from_db(prisma, counter) == 0.0
native.auth_spend.assert_awaited_once_with(kind, entity_id)
@pytest.mark.asyncio
async def test_native_reseed_distinguishes_missing_from_failure():
native = SimpleNamespace(auth_spend=AsyncMock(side_effect=[None, RuntimeError("offline")]))
prisma = SimpleNamespace(native_auth_database=native)
assert await SpendCounterReseed.from_db(prisma, "spend:user:u") is None
with pytest.raises(NativeAuthFailure):
await SpendCounterReseed.from_db(prisma, "spend:user:u")
assert native.auth_spend.await_count == 2
@pytest.mark.asyncio
async def test_native_end_user_floor_preserves_zero():
native = SimpleNamespace(auth_spend=AsyncMock(return_value=0.0))
assert (
await SpendCounterReseed.end_user_from_db(SimpleNamespace(native_auth_database=native), "spend:end_user:u")
== 0.0
)
native.auth_spend.assert_awaited_once_with("end_user", "u")
@pytest.mark.asyncio
@pytest.mark.parametrize("row", [None, {"spend": 99.0, "window_start": WINDOW_START - timedelta(days=1)}])
async def test_native_stale_or_missing_window_uses_native_spend_aggregate(row):
native = SimpleNamespace(auth_read=AsyncMock(side_effect=[[row] if row else [], [{"spend": 4.5}]]))
value = await SpendCounterReseed.window_from_db(
SimpleNamespace(native_auth_database=native), "Key", "hash", "1d", WINDOW_START
)
assert value == 4.5
assert native.auth_read.call_args_list[0].args == ("budget_window", ["Key", "hash", "1d"])
assert native.auth_read.call_args_list[1].args == ("window_spend", ["Key", "hash", WINDOW_START.isoformat()])
@pytest.mark.asyncio
async def test_native_current_window_does_not_range_scan_and_failure_does_not_fallback():
native = SimpleNamespace(
auth_read=AsyncMock(side_effect=[[{"spend": 0.0, "window_start": WINDOW_START}], RuntimeError("offline")])
)
prisma = SimpleNamespace(native_auth_database=native)
assert await SpendCounterReseed.window_from_db(prisma, "Key", "hash", "1d", WINDOW_START) == 0.0
with pytest.raises(NativeAuthFailure):
await SpendCounterReseed.window_from_db(prisma, "Key", "hash", "1d", WINDOW_START)
assert native.auth_read.await_count == 2
class _FakeFindUniqueTable:
def __init__(self, row: SimpleNamespace | None, error: Exception | None = None) -> None:
self._row = row

104
tests/unit/repositories/test_native_auth_repository.py Normal file
View file

@ -0,0 +1,104 @@
import asyncio
from datetime import datetime, timezone
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock
import pytest
from litellm._auth.failure import NativeAuthFailure
from litellm.models.budget import LiteLLM_BudgetTable
from litellm.models.team_membership import LiteLLM_TeamMembership
from litellm.repositories.native_auth_repository import NativeAuthRepository, native_auth_repository
@pytest.mark.asyncio
async def test_native_key_preserves_joined_budget_and_membership_projection():
raw = {
"token": "token-hash",
"user_id": "user",
"organization_id": "key-org",
"org_id": "team-org",
"team_models": None,
"team_blocked": None,
"max_budget": None,
"litellm_budget_table_max_budget": 25.0,
"model_max_budget": {},
"litellm_budget_table_model_max_budget": {"model": {"max_budget": 1.0}},
"team_model_aliases": {"alias": "model"},
"team_members_with_roles": [{"user_id": "user", "role": "user"}, {"user_id": "user", "role": "admin"}],
"expires": datetime(2027, 1, 1, tzinfo=timezone.utc),
}
handle = SimpleNamespace(auth_read=AsyncMock(return_value=[raw]))
row = await NativeAuthRepository(handle).key("token-hash", writer=True)
assert row is not None
assert (row.org_id, row.team_models, row.team_blocked, row.max_budget) == ("key-org", [], False, 25.0)
assert row.model_max_budget == {"model": {"max_budget": 1.0}}
assert row.team_model_aliases == {"alias": "model"}
assert row.team_member is not None and row.team_member.role == "admin"
assert row.expires == "2027-01-01T00:00:00+00:00"
assert row.last_refreshed_at is not None
assert raw["max_budget"] is None and raw["team_models"] is None
handle.auth_read.assert_awaited_once_with("key", ["token-hash"], writer=True)
@pytest.mark.asyncio
async def test_native_rows_keep_null_limits_and_make_request_local_models():
raw = {"user_id": "u", "team_id": "t", "spend": None, "litellm_budget_table": {"max_budget": None}}
handle = SimpleNamespace(auth_read=AsyncMock(return_value=[raw]))
repo = NativeAuthRepository(handle)
first = await repo.one("member", ["u", "t"], LiteLLM_TeamMembership)
second = await repo.one("member", ["u", "t"], LiteLLM_TeamMembership)
assert first is not None and second is not None
assert first.spend is None and second.spend is None
assert first.litellm_budget_table is not None and second.litellm_budget_table is not None
first.litellm_budget_table.max_budget = 3.0
assert second.litellm_budget_table.max_budget is None
assert raw["litellm_budget_table"] == {"max_budget": None}
@pytest.mark.asyncio
async def test_native_missing_row_is_none_and_zero_budget_is_not_missing():
handle = SimpleNamespace(auth_read=AsyncMock(side_effect=[[], [{"budget_id": "b", "max_budget": 0.0}]]))
repo = NativeAuthRepository(handle)
assert await repo.one("budget", ["missing"], LiteLLM_BudgetTable) is None
found = await repo.one("budget", ["b"], LiteLLM_BudgetTable)
assert found is not None and found.max_budget == 0.0
@pytest.mark.asyncio
@pytest.mark.parametrize("answer", [None, {}, ["bad"], [{"max_budget": "bad"}]])
async def test_native_invalid_rows_are_terminal(answer):
handle = SimpleNamespace(auth_read=AsyncMock(return_value=answer))
with pytest.raises(NativeAuthFailure):
await NativeAuthRepository(handle).one("budget", ["b"], LiteLLM_BudgetTable)
assert handle.auth_read.await_count == 1
@pytest.mark.asyncio
async def test_native_failure_and_cancellation_are_not_clean_misses():
handle = SimpleNamespace(auth_read=AsyncMock(side_effect=RuntimeError("connection closed")))
with pytest.raises(NativeAuthFailure, match="read failed"):
await NativeAuthRepository(handle).key("hash")
cancelled = SimpleNamespace(auth_read=AsyncMock(side_effect=asyncio.CancelledError))
with pytest.raises(asyncio.CancelledError):
await NativeAuthRepository(cancelled).key("hash")
@pytest.mark.asyncio
async def test_native_spend_preserves_missing_and_authoritative_zero():
handle = SimpleNamespace(auth_spend=AsyncMock(side_effect=[None, 0.0, RuntimeError("offline")]))
repo = NativeAuthRepository(handle)
assert await repo.spend("user", "u") is None
assert await repo.spend("user", "u") == 0.0
with pytest.raises(NativeAuthFailure):
await repo.spend("user", "u")
def test_repository_selection_requires_an_explicit_handle():
assert native_auth_repository(SimpleNamespace(native_auth_database=None)) is None
assert native_auth_repository(MagicMock()) is None
handle = SimpleNamespace(auth_read=AsyncMock())
repo = native_auth_repository(SimpleNamespace(native_auth_database=handle))
assert repo is not None and repo.database is handle