diff --git a/litellm/proxy/auth/resolvers/store.py b/litellm/proxy/auth/resolvers/store.py index 0c8d05e0b94..2924aa1372b 100644 --- a/litellm/proxy/auth/resolvers/store.py +++ b/litellm/proxy/auth/resolvers/store.py @@ -1,10 +1,12 @@ from __future__ import annotations +import time from collections.abc import Sequence from typing import TYPE_CHECKING, Final from pydantic import BaseModel +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm.proxy._types import UserAPIKeyAuth @@ -36,7 +38,12 @@ from litellm.proxy.auth.resolvers.models import ( UserIdentity, ) from litellm.proxy.auth.roles import TeamRole, map_role, team_role -from litellm.proxy.common_utils.user_api_key_cache import AUTH_OBJECTS_TARGET +from litellm.proxy.common_utils.user_api_key_cache import ( + AUTH_OBJECTS_TARGET, + UserApiKeyCache, + get_management_object_ttl, +) +from litellm.repositories.native_auth_repository import NativeAuthDatabaseError if TYPE_CHECKING: from litellm.caching.caching import DualCache @@ -108,10 +115,25 @@ class IdentityStore: async def _resolve_key(self, hashed_token: str) -> UserAPIKeyAuth: if self._prisma is None: raise NoDatabaseConnectionError() + first: Final = await self._resolve_key_once(hashed_token) + if first is not None: + return first + retried: Final = await self._resolve_key_once(hashed_token) + if retried is not None: + return retried + raise NativeAuthRuntimeError("Auth identity changed during lookup") + async def _resolve_key_once(self, hashed_token: str) -> UserAPIKeyAuth | None: + assert self._prisma is not None + cache: Final = self._cache if isinstance(self._cache, UserApiKeyCache) else None + generation: Final = cache.begin_fill(hashed_token) if cache is not None else None cached: Final = await self._cache.async_get_cache(key=hashed_token, model_type=UserAPIKeyAuth) if cached is not None: - return copy_user_api_key_auth_for_cache(user_api_key_obj=cached) + return ( + copy_user_api_key_auth_for_cache(user_api_key_obj=cached) + if cache is None or cache.current_generation() == generation + else None + ) if self._check_cache_only: raise KeyNotInCacheError(hashed_token) @@ -137,12 +159,21 @@ class IdentityStore: proxy_logging_obj=self._proxy_logging_obj, ) except Exception as e: + if isinstance(e, NativeAuthFailure) and not isinstance(e, NativeAuthDatabaseError): + raise verbose_proxy_logger.debug( "Failed to load object_permission for key with object_permission_id=%s: %s", key.object_permission_id, e, ) + if cache is not None and generation is not None: + refreshed: Final = key.model_copy(update={"last_refreshed_at": time.time()}) + stored: Final = cache.set_if_generation( + hashed_token, refreshed, generation, ttl=get_management_object_ttl(cache) + ) + return refreshed if stored or cache.current_generation() == generation else None + await cache_key_object( hashed_token=hashed_token, user_api_key_obj=key, diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 2a55face8ce..4b1f7e513f6 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2355,11 +2355,18 @@ class _BuilderEffects: return TeamMemberBudget(spend=spend, max_budget=max_budget) async def _load_team_membership(self, user_id: str, team_id: str) -> LiteLLM_TeamMembership | None: + from litellm.repositories.native_auth_repository import native_auth_repository + if self._prisma_client is None: return None - row: Final = await TeamMembershipRepository(self._prisma_client).table.find_first( - where={"user_id": user_id, "team_id": team_id}, - include={"litellm_budget_table": True}, + native: Final = native_auth_repository(self._prisma_client) + row: Final = ( + await native.one("member", [user_id, team_id], LiteLLM_TeamMembership) + if native is not None + else await TeamMembershipRepository(self._prisma_client).table.find_first( + where={"user_id": user_id, "team_id": team_id}, + include={"litellm_budget_table": True}, + ) ) if row is None: return None diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3ad323f8c1f..0cdc4ed1575 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -8305,6 +8305,9 @@ class ProxyConfig: Returns: dict: router_settings, or None if no settings found """ + from litellm._auth.failure import NativeAuthFailure + from litellm.repositories.native_auth_repository import NativeAuthDatabaseError + # 1. Try key-level router_settings # user_api_key_dict is already the cached/authenticated key object — # no DB call needed. @@ -8327,6 +8330,10 @@ class ProxyConfig: team_settings: Final = self._parse_router_settings_value(getattr(team_obj, "router_settings", None)) if team_settings is not None: return team_settings + except NativeAuthDatabaseError: + pass + except NativeAuthFailure: + raise except Exception: # If team lookup fails, no team-level settings available pass diff --git a/litellm/repositories/database.py b/litellm/repositories/database.py index 5beb75b1ca6..ab6e677bcd4 100644 --- a/litellm/repositories/database.py +++ b/litellm/repositories/database.py @@ -83,17 +83,22 @@ async def _refresh_password(native: NativeDatabase, settings: NativeDatabaseSett return await asyncio.sleep(_refresh_delay(settings, settings.password)) while True: - try: - password: Final = await asyncio.to_thread(settings.password_provider) - if reader: - native.set_passwords(reader_password=password) - else: - native.set_passwords(writer_password=password) - except Exception: - verbose_proxy_logger.warning("Native database credential refresh failed") - await asyncio.sleep(30.0) + await asyncio.sleep(await _rotate_password(native, settings, reader)) + + +async def _rotate_password(native: NativeDatabase, settings: NativeDatabaseSettings, reader: bool) -> float: + if settings.password_provider is None: + return 600.0 + try: + password: Final = await asyncio.to_thread(settings.password_provider) + if reader: + native.set_passwords(reader_password=password) else: - await asyncio.sleep(_refresh_delay(settings, password)) + native.set_passwords(writer_password=password) + except Exception: + verbose_proxy_logger.warning("Native database credential refresh failed") + return 30.0 + return _refresh_delay(settings, password) @asynccontextmanager diff --git a/tests/integration/_support/process.py b/tests/integration/_support/process.py index 262114ee32b..327264ba2fa 100644 --- a/tests/integration/_support/process.py +++ b/tests/integration/_support/process.py @@ -1,18 +1,18 @@ import json -import threading import os import re import signal import subprocess import sys +import threading import time import uuid from collections.abc import Callable, Generator, Iterable, Iterator, Mapping from concurrent.futures import ThreadPoolExecutor -from contextlib import AbstractContextManager, contextmanager, suppress +from contextlib import AbstractContextManager, ExitStack, contextmanager, suppress from dataclasses import dataclass -from enum import Enum from datetime import UTC, datetime, timedelta +from enum import Enum from ipaddress import IPv4Address from itertools import accumulate from pathlib import Path @@ -417,7 +417,7 @@ def kill_stray_process(pid: int) -> None: @contextmanager -def paused(process: subprocess.Popen[bytes] | subprocess.Popen[str]) -> Iterator[None]: +def paused(process: subprocess.Popen[bytes] | subprocess.Popen[str] | psutil.Process) -> Iterator[None]: target: Final = psutil.Process(process.pid) _suspend(target) try: @@ -427,6 +427,20 @@ def paused(process: subprocess.Popen[bytes] | subprocess.Popen[str]) -> Iterator target.resume() +@contextmanager +def paused_prisma_engines(worker_pid: int) -> Iterator[None]: + engines: Final = tuple( + child + for child in psutil.Process(worker_pid).children(recursive=True) + if "query-engine" in child.name() + ) + assert engines, f"No Prisma query engine under worker {worker_pid}" + with ExitStack() as resources: + for engine in engines: + resources.enter_context(paused(engine)) + yield + + @contextmanager def owned_proxy( gateway: Gateway, diff --git a/tests/integration/authorization/test_native_auth_io.py b/tests/integration/authorization/test_native_auth_io.py new file mode 100644 index 00000000000..ac69fbcb14c --- /dev/null +++ b/tests/integration/authorization/test_native_auth_io.py @@ -0,0 +1,149 @@ +import os +import uuid +from hashlib import sha256 +from typing import Final + +import httpx +import psycopg +import pytest +from pydantic import JsonValue, TypeAdapter +from redis import Redis + +from tests.integration._support.client import Gateway, eventually +from tests.integration._support.process import paused_prisma_engines + +_OBSERVATIONS: Final = TypeAdapter(dict[str, list[dict[str, JsonValue]]]) + + +def _worker_pid(gateway: Gateway) -> int: + response: Final = gateway.request("GET", "/health/integration-readiness") + assert response.status_code == 200, response.text + return int(response.headers["x-integration-worker-pid"]) + + +def _chat(gateway: Gateway, model: str, key: str) -> httpx.Response: + return gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": f"native auth storage {uuid.uuid4().hex}"}]}, + key=key, + ) + + +def _observations(gateway: Gateway) -> list[dict[str, JsonValue]]: + response: Final = httpx.get(f"{gateway.upstream_url}/__observations", timeout=5, trust_env=False) + response.raise_for_status() + return _OBSERVATIONS.validate_json(response.content)["requests"] + + +def test_cold_and_warm_context_reads_do_not_need_prisma(gateway: Gateway, peer: Gateway) -> None: + primary_pid: Final = _worker_pid(gateway) + peer_pid: Final = _worker_pid(peer) + assert primary_pid != peer_pid + with gateway.scenario() as scenario: + model: Final = scenario.model() + team: Final = scenario.team(models=[model]) + owner: Final = scenario.member(team) + project: Final = scenario.project(team, models=[model]) + key: Final = scenario.key(team_id=team, user_id=owner, project_id=project, models=[model]) + with paused_prisma_engines(peer_pid): + cold: Final = _chat(peer, model, key) + assert cold.status_code == 200, cold.text + warm: Final = _chat(peer, model, key) + assert warm.status_code == 200, warm.text + assert len(_observations(gateway)) == 2 + + +def test_warm_key_read_does_not_wait_for_the_key_table(gateway: Gateway, peer: Gateway) -> None: + database_url: Final = gateway.database_url or os.environ["DATABASE_URL"] + with gateway.scenario() as scenario: + model: Final = scenario.model() + key: Final = scenario.key(models=[model]) + assert _chat(peer, model, key).status_code == 200 + with psycopg.connect(database_url) as connection: + connection.execute('LOCK TABLE "LiteLLM_VerificationToken" IN ACCESS EXCLUSIVE MODE') + warm: Final = _chat(peer, model, key) + assert warm.status_code == 200, warm.text + assert len(_observations(gateway)) == 2 + + +def test_revocation_reaches_a_warmed_peer_without_prisma(gateway: Gateway, peer: Gateway) -> None: + peer_pid: Final = _worker_pid(peer) + assert _worker_pid(gateway) != peer_pid + with gateway.scenario() as scenario: + model: Final = scenario.model() + key: Final = scenario.key(models=[model]) + for worker in (gateway, peer): + response: Final = _chat(worker, model, key) + assert response.status_code == 200, response.text + with paused_prisma_engines(peer_pid): + gateway.post("/key/block", {"key": key}) + denied: Final = eventually( + lambda: _chat(peer, model, key), lambda response: response.status_code == 401, seconds=10 + ) + before: Final = _observations(gateway) + assert len(before) >= 2 + assert _chat(peer, model, key).status_code == 401 + assert _observations(gateway) == before + assert denied.json()["error"]["type"] == "auth_error", denied.text + rows: Final = gateway.rows( + 'SELECT blocked FROM "LiteLLM_VerificationToken" WHERE token=%s', (sha256(key.encode()).hexdigest(),) + ) + assert rows == [{"blocked": True}] + + +@pytest.mark.parametrize("context", ("team", "project", "user")) +def test_context_changes_reach_a_warmed_peer(gateway: Gateway, peer: Gateway, context: str) -> None: + peer_pid: Final = _worker_pid(peer) + with gateway.scenario() as scenario: + model: Final = scenario.model() + team: Final = scenario.team(models=[model]) + owner: Final = scenario.member(team) + project: Final = scenario.project(team, models=[model]) + key: Final = scenario.key(team_id=team, user_id=owner, project_id=project, models=[model]) + warm: Final = _chat(peer, model, key) + assert warm.status_code == 200, warm.text + with paused_prisma_engines(peer_pid): + if context == "team": + gateway.post("/team/update", {"team_id": team, "blocked": True}) + elif context == "project": + gateway.post("/project/update", {"project_id": project, "blocked": True}) + else: + gateway.post("/user/update", {"user_id": owner, "metadata": {"scim_active": False}}) + denied: Final = eventually( + lambda: _chat(peer, model, key), lambda response: response.status_code in (401, 422), seconds=10 + ) + before: Final = _observations(gateway) + assert before + repeated: Final = _chat(peer, model, key) + assert repeated.status_code == denied.status_code, repeated.text + assert _observations(gateway) == before + + +@pytest.mark.parametrize("source", ("redis", "database")) +def test_spend_reads_share_native_counters_and_reseed_without_prisma( + gateway: Gateway, peer: Gateway, source: str +) -> None: + database_url: Final = gateway.database_url or os.environ["DATABASE_URL"] + peer_pid: Final = _worker_pid(peer) + with gateway.scenario() as scenario: + model: Final = scenario.model() + control: Final = scenario.key(models=[model]) + allowed: Final = _chat(peer, model, control) + assert allowed.status_code == 200, allowed.text + key: Final = scenario.key(models=[model], max_budget=1) + hashed: Final = sha256(key.encode()).hexdigest() + counter: Final = f"spend:key:{hashed}" + with Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"]), socket_timeout=2) as cache: + if source == "redis": + cache.set(counter, 2, ex=60) + else: + with psycopg.connect(database_url) as connection: + connection.execute('UPDATE "LiteLLM_VerificationToken" SET spend=2 WHERE token=%s', (hashed,)) + cache.delete(counter) + with paused_prisma_engines(peer_pid): + denied: Final = _chat(peer, model, key) + assert denied.status_code == 422, denied.text + value: Final = cache.get(counter) + assert value is not None and float(value) == 2 + assert len(_observations(gateway)) == 1 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 1a225b2571c..e2696f89c2a 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 @@ -20,6 +20,7 @@ import pytest from fastapi import Request, status import litellm +from litellm.caching.dual_cache import DualCache from litellm.proxy._types import ( CallInfo, Litellm_EntityType, @@ -881,7 +882,7 @@ class TestAuthCacheRedisWritePolicy: """ @pytest.mark.asyncio - async def test_get_key_object_db_load_publishes_to_redis(self): + async def test_get_key_object_db_load_stays_local_and_is_reused_without_redis_publication(self): mock_prisma_client = MagicMock() mock_prisma_client.get_data = AsyncMock(return_value=UserAPIKeyAuth(token="hashed-token-db")) @@ -896,10 +897,14 @@ class TestAuthCacheRedisWritePolicy: ) assert key_obj.token == "hashed-token-db" - fake_redis.async_set_cache.assert_awaited_once() - assert ( - fake_redis.async_set_cache.await_args.kwargs.get("key") or fake_redis.async_set_cache.await_args.args[0] - ) == "hashed-token-db" + cached: Final = await get_key_object( + hashed_token="hashed-token-db", + prisma_client=mock_prisma_client, + user_api_key_cache=cache, + ) + assert cached.token == key_obj.token + mock_prisma_client.get_data.assert_awaited_once() + fake_redis.async_set_cache.assert_not_awaited() def test_get_cli_jwt_auth_token_default_expiration(valid_sso_user_defined_values): @@ -7409,7 +7414,7 @@ class TestManagementObjectTTLHonored: @pytest.mark.asyncio async def test_uses_configured_user_api_key_cache_ttl(self): mem = _TTLCapturingInMemoryCache() - cache = UserApiKeyCache(in_memory_cache=mem, default_in_memory_ttl=300) + cache = DualCache(in_memory_cache=mem, default_in_memory_ttl=300) await _cache_management_object( key="team_id:lit-3338", @@ -7424,7 +7429,7 @@ class TestManagementObjectTTLHonored: @pytest.mark.asyncio async def test_falls_back_to_management_default_when_unconfigured(self): mem = _TTLCapturingInMemoryCache() - cache = UserApiKeyCache(in_memory_cache=mem) + cache = DualCache(in_memory_cache=mem) assert cache.default_in_memory_ttl is None await _cache_management_object( diff --git a/tests/unit/proxy/auth/test_resolvers_store.py b/tests/unit/proxy/auth/test_resolvers_store.py index 607162000ab..bd0fd9a452d 100644 --- a/tests/unit/proxy/auth/test_resolvers_store.py +++ b/tests/unit/proxy/auth/test_resolvers_store.py @@ -1,6 +1,7 @@ from __future__ import annotations -from typing import Dict, Optional +from types import SimpleNamespace +from typing import Dict, Final, Optional import pytest @@ -14,6 +15,8 @@ from litellm.proxy.auth.resolvers.exceptions import ( from litellm.proxy.auth.auth_method import AuthMethod from litellm.proxy.auth.resolvers.models import Principal, PrincipalType from litellm.proxy.auth.resolvers.store import IdentityStore +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm._auth.failure import NativeAuthRuntimeError class _FakeCache: @@ -92,3 +95,39 @@ async def test_resolve_raises_a_generic_key_not_found_carrying_the_hash_when_no_ assert exc_info.value.message == INVALID_API_KEY_ERROR_MESSAGE assert missing not in str(exc_info.value) assert exc_info.value.hashed_token == missing + + +class _InvalidatedNativeRead: + def __init__(self, cache: UserApiKeyCache, *, always_invalidate: bool = False) -> None: + self.cache = cache + self.calls = 0 + self.always_invalidate = always_invalidate + + async def auth_read(self, kind: str, ids: list[str], writer: bool = False) -> list[dict[str, object]]: + self.calls += 1 + if self.calls == 1 or self.always_invalidate: + self.cache.delete_local(ids[0]) + return [{"token": ids[0], "blocked": self.calls > 1, "models": ["allowed"]}] + + +async def test_invalidated_key_fill_reloads_instead_of_returning_or_caching_the_old_identity(): + cache: Final = UserApiKeyCache() + handle: Final = _InvalidatedNativeRead(cache) + token: Final = hash_token("sk-fill-fence") + store: Final = IdentityStore(SimpleNamespace(native_auth_database=handle), cache) + principal: Final = await store.resolve(token) + assert IdentityStore.key_from_principal(principal).blocked is True + assert handle.calls == 2 + cached: Final = await cache.async_get_cache(token, model_type=UserAPIKeyAuth) + assert cached is not None and cached.blocked is True + + +async def test_continuously_invalidated_key_fills_stop_without_caching_a_stale_identity(): + cache: Final = UserApiKeyCache() + handle: Final = _InvalidatedNativeRead(cache, always_invalidate=True) + token: Final = hash_token("sk-contended-fill") + store: Final = IdentityStore(SimpleNamespace(native_auth_database=handle), cache) + with pytest.raises(NativeAuthRuntimeError, match="changed during lookup"): + await store.resolve(token) + assert handle.calls == 2 + assert cache.get_local(token) is None