From b3bd09704e2247e8155fb42f338b66b3c714041f Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 9 Nov 2026 02:23:57 +0000 Subject: [PATCH] feat(auth): complete native effect reads and preserve spend IO fallback --- .../auth/agent_permission_handler.py | 10 ++- .../proxy/agent_endpoints/identity_store.py | 11 +++ litellm/proxy/auth/auth_checks.py | 24 ++++-- litellm/proxy/auth/end_user_context_io.py | 4 +- .../common_utils/registry_read_through.py | 14 +++- litellm/proxy/db/spend_counter_reseed.py | 48 ++++++++--- .../repositories/native_auth_repository.py | 71 +++++++++++++--- .../auth/test_agent_permission_handler.py | 26 ++++++ .../agent_endpoints/test_identity_store.py | 36 ++++++++- ...st_auth_checks_object_access_and_lookup.py | 34 ++++++++ .../test_registry_read_through.py | 80 ++++++++++++++++++- .../proxy/db/test_spend_counter_reseed.py | 34 +++++++- .../test_native_auth_repository.py | 45 +++++++++++ 13 files changed, 403 insertions(+), 34 deletions(-) diff --git a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py index d7e036b2915..e628fce03d4 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py +++ b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py @@ -32,6 +32,7 @@ from litellm.proxy.agent_endpoints.auth.agent_caller import agent_caller_auth from litellm.proxy.agent_endpoints.auth.managed_authorization import managed_agent_policy from litellm.proxy.mcp_server.ui_session_utils import build_effective_auth_contexts from litellm.proxy.utils import PrismaClient +from litellm.repositories.native_auth_repository import AgentIdRow, native_auth_repository from litellm.repositories.table_repositories import AgentsRepository from litellm.types.agents import AgentResponse @@ -597,8 +598,13 @@ class AgentRequestHandler: if not access_groups or prisma_client is None: return set() - agents: Final = await AgentsRepository(prisma_client, use_writer=check_db_only).table.find_many( - where={"agent_access_groups": {"hasSome": access_groups}} + native: Final = native_auth_repository(prisma_client) + agents: Final = ( + await native.many("agent_access_group", access_groups, AgentIdRow, writer=check_db_only) + if native is not None + else await AgentsRepository(prisma_client, use_writer=check_db_only).table.find_many( + where={"agent_access_groups": {"hasSome": access_groups}} + ) ) return {agent.agent_id for agent in agents} diff --git a/litellm/proxy/agent_endpoints/identity_store.py b/litellm/proxy/agent_endpoints/identity_store.py index c3e2064cbd5..9767ca94224 100644 --- a/litellm/proxy/agent_endpoints/identity_store.py +++ b/litellm/proxy/agent_endpoints/identity_store.py @@ -3,9 +3,11 @@ from collections.abc import Mapping from datetime import datetime, timezone from typing import TYPE_CHECKING, Final +from litellm._auth.failure import NativeAuthFailure from litellm._internal_context import with_service_target from litellm.proxy.agent_endpoints.managed_identity import classify_agent_subject from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, get_management_object_ttl +from litellm.repositories.native_auth_repository import NativeAuthRepository, native_auth_repository from litellm.repositories.table_repositories import ( AgentIdentityRepository, AgentsRepository, @@ -48,6 +50,7 @@ class AgentIdentityStore: RetiredAgentIdentityRepository(client, use_writer=True), RetiredAgentRepository(client, use_writer=True), cache=cache, + native=native_auth_repository(client), ) def __init__( @@ -59,6 +62,7 @@ class AgentIdentityStore: retired_agents: RetiredAgentRepository | None = None, *, cache: UserApiKeyCache | None = None, + native: NativeAuthRepository | None = None, ) -> None: self.agents = agents self.identities = identities @@ -66,9 +70,12 @@ class AgentIdentityStore: self.retired = retired self.retired_agents = retired_agents self.cache = cache + self.native = native async def agent(self, agent_id: str) -> AgentResponse | AgentIdentityFailure | None: try: + if self.native is not None: + return await self.native.one("agent", [agent_id], AgentResponse, writer=True) where: Final[LiteLLM_AgentsTableWhereUniqueInput] = {"agent_id": agent_id} include: Final[LiteLLM_AgentsTableInclude] = { "identity": True, @@ -78,6 +85,8 @@ class AgentIdentityStore: if row is None: return None return AgentResponse.model_validate(row.model_dump()) + except NativeAuthFailure: + raise except Exception: return AgentIdentityFailure(code="policy_unavailable", message="Agent policy could not be loaded") @@ -185,6 +194,8 @@ class AgentIdentityStore: return AgentIdentityFailure(code="policy_unavailable", message="Subject classification is unavailable") async def retired_agent(self, agent_id: str) -> bool | AgentIdentityFailure: + if self.native is not None: + return bool(await self.native.read("retired_agent", [agent_id], writer=True)) if self.retired_agents is None: return AgentIdentityFailure(code="policy_unavailable", message="Agent history is unavailable") try: diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index e33f87494f4..5ff44d739ac 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -2542,7 +2542,12 @@ async def _backfill_null_user_email( user_id=user_row.user_id, user_email=user_email, ) - db_row: Final = await user_repo.find_by_id(user_row.user_id) + native: Final = native_auth_repository(prisma_client) + db_row: Final = ( + await native.one("user", [user_row.user_id], LiteLLM_UserTable, writer=True) + if native is not None + else await user_repo.find_by_id(user_row.user_id) + ) if db_row is None: return user_row email_update: Final = {"user_email": db_row.user_email} @@ -4155,13 +4160,22 @@ async def get_managed_vector_store_rows_by_uuids( if not cache_misses: return result - rows: Final = await _vector_store_table(ManagedVectorStoresRepository(prisma_client)).find_many( - where={"vector_store_id": {"in": cache_misses}}, - take=len(cache_misses), + native: Final = native_auth_repository(prisma_client) + rows: Final = ( + await native.many("managed_vector_store", cache_misses, LiteLLM_ManagedVectorStoresTable) + if native is not None + else await _vector_store_table(ManagedVectorStoresRepository(prisma_client)).find_many( + where={"vector_store_id": {"in": cache_misses}}, + take=len(cache_misses), + ) ) for row in rows: - row_dict = row.model_dump() if hasattr(row, "model_dump") else (row.dict() if hasattr(row, "dict") else None) + row_dict = ( + row.model_dump() + if isinstance(row, BaseModel) or hasattr(row, "model_dump") + else (row.dict() if hasattr(row, "dict") else None) + ) if not isinstance(row_dict, dict) or not row_dict: row_dict = dict(row) if hasattr(row, "__dict__") else {} if not row_dict: diff --git a/litellm/proxy/auth/end_user_context_io.py b/litellm/proxy/auth/end_user_context_io.py index e39ad797bea..82222f51a99 100644 --- a/litellm/proxy/auth/end_user_context_io.py +++ b/litellm/proxy/auth/end_user_context_io.py @@ -34,9 +34,9 @@ async def read_default_budget(budget_id: str, io: EndUserIO) -> LiteLLM_BudgetTa 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( + else await auth_checks._dictable_table(BudgetRepository(io.prisma), "budget").find_unique( # pyright: ignore[reportPrivateUsage] # shared raw repository projection where={"budget_id": budget_id} - ) # pyright: ignore[reportPrivateUsage] # shared raw repository projection + ) ) if record is None: verbose_proxy_logger.warning( diff --git a/litellm/proxy/common_utils/registry_read_through.py b/litellm/proxy/common_utils/registry_read_through.py index d58c58fa1a5..1adbc1b0620 100644 --- a/litellm/proxy/common_utils/registry_read_through.py +++ b/litellm/proxy/common_utils/registry_read_through.py @@ -19,10 +19,12 @@ from typing import TYPE_CHECKING, Final, Literal, TypeAlias from pydantic import TypeAdapter, ValidationError +from litellm._auth.failure import NativeAuthFailure from litellm._logging import verbose_proxy_logger from litellm.caching.in_memory_cache import InMemoryCache from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper +from litellm.repositories.native_auth_repository import native_auth_repository if TYPE_CHECKING: from prisma import models as prisma_models @@ -106,6 +108,8 @@ class RegistryReadThrough: return False try: found: Final = await self._resync(key) + except NativeAuthFailure: + raise except Exception as e: # noqa: BLE001 # a failed read-through must surface the original miss error, not a 500 verbose_proxy_logger.warning("registry read-through for %r failed: %s", key, e) return False @@ -280,7 +284,6 @@ async def _resync_agents(agent_id_or_name: str) -> bool: return True prisma_client: Final = proxy_server.prisma_client assert prisma_client is not None - table: Final = agents_table(prisma_client) id_filter: Final[LiteLLM_AgentsTableWhereUniqueInput] = {"agent_id": agent_id_or_name} name_filter: Final[LiteLLM_AgentsTableWhereUniqueInput] = {"agent_name": agent_id_or_name} include_permission: Final[LiteLLM_AgentsTableInclude] = { @@ -290,8 +293,13 @@ async def _resync_agents(agent_id_or_name: str) -> bool: async with AGENT_RECONCILE_LOCK: if _agent_from_registry(agent_id_or_name) is not None: return True - row: Final = await table.find_unique(where=id_filter, include=include_permission) or await table.find_unique( - where=name_filter, include=include_permission + native: Final = native_auth_repository(prisma_client) + row: Final = ( + await native.one("agent", [agent_id_or_name], AgentResponse) + or await native.one("agent_name", [agent_id_or_name], AgentResponse) + if native is not None + else await agents_table(prisma_client).find_unique(where=id_filter, include=include_permission) + or await agents_table(prisma_client).find_unique(where=name_filter, include=include_permission) ) if row is None: return False diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index a944bd6b513..c44d0e2d47f 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -30,7 +30,12 @@ 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.native_auth_repository import ( + BudgetWindowRow, + NativeAuthDatabaseError, + SpendRow, + native_auth_repository, +) from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.project_repository import ProjectRepository from litellm.repositories.table_repositories import ( @@ -148,7 +153,11 @@ class SpendCounterReseed: 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: + return await native.spend("member" if kind == "team_member" else kind, entity_id) + except NativeAuthDatabaseError: + verbose_proxy_logger.exception("SpendCounterReseed.from_db: native read failed for %s", counter_key) + return None try: row: Final = await bounded_db_lookup( SpendCounterReseed._counter_row(prisma_client, counter_key), name="spend_counter" @@ -198,7 +207,13 @@ class SpendCounterReseed: 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) :]) + try: + return await native.spend("end_user", counter_key[len(END_USER_COUNTER_PREFIX) :]) + except NativeAuthDatabaseError: + verbose_proxy_logger.exception( + "SpendCounterReseed.end_user_from_db: native read failed for %s", counter_key + ) + return None where: Final[LiteLLM_EndUserTableWhereUniqueInput] = {"user_id": counter_key[len(END_USER_COUNTER_PREFIX) :]} try: row: Final = await bounded_db_lookup( @@ -316,9 +331,18 @@ class SpendCounterReseed: 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 - ) + try: + window: Final = await native.one( + "budget_window", [entity_type, entity_id, window_duration], BudgetWindowRow + ) + except NativeAuthDatabaseError: + verbose_proxy_logger.exception( + "SpendCounterReseed.window_from_table: native read failed for %s=%s window=%s", + entity_type, + entity_id, + window_duration, + ) + return None if window is None or _as_utc(window.window_start) < _as_utc(expected_window_start): return None return float(window.spend or 0.0) @@ -396,9 +420,15 @@ class SpendCounterReseed: 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 - ) + try: + aggregate: Final = await native.one( + "window_spend", [entity_type, entity_id, window_start.isoformat()], SpendRow + ) + except NativeAuthDatabaseError: + verbose_proxy_logger.exception( + "SpendCounterReseed.window_from_spend_logs: native read failed for %s=%s", entity_type, entity_id + ) + return None return float(aggregate.spend or 0.0) if aggregate is not None else 0.0 where: Final = { group_field: entity_id, diff --git a/litellm/repositories/native_auth_repository.py b/litellm/repositories/native_auth_repository.py index 74f3e09ece2..e3d9a2f4752 100644 --- a/litellm/repositories/native_auth_repository.py +++ b/litellm/repositories/native_auth_repository.py @@ -1,8 +1,9 @@ import time -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from dataclasses import dataclass from datetime import datetime from inspect import getattr_static +from itertools import chain from typing import Final, Protocol, TypeVar from pydantic import BaseModel, ConfigDict, TypeAdapter @@ -14,10 +15,30 @@ 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) +_Model = TypeVar("_Model", bound=BaseModel) # rebind-ok: TypeVar declarations cannot be annotated Final _ROWS: Final = TypeAdapter(list[dict[str, object]]) _SPEND: Final[TypeAdapter[float | None]] = TypeAdapter(float | None) -_MEMBERS: Final = TypeAdapter(list[dict[str, object]]) +_BATCH_ID_KINDS: Final = frozenset( + { + "key", + "user", + "team", + "org", + "project", + "end_user", + "budget", + "object_permission", + "access_group", + "tag", + "model_access_group", + "agent", + "agent_name", + "retired_agent", + "managed_vector_store", + "agent_access_group", + } +) +_BATCH_SIZE: Final = 1000 class ModelAccessGroupRow(BaseModel): @@ -37,6 +58,11 @@ class TagRegistryRow(BaseModel): tag_name: str +class AgentIdRow(BaseModel): + model_config = ConfigDict(frozen=True) + agent_id: str + + class SpendRow(BaseModel): model_config = ConfigDict(frozen=True) spend: float | None = None @@ -46,12 +72,21 @@ class BudgetWindowRow(SpendRow): window_start: datetime -class NativeAuthDeadlineExceeded(NativeAuthFailure, DBLookupDeadlineExceeded): +class NativeAuthDatabaseError(NativeAuthFailure): + pass + + +class NativeAuthDeadlineExceeded(NativeAuthDatabaseError, DBLookupDeadlineExceeded): pass class AuthReadHandle(Protocol): - async def auth_read(self, kind: str, ids: list[str], writer: bool = False) -> object: ... + async def auth_read( + self, + kind: str, + ids: list[str], # mutable-ok: PyO3 requires a Python list at this boundary + writer: bool = False, + ) -> object: ... async def auth_spend(self, kind: str, id: str) -> object: ... @@ -71,7 +106,7 @@ class NativeAuthRepository: return None try: row: Final = rows[0] - members: Final = _MEMBERS.validate_python(row.get("team_members_with_roles") or []) + members: Final = _ROWS.validate_python(row.get("team_members_with_roles") or []) matches: Final = tuple( Member.model_validate(member) for member in members @@ -94,7 +129,7 @@ class NativeAuthRepository: 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]]: + async def read(self, kind: str, ids: Sequence[str], writer: bool = False) -> Sequence[Mapping[str, object]]: try: value: Final = await bounded_db_lookup(self._read(kind, ids, writer), name=kind) return _ROWS.validate_python(value) @@ -102,6 +137,8 @@ class NativeAuthRepository: raise except DBLookupDeadlineExceeded as error: raise NativeAuthDeadlineExceeded(error.lookup, error.deadline_seconds) from error + except OSError as error: + raise NativeAuthDatabaseError(f"Native auth {kind} database read failed") 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 @@ -113,13 +150,25 @@ class NativeAuthRepository: 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) + async def many( + self, kind: str, ids: Sequence[str], model: type[_Model], *, writer: bool = False + ) -> Sequence[_Model]: + rows: Final = await self._many_rows(kind, ids, writer=writer) try: - return [model.model_validate(row) for row in rows] + return tuple(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 _many_rows(self, kind: str, ids: Sequence[str], *, writer: bool) -> Sequence[Mapping[str, object]]: + if kind not in _BATCH_ID_KINDS or len(ids) <= _BATCH_SIZE: + return await self.read(kind, ids, writer=writer) + unique_ids: Final = tuple(dict.fromkeys(ids)) + chunks: Final = [ + await self.read(kind, unique_ids[start : start + _BATCH_SIZE], writer=writer) + for start in range(0, len(unique_ids), _BATCH_SIZE) + ] + return tuple(chain.from_iterable(chunks)) + 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") @@ -128,6 +177,8 @@ class NativeAuthRepository: raise except DBLookupDeadlineExceeded as error: raise NativeAuthDeadlineExceeded(error.lookup, error.deadline_seconds) from error + except OSError as error: + raise NativeAuthDatabaseError(f"Native auth {kind} database spend read failed") 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 diff --git a/tests/unit/proxy/agent_endpoints/auth/test_agent_permission_handler.py b/tests/unit/proxy/agent_endpoints/auth/test_agent_permission_handler.py index c6406b82c15..bb967550c96 100644 --- a/tests/unit/proxy/agent_endpoints/auth/test_agent_permission_handler.py +++ b/tests/unit/proxy/agent_endpoints/auth/test_agent_permission_handler.py @@ -57,6 +57,32 @@ async def _single_context(user_api_key_auth: UserAPIKeyAuth) -> list[UserAPIKeyA return [user_api_key_auth] +@pytest.mark.asyncio +@pytest.mark.parametrize("writer", [False, True]) +async def test_native_legacy_agent_group_lookup_reads_selected_pool_without_prisma(writer): + from types import SimpleNamespace + + native = SimpleNamespace(auth_read=AsyncMock(return_value=[{"agent_id": "agent-a"}, {"agent_id": "agent-b"}])) + ids = await AgentRequestHandler._get_db_agent_ids_for_access_groups( + SimpleNamespace(native_auth_database=native), ["legacy-group"], check_db_only=writer + ) + assert ids == {"agent-a", "agent-b"} + native.auth_read.assert_awaited_once_with("agent_access_group", ["legacy-group"], writer=writer) + + +@pytest.mark.asyncio +async def test_native_legacy_agent_group_error_is_not_empty_grants(): + from types import SimpleNamespace + + from litellm._auth.failure import NativeAuthFailure + + native = SimpleNamespace(auth_read=AsyncMock(side_effect=OSError("offline"))) + with pytest.raises(NativeAuthFailure): + await AgentRequestHandler._get_db_agent_ids_for_access_groups( + SimpleNamespace(native_auth_database=native), ["legacy-group"] + ) + + @pytest.mark.asyncio class TestAgentRequestHandler: """ diff --git a/tests/unit/proxy/agent_endpoints/test_identity_store.py b/tests/unit/proxy/agent_endpoints/test_identity_store.py index 005f0b4c074..4ea0126e0ce 100644 --- a/tests/unit/proxy/agent_endpoints/test_identity_store.py +++ b/tests/unit/proxy/agent_endpoints/test_identity_store.py @@ -74,13 +74,47 @@ def setup_store( ) ) return ( - AgentIdentityStore(AgentsRepository(db), AgentIdentityRepository(db), VerifiedSubjectRepository(db), cache=cache), + AgentIdentityStore( + AgentsRepository(db), AgentIdentityRepository(db), VerifiedSubjectRepository(db), cache=cache + ), agents, identities, humans, ) +@pytest.mark.asyncio +async def test_native_agent_policy_reads_writer_and_preserves_binding_without_prisma() -> None: + expected: Final = stored_agent(enabled=False, object_permission={"models": ["approved"]}) + native: Final = SimpleNamespace(auth_read=AsyncMock(return_value=[expected.model_dump()])) + store: Final = AgentIdentityStore.from_client(SimpleNamespace(native_auth_database=native)) + actual: Final = await store.agent("agent-one") + assert actual == expected + native.auth_read.assert_awaited_once_with("agent", ["agent-one"], writer=True) + + +@pytest.mark.asyncio +async def test_native_retired_agent_lookup_preserves_missing_and_found_without_prisma() -> None: + native: Final = SimpleNamespace(auth_read=AsyncMock(side_effect=[[], [{"original_agent_id": "agent-one"}]])) + store: Final = AgentIdentityStore.from_client(SimpleNamespace(native_auth_database=native)) + assert await store.retired_agent("missing") is False + assert await store.retired_agent("agent-one") is True + native.auth_read.assert_awaited_with("retired_agent", ["agent-one"], writer=True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("error", [OSError("offline"), RuntimeError("invalid native state")]) +async def test_native_agent_failure_cannot_enter_legacy_fallback(error: Exception) -> None: + from litellm._auth.failure import NativeAuthFailure + + native: Final = SimpleNamespace(auth_read=AsyncMock(side_effect=error)) + store: Final = AgentIdentityStore.from_client(SimpleNamespace(native_auth_database=native)) + with pytest.raises(NativeAuthFailure): + await store.agent("agent-one") + with pytest.raises(NativeAuthFailure): + await store.retired_agent("agent-one") + + @pytest.mark.asyncio async def test_application_authentication_has_no_fabricated_human() -> None: store, _, _, humans = setup_store() diff --git a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py index 8f4404566fd..1a225b2571c 100644 --- a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py +++ b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py @@ -9726,6 +9726,40 @@ async def test_get_managed_vector_store_rows_by_uuids_db_fetch_validates_rows(): assert result[0].custom_llm_provider == "openai" +@pytest.mark.asyncio +async def test_native_managed_vector_store_read_preserves_owner_and_warms_cache(): + from litellm.proxy.auth.auth_checks import get_managed_vector_store_rows_by_uuids + + native = SimpleNamespace(auth_read=AsyncMock(return_value=[{ + "vector_store_id": "vs-native", "custom_llm_provider": "openai", + "team_id": "team", "user_id": "owner", "litellm_params": {"vector_store_id": "upstream"}, + }])) + prisma = SimpleNamespace(native_auth_database=native) + cache = UserApiKeyCache() + rows = await get_managed_vector_store_rows_by_uuids(["vs-native"], prisma, cache) + cached = await get_managed_vector_store_rows_by_uuids(["vs-native"], prisma, cache) + assert len(rows) == len(cached) == 1 + assert rows[0].team_id == cached[0].team_id == "team" + assert rows[0].user_id == cached[0].user_id == "owner" + assert rows[0].litellm_params == {"vector_store_id": "upstream"} + native.auth_read.assert_awaited_once_with("managed_vector_store", ["vs-native"], writer=False) + + +@pytest.mark.asyncio +async def test_native_email_backfill_reads_the_writer_after_the_atomic_update(): + from litellm.proxy.auth.auth_checks import _backfill_null_user_email + + native = SimpleNamespace(auth_read=AsyncMock(return_value=[{"user_id": "u", "user_email": "winner@example.test"}])) + prisma = MagicMock(native_auth_database=native) + prisma.db.litellm_usertable.update_many = AsyncMock(return_value=0) + prisma.db.litellm_usertable.find_unique = AsyncMock(side_effect=AssertionError("legacy read")) + row = await _backfill_null_user_email(prisma, UserApiKeyCache(), LiteLLM_UserTable(user_id="u"), "loser@example.test") + assert row.user_email == "winner@example.test" + native.auth_read.assert_awaited_once_with("user", ["u"], writer=True) + prisma.db.litellm_usertable.update_many.assert_awaited_once() + prisma.db.litellm_usertable.find_unique.assert_not_awaited() + + @pytest.mark.asyncio async def test_get_project_object_db_fetch_returns_cached_obj(): from litellm.proxy._types import LiteLLM_ProjectTableCachedObj diff --git a/tests/unit/proxy/common_utils/test_registry_read_through.py b/tests/unit/proxy/common_utils/test_registry_read_through.py index 6c9104bb711..c5c31dcfa21 100644 --- a/tests/unit/proxy/common_utils/test_registry_read_through.py +++ b/tests/unit/proxy/common_utils/test_registry_read_through.py @@ -80,6 +80,19 @@ async def test_resync_exception_returns_false_without_negative_caching(): assert spy.calls == ["new-model", "new-model"] +@pytest.mark.asyncio +async def test_native_read_through_failure_is_terminal_and_not_negative_cached(): + from litellm._auth.failure import NativeAuthRuntimeError + + spy: Final = ResyncSpy(error=NativeAuthRuntimeError("native connection lost")) + read_through: Final = RegistryReadThrough(resync=spy, is_loaded=nothing_loaded) + with pytest.raises(NativeAuthRuntimeError): + await read_through.attempt("agent") + with pytest.raises(NativeAuthRuntimeError): + await read_through.attempt("agent") + assert spy.calls == ["agent", "agent"] + + @pytest.mark.asyncio async def test_concurrent_attempts_for_missing_key_resync_once(): class SlowResyncSpy(ResyncSpy): @@ -257,6 +270,69 @@ async def test_get_agent_with_read_through_recovers_agent_by_name(clean_agent_re ) +@pytest.mark.asyncio +@pytest.mark.parametrize("by_name", [False, True]) +async def test_native_agent_read_through_preserves_identity_and_permission_without_prisma( + by_name, clean_agent_registry, fresh_agent_read_through, monkeypatch +): + from types import SimpleNamespace + from unittest.mock import AsyncMock + + from litellm.proxy.common_utils.registry_read_through import get_agent_with_read_through + + binding = { + "agent_id": "native-agent", + "provider": "microsoft_entra", + "tenant_id": "tenant", + "client_id": "client", + "issuer": "https://login.microsoftonline.com/tenant/v2.0", + "revision": "revision", + } + row = { + **FakeAgentRow("native-agent", "Native agent name").model_dump(), + "identity": binding, + "enabled": False, + "access_group_ids": ["group"], + "object_permission": {"models": ["approved-model"]}, + } + native = SimpleNamespace(auth_read=AsyncMock(side_effect=[[], [row]] if by_name else [[row]])) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", SimpleNamespace(native_auth_database=native)) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) + lookup = "Native agent name" if by_name else "native-agent" + + agent = await get_agent_with_read_through(lookup) + assert agent is not None and agent.identity is not None + assert agent.identity.model_dump(include=set(binding)) == binding + assert agent.object_permission == {"models": ["approved-model"]} + assert agent.access_group_ids == ["group"] and agent.enabled is False + assert clean_agent_registry.get_agent_by_id(agent_id="native-agent") == agent + assert [call.args[0] for call in native.auth_read.call_args_list] == ( + ["agent", "agent_name"] if by_name else ["agent"] + ) + + await get_agent_with_read_through(lookup) + assert native.auth_read.await_count == (2 if by_name else 1) + + +@pytest.mark.asyncio +async def test_native_agent_read_through_error_does_not_retry_by_name_or_return_miss( + clean_agent_registry, fresh_agent_read_through, monkeypatch +): + from types import SimpleNamespace + from unittest.mock import AsyncMock + + from litellm._auth.failure import NativeAuthFailure + from litellm.proxy.common_utils.registry_read_through import get_agent_with_read_through + + native = SimpleNamespace(auth_read=AsyncMock(side_effect=RuntimeError("offline"))) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", SimpleNamespace(native_auth_database=native)) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) + with pytest.raises(NativeAuthFailure): + await get_agent_with_read_through("native-failed-agent") + native.auth_read.assert_awaited_once_with("agent", ["native-failed-agent"], writer=False) + assert clean_agent_registry.get_agent_by_id(agent_id="native-failed-agent") is None + + @pytest.mark.asyncio async def test_get_agent_with_read_through_returns_none_for_unknown_agent( clean_agent_registry, fresh_agent_read_through, monkeypatch @@ -1083,7 +1159,9 @@ async def test_resync_model_deployments_by_ids_reloads_credentials_only_for_one_ model_params: Final = {"model": "openai/gpt-backup"} row: Final = _stored_model_row( "backup-id", - model_params if stored_credential_name is None else model_params | {"litellm_credential_name": stored_credential_name}, + model_params + if stored_credential_name is None + else model_params | {"litellm_credential_name": stored_credential_name}, ) stored_credential: Final = CredentialItem( credential_name="shared-openai", diff --git a/tests/unit/proxy/db/test_spend_counter_reseed.py b/tests/unit/proxy/db/test_spend_counter_reseed.py index 62ceff452e7..cfb915db8e4 100644 --- a/tests/unit/proxy/db/test_spend_counter_reseed.py +++ b/tests/unit/proxy/db/test_spend_counter_reseed.py @@ -48,7 +48,7 @@ async def test_native_reseed_reads_authoritative_zero_without_a_prisma_handle(co @pytest.mark.asyncio -async def test_native_reseed_distinguishes_missing_from_failure(): +async def test_native_reseed_distinguishes_missing_from_invariant_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 @@ -57,6 +57,38 @@ async def test_native_reseed_distinguishes_missing_from_failure(): assert native.auth_spend.await_count == 2 +@pytest.mark.asyncio +@pytest.mark.parametrize("error", [OSError("connection lost"), asyncio.TimeoutError()]) +async def test_native_spend_io_failure_preserves_cached_fallback_without_prisma(error): + native = SimpleNamespace(auth_spend=AsyncMock(side_effect=error)) + prisma = SimpleNamespace(native_auth_database=native) + assert await SpendCounterReseed.from_db(prisma, "spend:key:hash") is None + assert await SpendCounterReseed.end_user_from_db(prisma, "spend:end_user:u") is None + assert native.auth_spend.await_count == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("aggregate", [[{"spend": 4.5}], OSError("connection still lost")]) +async def test_native_window_io_failure_uses_only_native_aggregate_and_preserves_unavailable(aggregate): + native = SimpleNamespace(auth_read=AsyncMock(side_effect=[OSError("connection lost"), aggregate])) + result = await SpendCounterReseed.window_from_db( + SimpleNamespace(native_auth_database=native), "Key", "hash", "1d", WINDOW_START + ) + assert result == (None if isinstance(aggregate, OSError) else 4.5) + assert [call.args[0] for call in native.auth_read.call_args_list] == ["budget_window", "window_spend"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("payload", ["not a number", {}, []]) +async def test_native_invalid_spend_cannot_become_cached_fallback(payload): + native = SimpleNamespace(auth_spend=AsyncMock(return_value=payload)) + prisma = SimpleNamespace(native_auth_database=native) + with pytest.raises(NativeAuthFailure): + await SpendCounterReseed.from_db(prisma, "spend:key:hash") + with pytest.raises(NativeAuthFailure): + await SpendCounterReseed.end_user_from_db(prisma, "spend:end_user:u") + + @pytest.mark.asyncio async def test_native_end_user_floor_preserves_zero(): native = SimpleNamespace(auth_spend=AsyncMock(return_value=0.0)) diff --git a/tests/unit/repositories/test_native_auth_repository.py b/tests/unit/repositories/test_native_auth_repository.py index 686135e123e..b0e50a288d5 100644 --- a/tests/unit/repositories/test_native_auth_repository.py +++ b/tests/unit/repositories/test_native_auth_repository.py @@ -58,6 +58,40 @@ async def test_native_rows_keep_null_limits_and_make_request_local_models(): assert raw["litellm_budget_table"] == {"max_budget": None} +@pytest.mark.asyncio +async def test_native_large_id_batch_is_chunked_deduplicated_and_keeps_missing_rows_absent(): + async def read(kind, ids, writer=False): + assert kind == "budget" and writer is True + assert len(ids) <= 1000 + return [{"budget_id": entity_id} for entity_id in ids if entity_id != "missing"] + + ids = tuple(f"budget-{i}" for i in range(1001)) + handle = SimpleNamespace(auth_read=AsyncMock(side_effect=read)) + rows = await NativeAuthRepository(handle).many( + "budget", (*ids, ids[0], "missing"), LiteLLM_BudgetTable, writer=True + ) + assert tuple(row.budget_id for row in rows) == ids + assert [len(call.args[1]) for call in handle.auth_read.call_args_list] == [1000, 2] + + +@pytest.mark.asyncio +async def test_native_large_batch_error_never_returns_partial_rows(): + from litellm.repositories.native_auth_repository import NativeAuthDatabaseError + + handle = SimpleNamespace(auth_read=AsyncMock(side_effect=[[{"budget_id": "first"}], OSError("offline")])) + with pytest.raises(NativeAuthDatabaseError): + await NativeAuthRepository(handle).many("budget", tuple(str(i) for i in range(1001)), LiteLLM_BudgetTable) + assert handle.auth_read.await_count == 2 + + +@pytest.mark.asyncio +async def test_native_empty_batch_preserves_native_empty_read_contract(): + handle = SimpleNamespace(auth_read=AsyncMock(return_value=[])) + rows = await NativeAuthRepository(handle).many("budget", (), LiteLLM_BudgetTable) + assert rows == () + handle.auth_read.assert_awaited_once_with("budget", [], writer=False) + + @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}]])) @@ -86,6 +120,17 @@ async def test_native_failure_and_cancellation_are_not_clean_misses(): await NativeAuthRepository(cancelled).key("hash") +@pytest.mark.asyncio +async def test_native_timeout_keeps_deadline_identity_and_terminal_native_semantics(): + from litellm.proxy.db.db_lookup_gate import DBLookupDeadlineExceeded + + handle = SimpleNamespace(auth_read=AsyncMock(side_effect=asyncio.TimeoutError)) + with pytest.raises(NativeAuthFailure) as raised: + await NativeAuthRepository(handle).key("hash") + assert isinstance(raised.value, DBLookupDeadlineExceeded) + assert raised.value.lookup == "key" + + @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")]))