feat(auth): fence native key fills and prove cross-worker storage

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,

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

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

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

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,

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

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(

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