mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
feat(auth): complete native effect reads and preserve spend IO fallback
This commit is contained in:
Devin AI
parent
e00029bd69
commit
e5c28a5b1f
13 changed files with 403 additions and 34 deletions
|
|
@ -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}
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
|
|
|||
|
|
@ -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")]))
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue