From f624f3402fdd8386eda53e65beadd8a655e9d72f Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 17:53:15 +0000 Subject: [PATCH 01/17] refactor(redis): extract shared health store --- .../shared_health_check_manager.py | 131 +-- litellm/redis/stores/shared_health.py | 112 ++ tests/unit/proxy/test_shared_health_check.py | 997 +++++++----------- 3 files changed, 568 insertions(+), 672 deletions(-) create mode 100644 litellm/redis/stores/shared_health.py diff --git a/litellm/proxy/health_check_utils/shared_health_check_manager.py b/litellm/proxy/health_check_utils/shared_health_check_manager.py index d689ab2a537..567fc5823bf 100644 --- a/litellm/proxy/health_check_utils/shared_health_check_manager.py +++ b/litellm/proxy/health_check_utils/shared_health_check_manager.py @@ -1,12 +1,9 @@ import asyncio -import json import time from collections.abc import Mapping, Sequence +from datetime import timedelta from typing import TYPE_CHECKING, Final -from pydantic import ConfigDict, TypeAdapter, with_config -from typing_extensions import NotRequired, ReadOnly, TypedDict - from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm.caching.redis_cache import RedisCache @@ -14,25 +11,23 @@ from litellm.constants import ( DEFAULT_SHARED_HEALTH_CHECK_LOCK_TTL, DEFAULT_SHARED_HEALTH_CHECK_TTL, ) -from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy.health_check import perform_health_check +from litellm.redis.stores.shared_health import ( + CachedHealthCheckResults, + SharedHealthStore, + TransportSharedHealthStore, + lock_key, + model_lock_key, + model_results_key, + results_key, +) +from litellm.redis.transport.contract import RedisFailure from litellm.router_utils.health_state_cache import HEALTH_CHECKS_TARGET if TYPE_CHECKING: from litellm.router import Router -@with_config(ConfigDict(extra="allow", strict=True)) -class _CachedHealthCheckResults(TypedDict): - healthy_endpoints: NotRequired[ReadOnly[list[Mapping[str, object]]]] - unhealthy_endpoints: NotRequired[ReadOnly[list[Mapping[str, object]]]] - timestamp: NotRequired[ReadOnly[int | float]] - checked_by: NotRequired[ReadOnly[object]] - - -_CACHED_HEALTH_CHECK_RESULTS: Final = TypeAdapter(_CachedHealthCheckResults) - - class SharedHealthCheckManager: """ Manager for coordinating health checks across multiple pods using Redis. @@ -51,6 +46,9 @@ class SharedHealthCheckManager: lock_ttl: int = DEFAULT_SHARED_HEALTH_CHECK_LOCK_TTL, ): self.redis_cache = redis_cache + self._store: SharedHealthStore | None = ( + TransportSharedHealthStore(redis_cache.transport()) if redis_cache is not None else None + ) self.health_check_ttl = health_check_ttl self.lock_ttl = lock_ttl self.pod_id = f"pod_{int(time.time() * 1000)}" @@ -58,22 +56,22 @@ class SharedHealthCheckManager: @staticmethod def get_health_check_lock_key() -> str: """Get the Redis key for health check lock.""" - return "health_check_lock" + return lock_key() @staticmethod def get_health_check_cache_key() -> str: """Get the Redis key for health check results cache.""" - return "health_check_results" + return results_key() @staticmethod def get_model_health_check_lock_key(model_name: str) -> str: """Get the Redis key for model-specific health check lock.""" - return f"health_check_lock:{model_name}" + return model_lock_key(model_name) @staticmethod def get_model_health_check_cache_key(model_name: str) -> str: """Get the Redis key for model-specific health check results cache.""" - return f"health_check_results:{model_name}" + return model_results_key(model_name) @with_service_target(HEALTH_CHECKS_TARGET) async def acquire_health_check_lock(self) -> bool: @@ -83,18 +81,14 @@ class SharedHealthCheckManager: Returns: bool: True if lock was acquired, False otherwise """ - if self.redis_cache is None: + if self._store is None: verbose_proxy_logger.debug("redis_cache is None, skipping lock acquisition") return False try: - lock_key: Final = self.get_health_check_lock_key() - acquired: Final = await self.redis_cache.async_set_cache( - lock_key, - self.pod_id, - nx=True, # Only set if key doesn't exist - ttl=self.lock_ttl, - ) + acquired: Final = await self._store.acquire(self.pod_id, ttl=timedelta(seconds=int(self.lock_ttl))) + if isinstance(acquired, RedisFailure): + return False if acquired: verbose_proxy_logger.info("Pod %s acquired health check lock", self.pod_id) @@ -102,49 +96,38 @@ class SharedHealthCheckManager: verbose_proxy_logger.debug("Pod %s failed to acquire health check lock", self.pod_id) return bool(acquired) - except Exception as e: - verbose_proxy_logger.error("Error acquiring health check lock: %s", str(e)) + except Exception: return False @with_service_target(HEALTH_CHECKS_TARGET) async def release_health_check_lock(self) -> None: """Release the global health check lock.""" - if self.redis_cache is None: + if self._store is None: return try: - lock_key: Final = self.get_health_check_lock_key() - # Only release if we own the lock - current_owner: Final = await self.redis_cache.async_get_cache(lock_key) - if current_owner == self.pod_id: - await self.redis_cache.async_delete_cache(lock_key) + released: Final = await self._store.release(self.pod_id) + if released is True: verbose_proxy_logger.info("Pod %s released health check lock", self.pod_id) - except Exception as e: - verbose_proxy_logger.error("Error releasing health check lock: %s", str(e)) + except Exception: + return @with_service_target(HEALTH_CHECKS_TARGET) - async def get_cached_health_check_results(self) -> _CachedHealthCheckResults | None: + async def get_cached_health_check_results(self) -> CachedHealthCheckResults | None: """ Get cached health check results from Redis. Returns: Optional[Dict]: Cached health check results or None if not found/expired """ - if self.redis_cache is None: + if self._store is None: return None try: - cache_key: Final = self.get_health_check_cache_key() - cached_data: Final = await self.redis_cache.async_get_cache(cache_key) - - if cached_data is None: + cached_results: Final = await self._store.get() + if cached_results is None or isinstance(cached_results, RedisFailure): return None - # Parse the cached data - cached_results: Final = _CACHED_HEALTH_CHECK_RESULTS.validate_python( - json.loads(cached_data) if isinstance(cached_data, str) else cached_data - ) - # Check if the cache is still valid cache_timestamp: Final = cached_results.get("timestamp", 0) current_time: Final = time.time() @@ -156,8 +139,7 @@ class SharedHealthCheckManager: verbose_proxy_logger.debug("Using cached health check results") return cached_results - except Exception as e: - verbose_proxy_logger.error("Error getting cached health check results: %s", str(e)) + except Exception: return None @with_service_target(HEALTH_CHECKS_TARGET) @@ -173,7 +155,7 @@ class SharedHealthCheckManager: healthy_endpoints: List of healthy endpoints unhealthy_endpoints: List of unhealthy endpoints """ - if self.redis_cache is None: + if self._store is None: return try: @@ -186,12 +168,9 @@ class SharedHealthCheckManager: "checked_by": self.pod_id, } - cache_key: Final = self.get_health_check_cache_key() - await self.redis_cache.async_set_cache( - cache_key, - safe_dumps(cache_data), - ttl=self.health_check_ttl, - ) + written: Final = await self._store.put(cache_data, ttl=timedelta(seconds=int(self.health_check_ttl))) + if isinstance(written, RedisFailure): + return verbose_proxy_logger.info( "Cached health check results for %d healthy and %d unhealthy endpoints", @@ -199,8 +178,8 @@ class SharedHealthCheckManager: len(unhealthy_endpoints), ) - except Exception as e: - verbose_proxy_logger.error("Error caching health check results: %s", str(e)) + except Exception: + return @with_service_target(HEALTH_CHECKS_TARGET) async def perform_shared_health_check( @@ -274,7 +253,7 @@ class SharedHealthCheckManager: else: # If Redis is not configured, skip polling — there is no cache # to wait for. - if self.redis_cache is None: + if self._store is None: return await perform_health_check( model_list=model_list, details=details, @@ -312,9 +291,12 @@ class SharedHealthCheckManager: # Check if the lock is still held — if it was released without # caching (e.g. the holder crashed), stop waiting early. try: - lock_key = self.get_health_check_lock_key() - current_owner = await self.redis_cache.async_get_cache(lock_key) - if current_owner is None: + if ( + isinstance(current_owner := await self._store.owner(), RedisFailure) + and current_owner.kind == "breaker_open" + ): + continue + if current_owner is None or isinstance(current_owner, RedisFailure): verbose_proxy_logger.debug( "Pod %s detected lock released without cache, stopping wait", self.pod_id, @@ -347,15 +329,15 @@ class SharedHealthCheckManager: Returns: bool: True if health check is in progress, False otherwise """ - if self.redis_cache is None: + if self._store is None: return False try: - lock_key: Final = self.get_health_check_lock_key() - current_owner: Final = await self.redis_cache.async_get_cache(lock_key) - return current_owner is not None and current_owner != self.pod_id - except Exception as e: - verbose_proxy_logger.error("Error checking health check lock status: %s", str(e)) + current_owner: Final = await self._store.owner() + if isinstance(current_owner, RedisFailure): + return False + return current_owner is not None and current_owner.value != self.pod_id + except Exception: return False @with_service_target(HEALTH_CHECKS_TARGET) @@ -373,11 +355,12 @@ class SharedHealthCheckManager: "cache_ttl": self.health_check_ttl, } - if self.redis_cache is not None: + if self._store is not None: try: - # Check if there's a current lock - lock_key: Final = self.get_health_check_lock_key() - current_owner: Final = await self.redis_cache.async_get_cache(lock_key) + owner: Final = await self._store.owner() + if isinstance(owner, RedisFailure) and owner.kind == "breaker_open": + return {**status, "error": owner.message} + current_owner: Final = None if owner is None or isinstance(owner, RedisFailure) else owner.value status["lock_owner"] = current_owner status["lock_in_progress"] = current_owner is not None diff --git a/litellm/redis/stores/shared_health.py b/litellm/redis/stores/shared_health.py new file mode 100644 index 00000000000..0ce75fd937a --- /dev/null +++ b/litellm/redis/stores/shared_health.py @@ -0,0 +1,112 @@ +import ast +import json +from collections.abc import Mapping +from dataclasses import dataclass +from datetime import timedelta +from typing import Final, Protocol + +from pydantic import ConfigDict, TypeAdapter, with_config +from typing_extensions import NotRequired, ReadOnly, TypedDict + +from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.redis.transport.contract import RedisFailure, RedisTransport + + +@with_config(ConfigDict(extra="allow", strict=True)) +class CachedHealthCheckResults(TypedDict): + healthy_endpoints: NotRequired[ReadOnly[list[Mapping[str, object]]]] + unhealthy_endpoints: NotRequired[ReadOnly[list[Mapping[str, object]]]] + timestamp: NotRequired[ReadOnly[int | float]] + checked_by: NotRequired[ReadOnly[object]] + + +_RESULTS: Final = TypeAdapter(CachedHealthCheckResults) +_VALUE: Final = TypeAdapter(object) + + +@dataclass(frozen=True, slots=True) +class SharedHealthLockOwner: + value: object + + +def lock_key() -> str: + return "health_check_lock" + + +def results_key() -> str: + return "health_check_results" + + +def model_lock_key(model_name: str) -> str: + return f"{lock_key()}:{model_name}" + + +def model_results_key(model_name: str) -> str: + return f"{results_key()}:{model_name}" + + +def _decode_cache_value(raw: bytes) -> object: + decoded: Final = raw.decode("utf-8") + try: + return _VALUE.validate_python(json.loads(decoded)) + except ValueError: + return _VALUE.validate_python(ast.literal_eval(decoded)) + + +class SharedHealthStore(Protocol): + async def acquire(self, owner: str, *, ttl: timedelta) -> bool | RedisFailure: ... + + async def owner(self) -> SharedHealthLockOwner | None | RedisFailure: ... + + async def release(self, owner: str) -> bool | RedisFailure: ... + + async def put(self, snapshot: Mapping[str, object], *, ttl: timedelta) -> None | RedisFailure: ... + + async def get(self) -> CachedHealthCheckResults | None | RedisFailure: ... + + +@dataclass(frozen=True, slots=True) +class TransportSharedHealthStore: + transport: RedisTransport + + async def acquire(self, owner: str, *, ttl: timedelta) -> bool | RedisFailure: + return await self.transport.set_seconds(lock_key(), json.dumps(owner), ttl=ttl, nx=True) + + async def owner(self) -> SharedHealthLockOwner | None | RedisFailure: + raw: Final = await self.transport.get(lock_key()) + if raw is None or isinstance(raw, RedisFailure): + return raw + try: + value: Final = _decode_cache_value(raw) + return None if value is None else SharedHealthLockOwner(value) + except (ValueError, TypeError, SyntaxError, RecursionError) as exc: + return RedisFailure(kind="decode", message=f"shared health owner is invalid: {exc}") + + async def release(self, owner: str) -> bool | RedisFailure: + held: Final = await self.owner() + if isinstance(held, RedisFailure): + return held + if held is None or held.value != owner: + return False + deleted: Final = await self.transport.delete(lock_key()) + return deleted if isinstance(deleted, RedisFailure) else bool(deleted) + + async def put(self, snapshot: Mapping[str, object], *, ttl: timedelta) -> None | RedisFailure: + try: + encoded: Final = json.dumps(safe_dumps(snapshot)) + except (ValueError, TypeError, RecursionError) as exc: + return RedisFailure(kind="decode", message=f"shared health snapshot could not be encoded: {exc}") + written: Final = await self.transport.set_seconds(results_key(), encoded, ttl=ttl) + return written if isinstance(written, RedisFailure) else None + + async def get(self) -> CachedHealthCheckResults | None | RedisFailure: + raw: Final = await self.transport.get(results_key()) + if raw is None or isinstance(raw, RedisFailure): + return raw + try: + value: Final = _decode_cache_value(raw) + if value is None: + return None + return _RESULTS.validate_python(json.loads(value) if isinstance(value, str) else value) + except (ValueError, TypeError, SyntaxError, RecursionError) as exc: + return RedisFailure(kind="decode", message=f"shared health snapshot is invalid: {exc}") diff --git a/tests/unit/proxy/test_shared_health_check.py b/tests/unit/proxy/test_shared_health_check.py index 100425a8c9f..02e132659d2 100644 --- a/tests/unit/proxy/test_shared_health_check.py +++ b/tests/unit/proxy/test_shared_health_check.py @@ -1,605 +1,406 @@ import asyncio import json import time -from unittest.mock import AsyncMock, MagicMock, patch +from collections.abc import AsyncIterator, Mapping +from dataclasses import dataclass +from itertools import chain +from typing import Final +from unittest.mock import patch +import fakeredis import pytest +import pytest_asyncio +from prometheus_client import REGISTRY +from pydantic import TypeAdapter +from redis.exceptions import TimeoutError as RedisTimeoutError -from litellm.proxy.health_check_utils.shared_health_check_manager import ( - SharedHealthCheckManager, +import litellm +from litellm.caching.redis_cache import RedisCache +from litellm.constants import REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD +from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.proxy.health_check_utils.shared_health_check_manager import SharedHealthCheckManager +from tests._support.fake_redis import fake_redis_cache + +NOW: Final = 1_700_000_000.25 +LOCK: Final = "shared:health_check_lock" +RESULTS: Final = "shared:health_check_results" +_VALUE: Final = TypeAdapter(object) + + +@dataclass(frozen=True) +class HealthRedis: + server: fakeredis.FakeServer + raw: fakeredis.FakeAsyncRedis + cache: RedisCache + manager: SharedHealthCheckManager + + +@pytest_asyncio.fixture(loop_scope="function") +async def health(monkeypatch: pytest.MonkeyPatch) -> AsyncIterator[HealthRedis]: + monkeypatch.setattr(time, "time", lambda: NOW) + server: Final = fakeredis.FakeServer() + raw: Final = fakeredis.FakeAsyncRedis(server=server) + cache: Final = fake_redis_cache(server, namespace="shared") + manager: Final = SharedHealthCheckManager(cache, health_check_ttl=300, lock_ttl=60) + manager.pod_id = 'pod_雪"' + yield HealthRedis(server, raw, cache, manager) + + +def test_public_key_helpers_keep_global_and_model_keys() -> None: + assert SharedHealthCheckManager.get_health_check_lock_key() == "health_check_lock" + assert SharedHealthCheckManager.get_health_check_cache_key() == "health_check_results" + assert SharedHealthCheckManager.get_model_health_check_lock_key("a:雪") == "health_check_lock:a:雪" + assert SharedHealthCheckManager.get_model_health_check_cache_key("a:雪") == "health_check_results:a:雪" + + +@pytest.mark.asyncio +async def test_old_lock_writer_new_status_and_owner_checked_release(health: HealthRedis) -> None: + await health.cache.async_set_cache("health_check_lock", "other_雪", ttl=60, nx=True) + assert await health.manager.acquire_health_check_lock() is False + assert await health.manager.is_health_check_in_progress() is True + assert await health.manager.get_health_check_status() == { + "pod_id": health.manager.pod_id, + "redis_available": True, + "lock_ttl": 60, + "cache_ttl": 300, + "lock_owner": "other_雪", + "lock_in_progress": True, + "cache_available": False, + } + await health.manager.release_health_check_lock() + assert await health.raw.get(LOCK) == b'"other_\\u96ea"' + health.manager.pod_id = "other_雪" + assert await health.manager.is_health_check_in_progress() is False + await health.manager.release_health_check_lock() + assert await health.raw.get(LOCK) is None + + +@pytest.mark.asyncio +async def test_new_lock_writer_old_reader_exact_bytes_namespace_and_expiry(health: HealthRedis) -> None: + assert await health.manager.acquire_health_check_lock() is True + assert await health.raw.get(LOCK) == b'"pod_\\u96ea\\""' + assert await health.cache.async_get_cache("health_check_lock") == health.manager.pod_id + assert await health.raw.pttl(LOCK) == 60_000 + assert await health.raw.keys("*") == [LOCK.encode()] + assert await health.manager.acquire_health_check_lock() is False + await health.manager.release_health_check_lock() + assert await health.raw.get(LOCK) is None + assert await health.manager.is_health_check_in_progress() is False + + +@pytest.mark.asyncio +async def test_snapshot_new_writer_old_reader_exact_double_encoding_and_ttl(health: HealthRedis) -> None: + healthy: Final = [{"model": "雪", "latency": 1.25, "count": 2, "text": "a\x00b"}] + unhealthy: Final = [{"model": "bad", "exception_status": 503}] + expected: Final = { + "healthy_endpoints": [{"model": "雪", "latency": 1.25, "count": 2, "text": "ab"}], + "unhealthy_endpoints": unhealthy, + "healthy_count": 1, + "unhealthy_count": 1, + "timestamp": NOW, + "checked_by": health.manager.pod_id, + } + await health.manager.cache_health_check_results(healthy, unhealthy) + await health.cache.async_set_cache("old_writer", safe_dumps(expected), ttl=300) + assert await health.raw.get(RESULTS) == await health.raw.get("shared:old_writer") + assert await health.raw.get(RESULTS) == json.dumps(json.dumps(expected)).encode() + assert await health.raw.pttl(RESULTS) == 300_000 + assert await health.raw.get("shared:shared:health_check_results") is None + assert _VALUE.validate_json(await health.cache.async_get_cache("health_check_results")) == expected + assert await health.manager.get_cached_health_check_results() == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize("encoding", ["double", "single", "literal"]) +async def test_old_snapshot_formats_preserve_extra_fields_and_cached_flow(health: HealthRedis, encoding: str) -> None: + snapshot: Final = { + "healthy_endpoints": [{"model": "雪", "latency": 1.25}], + "unhealthy_endpoints": [], + "timestamp": NOW - 300, + "checked_by": {"legacy": "pod"}, + "extra": [0, False, None], + } + if encoding == "double": + await health.cache.async_set_cache("health_check_results", safe_dumps(snapshot), ttl=300) + elif encoding == "single": + await health.cache.async_set_cache("health_check_results", snapshot, ttl=300) + else: + await health.raw.set(RESULTS, repr(snapshot), ex=300) + assert await health.manager.get_cached_health_check_results() == snapshot + assert await health.manager.perform_shared_health_check([]) == (snapshot["healthy_endpoints"], [], {}) + assert await health.manager.get_health_check_status() == { + "pod_id": health.manager.pod_id, + "redis_available": True, + "lock_ttl": 60, + "cache_ttl": 300, + "lock_owner": None, + "lock_in_progress": False, + "cache_available": True, + "cache_age_seconds": 300, + "last_checked_by": {"legacy": "pod"}, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "snapshot", + [ + {}, + {"timestamp": NOW - 301}, + {"timestamp": str(NOW)}, + {"timestamp": True}, + {"timestamp": None}, + {"timestamp": NOW, "healthy_endpoints": None}, + {"timestamp": NOW, "unhealthy_endpoints": [1]}, + ], ) - - -class TestSharedHealthCheckManager: - """Test cases for SharedHealthCheckManager""" - - @pytest.fixture - def mock_redis_cache(self): - """Mock Redis cache for testing""" - cache = AsyncMock() - cache.async_set_cache = AsyncMock() - cache.async_get_cache = AsyncMock() - cache.async_delete_cache = AsyncMock() - return cache - - @pytest.fixture - def shared_health_manager(self, mock_redis_cache): - """Create SharedHealthCheckManager instance for testing""" - return SharedHealthCheckManager( - redis_cache=mock_redis_cache, - health_check_ttl=300, - lock_ttl=60, - ) - - def test_initialization(self, mock_redis_cache): - """Test SharedHealthCheckManager initialization""" - manager = SharedHealthCheckManager( - redis_cache=mock_redis_cache, - health_check_ttl=300, - lock_ttl=60, - ) - - assert manager.redis_cache == mock_redis_cache - assert manager.health_check_ttl == 300 - assert manager.lock_ttl == 60 - assert manager.pod_id.startswith("pod_") - - def test_initialization_without_redis(self): - """Test SharedHealthCheckManager initialization without Redis""" - manager = SharedHealthCheckManager(redis_cache=None) - - assert manager.redis_cache is None - assert manager.health_check_ttl == 300 # Default value - assert manager.lock_ttl == 60 # Default value - - def test_get_health_check_lock_key(self): - """Test getting health check lock key""" - key = SharedHealthCheckManager.get_health_check_lock_key() - assert key == "health_check_lock" - - def test_get_health_check_cache_key(self): - """Test getting health check cache key""" - key = SharedHealthCheckManager.get_health_check_cache_key() - assert key == "health_check_results" - - def test_get_model_health_check_lock_key(self): - """Test getting model-specific health check lock key""" - key = SharedHealthCheckManager.get_model_health_check_lock_key("test-model") - assert key == "health_check_lock:test-model" - - def test_get_model_health_check_cache_key(self): - """Test getting model-specific health check cache key""" - key = SharedHealthCheckManager.get_model_health_check_cache_key("test-model") - assert key == "health_check_results:test-model" - - @pytest.mark.asyncio - async def test_acquire_health_check_lock_success( - self, shared_health_manager, mock_redis_cache - ): - """Test successful lock acquisition""" - mock_redis_cache.async_set_cache.return_value = True - - result = await shared_health_manager.acquire_health_check_lock() - - assert result is True - mock_redis_cache.async_set_cache.assert_called_once_with( - "health_check_lock", - shared_health_manager.pod_id, - nx=True, - ttl=60, - ) - - @pytest.mark.asyncio - async def test_acquire_health_check_lock_failure( - self, shared_health_manager, mock_redis_cache - ): - """Test failed lock acquisition""" - mock_redis_cache.async_set_cache.return_value = False - - result = await shared_health_manager.acquire_health_check_lock() - - assert result is False - - @pytest.mark.asyncio - async def test_acquire_health_check_lock_no_redis(self): - """Test lock acquisition without Redis""" - manager = SharedHealthCheckManager(redis_cache=None) - - result = await manager.acquire_health_check_lock() - - assert result is False - - @pytest.mark.asyncio - async def test_acquire_health_check_lock_exception( - self, shared_health_manager, mock_redis_cache - ): - """Test lock acquisition with exception""" - mock_redis_cache.async_set_cache.side_effect = Exception("Redis error") - - result = await shared_health_manager.acquire_health_check_lock() - - assert result is False - - @pytest.mark.asyncio - async def test_release_health_check_lock_success( - self, shared_health_manager, mock_redis_cache - ): - """Test successful lock release""" - mock_redis_cache.async_get_cache.return_value = shared_health_manager.pod_id - - await shared_health_manager.release_health_check_lock() - - mock_redis_cache.async_get_cache.assert_called_once_with("health_check_lock") - mock_redis_cache.async_delete_cache.assert_called_once_with("health_check_lock") - - @pytest.mark.asyncio - async def test_release_health_check_lock_wrong_owner( - self, shared_health_manager, mock_redis_cache - ): - """Test lock release when not the owner""" - mock_redis_cache.async_get_cache.return_value = "other_pod_id" - - await shared_health_manager.release_health_check_lock() - - mock_redis_cache.async_get_cache.assert_called_once_with("health_check_lock") - mock_redis_cache.async_delete_cache.assert_not_called() - - @pytest.mark.asyncio - async def test_release_health_check_lock_no_redis(self): - """Test lock release without Redis""" - manager = SharedHealthCheckManager(redis_cache=None) - - # Should not raise exception - await manager.release_health_check_lock() - - @pytest.mark.asyncio - async def test_get_cached_health_check_results_success( - self, shared_health_manager, mock_redis_cache - ): - """Test getting cached health check results successfully""" - current_time = time.time() - cached_data = { - "healthy_endpoints": [{"model": "test-model"}], - "unhealthy_endpoints": [], - "healthy_count": 1, - "unhealthy_count": 0, - "timestamp": current_time - 100, # 100 seconds ago - "checked_by": "test_pod", - } - mock_redis_cache.async_get_cache.return_value = json.dumps(cached_data) - - result = await shared_health_manager.get_cached_health_check_results() - - assert result is not None - assert result["healthy_count"] == 1 - assert result["unhealthy_count"] == 0 - - @pytest.mark.asyncio - async def test_get_cached_health_check_results_expired( - self, shared_health_manager, mock_redis_cache - ): - """Test getting expired cached health check results""" - current_time = time.time() - cached_data = { - "healthy_endpoints": [{"model": "test-model"}], - "unhealthy_endpoints": [], - "healthy_count": 1, - "unhealthy_count": 0, - "timestamp": current_time - 400, # 400 seconds ago (expired) - "checked_by": "test_pod", - } - mock_redis_cache.async_get_cache.return_value = json.dumps(cached_data) - - result = await shared_health_manager.get_cached_health_check_results() - - assert result is None - - @pytest.mark.asyncio - async def test_get_cached_health_check_results_no_cache( - self, shared_health_manager, mock_redis_cache - ): - """Test getting cached results when no cache exists""" - mock_redis_cache.async_get_cache.return_value = None - - result = await shared_health_manager.get_cached_health_check_results() - - assert result is None - - @pytest.mark.asyncio - async def test_get_cached_health_check_results_no_redis(self): - """Test getting cached results without Redis""" - manager = SharedHealthCheckManager(redis_cache=None) - - result = await manager.get_cached_health_check_results() - - assert result is None - - @pytest.mark.asyncio - async def test_cache_health_check_results_success( - self, shared_health_manager, mock_redis_cache - ): - """Test caching health check results successfully""" - healthy_endpoints = [{"model": "test-model-1"}] - unhealthy_endpoints = [{"model": "test-model-2"}] - - await shared_health_manager.cache_health_check_results( - healthy_endpoints, unhealthy_endpoints - ) - - mock_redis_cache.async_set_cache.assert_called_once() - call_args = mock_redis_cache.async_set_cache.call_args - assert call_args[0][0] == "health_check_results" # key - assert call_args[1]["ttl"] == 300 # ttl - - # Verify cached data structure - cached_data = json.loads(call_args[0][1]) - assert cached_data["healthy_endpoints"] == healthy_endpoints - assert cached_data["unhealthy_endpoints"] == unhealthy_endpoints - assert cached_data["healthy_count"] == 1 - assert cached_data["unhealthy_count"] == 1 - assert "timestamp" in cached_data - assert cached_data["checked_by"] == shared_health_manager.pod_id - - @pytest.mark.asyncio - async def test_cache_health_check_results_no_redis(self): - """Test caching results without Redis""" - manager = SharedHealthCheckManager(redis_cache=None) - - # Should not raise exception - await manager.cache_health_check_results([], []) - - @pytest.mark.asyncio - async def test_perform_shared_health_check_with_cache( - self, shared_health_manager, mock_redis_cache - ): - """Test performing shared health check when cache is available""" - # Mock cached results - cached_data = { - "healthy_endpoints": [{"model": "cached-model"}], - "unhealthy_endpoints": [], - "healthy_count": 1, - "unhealthy_count": 0, - "timestamp": time.time() - 100, - } - mock_redis_cache.async_get_cache.return_value = json.dumps(cached_data) - - model_list = [ - {"model_name": "test-model", "litellm_params": {"model": "test-model"}} - ] - - with patch( - "litellm.proxy.health_check_utils.shared_health_check_manager.perform_health_check" - ) as mock_perform: - healthy, unhealthy, _ = ( - await shared_health_manager.perform_shared_health_check( - model_list, details=True - ) - ) - - # Should return cached results, not call perform_health_check - assert healthy == [{"model": "cached-model"}] - assert unhealthy == [] - mock_perform.assert_not_called() - - @pytest.mark.asyncio - async def test_perform_shared_health_check_with_lock_acquisition( - self, shared_health_manager, mock_redis_cache - ): - """Test performing shared health check when acquiring lock""" - # No cached results - mock_redis_cache.async_get_cache.return_value = None - # Lock acquisition succeeds - mock_redis_cache.async_set_cache.return_value = True - - model_list = [ - {"model_name": "test-model", "litellm_params": {"model": "test-model"}} - ] - expected_healthy = [{"model": "test-model", "status": "healthy"}] - expected_unhealthy = [] - - with patch( - "litellm.proxy.health_check_utils.shared_health_check_manager.perform_health_check" - ) as mock_perform: - mock_perform.return_value = (expected_healthy, expected_unhealthy, {}) - - healthy, unhealthy, _ = ( - await shared_health_manager.perform_shared_health_check( - model_list, details=True - ) - ) - - # Should call perform_health_check and cache results - mock_perform.assert_called_once_with( - model_list=model_list, - details=True, - max_concurrency=None, - health_check_skip_disabled_background_models=False, - router=None, - ) - assert healthy == expected_healthy - assert unhealthy == expected_unhealthy - - # Should cache the results - assert mock_redis_cache.async_set_cache.call_count >= 2 # Lock + cache - - @pytest.mark.asyncio - async def test_perform_shared_health_check_lock_failed_then_cache( - self, shared_health_manager, mock_redis_cache - ): - """Test performing shared health check when lock fails but cache becomes available during polling""" - # First call: no cache, lock fails - # Polling finds cache on first iteration - mock_redis_cache.async_get_cache.side_effect = [ - None, # No cache initially (get_cached_health_check_results) - json.dumps( - { # Cache available on first poll iteration - "healthy_endpoints": [{"model": "cached-model"}], - "unhealthy_endpoints": [], - "healthy_count": 1, - "unhealthy_count": 0, - "timestamp": time.time() - 100, - } - ), - ] - mock_redis_cache.async_set_cache.return_value = False # Lock acquisition fails - - model_list = [ - {"model_name": "test-model", "litellm_params": {"model": "test-model"}} - ] - - with patch("asyncio.sleep") as mock_sleep: # Mock sleep to avoid actual delay - healthy, unhealthy, _ = ( - await shared_health_manager.perform_shared_health_check( - model_list, details=True - ) - ) - - # Should poll once (5s interval) and find cached results - mock_sleep.assert_called_once_with(5) - assert healthy == [{"model": "cached-model"}] - assert unhealthy == [] - - @pytest.mark.asyncio - async def test_perform_shared_health_check_fallback(self, mock_redis_cache): - """Test performing shared health check with fallback to local health check""" - # Use short lock_ttl so the polling loop only runs 2 iterations - manager = SharedHealthCheckManager( - redis_cache=mock_redis_cache, - health_check_ttl=300, - lock_ttl=10, - ) - - # No cache ever, lock always held by another pod - mock_redis_cache.async_get_cache.side_effect = [ - None, # Initial cache check - None, # Iteration 1: cache check - "other_pod", # Iteration 1: lock check (still held) - None, # Iteration 2: cache check - "other_pod", # Iteration 2: lock check (still held) - ] - mock_redis_cache.async_set_cache.return_value = False # Lock acquisition fails - - model_list = [ - {"model_name": "test-model", "litellm_params": {"model": "test-model"}} - ] - expected_healthy = [{"model": "test-model", "status": "healthy"}] - expected_unhealthy = [] - - with ( - patch("asyncio.sleep") as mock_sleep, - patch( - "litellm.proxy.health_check_utils.shared_health_check_manager.perform_health_check" - ) as mock_perform, - ): - mock_perform.return_value = (expected_healthy, expected_unhealthy, {}) - - healthy, unhealthy, _ = await manager.perform_shared_health_check( - model_list, details=True - ) - - # Should poll twice (5s * 2 = 10s >= lock_ttl) then fall back - assert mock_sleep.call_count == 2 - mock_sleep.assert_called_with(5) - mock_perform.assert_called_once_with( - model_list=model_list, - details=True, - max_concurrency=None, - health_check_skip_disabled_background_models=False, - router=None, - ) - assert healthy == expected_healthy - assert unhealthy == expected_unhealthy - - @pytest.mark.asyncio - async def test_perform_shared_health_check_early_exit_orphaned_lock( - self, shared_health_manager, mock_redis_cache - ): - """Test that polling exits early when the lock disappears without a cache write (crash recovery)""" - mock_redis_cache.async_get_cache.side_effect = [ - None, # Initial cache check - None, # Iteration 1: cache check (still no cache) - None, # Iteration 1: lock check -> lock gone (holder crashed) - ] - mock_redis_cache.async_set_cache.return_value = False # Lock acquisition fails - - model_list = [ - {"model_name": "test-model", "litellm_params": {"model": "test-model"}} - ] - expected_healthy = [{"model": "test-model", "status": "healthy"}] - expected_unhealthy = [] - - with ( - patch("asyncio.sleep") as mock_sleep, - patch( - "litellm.proxy.health_check_utils.shared_health_check_manager.perform_health_check" - ) as mock_perform, - ): - mock_perform.return_value = (expected_healthy, expected_unhealthy, {}) - - healthy, unhealthy, _ = ( - await shared_health_manager.perform_shared_health_check( - model_list, details=True - ) - ) - - # Should detect orphaned lock after 1 iteration and fall back immediately - mock_sleep.assert_called_once_with(5) - mock_perform.assert_called_once_with( - model_list=model_list, - details=True, - max_concurrency=None, - health_check_skip_disabled_background_models=False, - router=None, - ) - assert healthy == expected_healthy - assert unhealthy == expected_unhealthy - - @pytest.mark.asyncio - async def test_perform_shared_health_check_redis_error_during_polling( - self, shared_health_manager, mock_redis_cache - ): - """Test that a transient Redis error during lock polling doesn't crash the loop""" - cached_data = json.dumps( - { - "healthy_endpoints": [{"model": "cached-model"}], - "unhealthy_endpoints": [], - "healthy_count": 1, - "unhealthy_count": 0, - "timestamp": time.time() - 100, - } - ) - mock_redis_cache.async_get_cache.side_effect = [ - None, # Initial cache check - None, # Iteration 1: cache check - Exception("Redis connection lost"), # Iteration 1: lock check errors - cached_data, # Iteration 2: cache check -> found! - ] - mock_redis_cache.async_set_cache.return_value = False # Lock acquisition fails - - model_list = [ - {"model_name": "test-model", "litellm_params": {"model": "test-model"}} - ] - - with patch("asyncio.sleep") as mock_sleep: - healthy, unhealthy, _ = ( - await shared_health_manager.perform_shared_health_check( - model_list, details=True - ) - ) - - # Should survive the Redis error and find cache on iteration 2 - assert mock_sleep.call_count == 2 - assert healthy == [{"model": "cached-model"}] - assert unhealthy == [] - - @pytest.mark.asyncio - async def test_perform_shared_health_check_no_redis_skips_polling(self): - """Test that polling is skipped entirely when redis_cache is None""" - manager = SharedHealthCheckManager(redis_cache=None) - - model_list = [ - {"model_name": "test-model", "litellm_params": {"model": "test-model"}} - ] - expected_healthy = [{"model": "test-model", "status": "healthy"}] - expected_unhealthy = [] - - with ( - patch("asyncio.sleep") as mock_sleep, - patch( - "litellm.proxy.health_check_utils.shared_health_check_manager.perform_health_check" - ) as mock_perform, - ): - mock_perform.return_value = (expected_healthy, expected_unhealthy, {}) - - healthy, unhealthy, _ = await manager.perform_shared_health_check( - model_list, details=True - ) - - # Should NOT sleep at all — falls back to local health check immediately - mock_sleep.assert_not_called() - mock_perform.assert_called_once_with( - model_list=model_list, - details=True, - max_concurrency=None, - health_check_skip_disabled_background_models=False, - router=None, - ) - assert healthy == expected_healthy - assert unhealthy == expected_unhealthy - - @pytest.mark.asyncio - async def test_is_health_check_in_progress_true( - self, shared_health_manager, mock_redis_cache - ): - """Test checking if health check is in progress when it is""" - mock_redis_cache.async_get_cache.return_value = "other_pod_id" - - result = await shared_health_manager.is_health_check_in_progress() - - assert result is True - - @pytest.mark.asyncio - async def test_is_health_check_in_progress_false( - self, shared_health_manager, mock_redis_cache - ): - """Test checking if health check is in progress when it's not""" - mock_redis_cache.async_get_cache.return_value = None - - result = await shared_health_manager.is_health_check_in_progress() - - assert result is False - - @pytest.mark.asyncio - async def test_is_health_check_in_progress_own_lock( - self, shared_health_manager, mock_redis_cache - ): - """Test checking if health check is in progress when we own the lock""" - mock_redis_cache.async_get_cache.return_value = shared_health_manager.pod_id - - result = await shared_health_manager.is_health_check_in_progress() - - assert result is False - - @pytest.mark.asyncio - async def test_get_health_check_status( - self, shared_health_manager, mock_redis_cache - ): - """Test getting health check status""" - current_time = time.time() - cached_data = { - "healthy_endpoints": [{"model": "test-model"}], - "unhealthy_endpoints": [], - "healthy_count": 1, - "unhealthy_count": 0, - "timestamp": current_time - 100, - "checked_by": "test_pod", - } - - mock_redis_cache.async_get_cache.side_effect = [ - "other_pod_id", # Lock owner - json.dumps(cached_data), # Cached results - ] - - status = await shared_health_manager.get_health_check_status() - - assert status["pod_id"] == shared_health_manager.pod_id - assert status["redis_available"] is True - assert status["lock_ttl"] == 60 - assert status["cache_ttl"] == 300 - assert status["lock_owner"] == "other_pod_id" - assert status["lock_in_progress"] is True - assert status["cache_available"] is True - assert status["last_checked_by"] == "test_pod" - assert "cache_age_seconds" in status - - @pytest.mark.asyncio - async def test_get_health_check_status_no_redis(self): - """Test getting health check status without Redis""" - manager = SharedHealthCheckManager(redis_cache=None) - - status = await manager.get_health_check_status() - - assert status["pod_id"] == manager.pod_id - assert status["redis_available"] is False - assert status["lock_ttl"] == 60 - assert status["cache_ttl"] == 300 +async def test_stale_or_schema_invalid_snapshot_is_a_miss(health: HealthRedis, snapshot: Mapping[str, object]) -> None: + await health.cache.async_set_cache("health_check_results", safe_dumps(snapshot), ttl=300) + assert await health.manager.get_cached_health_check_results() is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("raw", [b"broken", b'"broken"', b"[]", b"null", b"\xff", b'"null"']) +async def test_malformed_snapshot_and_null_are_misses(health: HealthRedis, raw: bytes) -> None: + await health.raw.set(RESULTS, raw) + assert await health.manager.get_cached_health_check_results() is None + + +@pytest.mark.asyncio +async def test_missing_endpoints_and_future_timestamp_remain_valid(health: HealthRedis) -> None: + snapshot: Final = {"timestamp": NOW + 1} + await health.cache.async_set_cache("health_check_results", safe_dumps(snapshot), ttl=300) + assert await health.manager.get_cached_health_check_results() == snapshot + assert await health.manager.perform_shared_health_check([]) == ([], [], {}) + + +@pytest.mark.asyncio +async def test_lock_holder_executes_health_check_caches_and_releases(health: HealthRedis) -> None: + assert await health.manager.perform_shared_health_check([]) == ([], [], {}) + assert await health.manager.get_cached_health_check_results() == { + "healthy_endpoints": [], + "unhealthy_endpoints": [], + "healthy_count": 0, + "unhealthy_count": 0, + "timestamp": NOW, + "checked_by": health.manager.pod_id, + } + assert await health.raw.get(LOCK) is None + assert await health.raw.pttl(RESULTS) == 300_000 + + +@pytest.mark.asyncio +async def test_waiter_reads_snapshot_published_during_poll(health: HealthRedis) -> None: + await health.cache.async_set_cache("health_check_lock", "other", ttl=60) + + async def publish(delay: float) -> None: + assert delay == 5 + await health.manager.cache_health_check_results([{"model": "published"}], []) + + with patch("asyncio.sleep", new=publish): + result: Final = await health.manager.perform_shared_health_check([]) + assert result == ([{"model": "published"}], [], {}) + assert await health.raw.get(LOCK) == b'"other"' + + +@pytest.mark.asyncio +@pytest.mark.parametrize("change", ["orphaned", "outage", "held"]) +async def test_wait_policy_falls_back_locally_without_writing_snapshot(health: HealthRedis, change: str) -> None: + health.manager.lock_ttl = 10 + await health.cache.async_set_cache("health_check_lock", "other", ttl=60) + + async def during_wait(delay: float) -> None: + assert delay == 5 + if change == "orphaned": + await health.raw.delete(LOCK) + elif change == "outage": + health.server.connected = False + + with patch("asyncio.sleep", wraps=during_wait) as sleep: + assert await health.manager.perform_shared_health_check([]) == ([], [], {}) + assert sleep.await_count == (2 if change == "held" else 1) + health.server.connected = True + assert await health.raw.get(RESULTS) is None + + +@pytest.mark.asyncio +async def test_outage_is_false_or_miss_then_breaker_refuses_after_recovery(health: HealthRedis) -> None: + health.server.connected = False + assert await health.manager.acquire_health_check_lock() is False + assert await health.manager.get_cached_health_check_results() is None + assert await health.manager.is_health_check_in_progress() is False + assert await health.manager.release_health_check_lock() is None + assert await health.manager.cache_health_check_results([], []) is None + health.server.connected = True + assert await health.manager.acquire_health_check_lock() is False + assert await health.raw.get(LOCK) is None + status: Final = await health.manager.get_health_check_status() + assert "error" in status + assert "lock_owner" not in status + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ttl", [0, -1]) +async def test_nonpositive_ttl_does_not_write_a_short_lived_lock_or_snapshot(health: HealthRedis, ttl: int) -> None: + health.manager.lock_ttl = ttl + health.manager.health_check_ttl = ttl + assert await health.manager.acquire_health_check_lock() is False + await health.manager.cache_health_check_results([], []) + assert await health.raw.get(LOCK) is None + assert await health.raw.get(RESULTS) is None + + +@pytest.mark.asyncio +async def test_lost_write_replies_remain_unknown_applied(health: HealthRedis) -> None: + async def applied_set( + name: str, value: str, *, px: int | None = None, ex: int | None = None, nx: bool = False + ) -> None: + await health.raw.set(name, value, px=px, ex=ex, nx=nx) + raise RedisTimeoutError("reply lost after set") + + with patch.object(health.cache.init_async_client(), "set", new=applied_set): + assert await health.manager.acquire_health_check_lock() is False + await health.manager.cache_health_check_results([], []) + assert await health.raw.get(LOCK) == json.dumps(health.manager.pod_id).encode() + assert await health.raw.pttl(LOCK) == 60_000 + assert await health.manager.get_cached_health_check_results() == { + "healthy_endpoints": [], + "unhealthy_endpoints": [], + "healthy_count": 0, + "unhealthy_count": 0, + "timestamp": NOW, + "checked_by": health.manager.pod_id, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("applied", [False, True]) +async def test_lost_delete_reply_is_best_effort(health: HealthRedis, applied: bool) -> None: + assert await health.manager.acquire_health_check_lock() is True + + async def delete_timeout(name: str) -> None: + if applied: + await health.raw.delete(name) + raise RedisTimeoutError("reply lost during delete") + + with patch.object(health.cache.init_async_client(), "delete", new=delete_timeout): + assert await health.manager.release_health_check_lock() is None + assert await health.raw.get(LOCK) == (None if applied else json.dumps(health.manager.pod_id).encode()) + + +@pytest.mark.asyncio +async def test_release_keeps_existing_get_then_delete_race(health: HealthRedis) -> None: + assert await health.manager.acquire_health_check_lock() is True + + async def replace_after_read(name: str) -> bytes | None: + previous: Final = await health.raw.get(name) + await health.raw.set(name, json.dumps("replacement"), ex=60) + return previous + + with patch.object(health.cache.init_async_client(), "get", new=replace_after_read): + await health.manager.release_health_check_lock() + assert await health.raw.get(LOCK) is None + + +@pytest.mark.asyncio +async def test_no_redis_runs_locally_without_polling() -> None: + manager: Final = SharedHealthCheckManager() + + async def no_wait(delay: float) -> None: + pytest.fail(f"unexpected local polling delay {delay}") + + with patch("asyncio.sleep", new=no_wait): + assert await manager.perform_shared_health_check([]) == ([], [], {}) + assert await manager.acquire_health_check_lock() is False + assert await manager.is_health_check_in_progress() is False + assert await manager.get_cached_health_check_results() is None + assert await manager.release_health_check_lock() is None + assert await manager.cache_health_check_results([], []) is None + assert await manager.get_health_check_status() == { + "pod_id": manager.pod_id, + "redis_available": False, + "lock_ttl": 60, + "cache_ttl": 300, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "raw, expected", + [ + (b"42", 42), + (b"[]", []), + (b"false", False), + (b'""', ""), + (b"'legacy'", "legacy"), + (b"null", None), + (b"broken", None), + (b"\xff", None), + ], +) +async def test_legacy_owner_decoding_keeps_lock_presence(health: HealthRedis, raw: bytes, expected: object) -> None: + await health.raw.set(LOCK, raw) + assert await health.manager.is_health_check_in_progress() is (expected is not None) + status: Final = await health.manager.get_health_check_status() + assert status["lock_owner"] == expected + assert status["lock_in_progress"] is (expected is not None) + await health.manager.release_health_check_lock() + assert await health.raw.get(LOCK) == raw + + +@pytest.mark.asyncio +async def test_keys_expire_at_caller_selected_seconds(health: HealthRedis, monkeypatch: pytest.MonkeyPatch) -> None: + assert await health.manager.acquire_health_check_lock() is True + await health.manager.cache_health_check_results([], []) + monkeypatch.setattr(time, "time", lambda: NOW + 61) + assert await health.raw.pttl(LOCK) == -2 + assert await health.raw.pttl(RESULTS) == 239_000 + assert await health.manager.is_health_check_in_progress() is False + monkeypatch.setattr(time, "time", lambda: NOW + 301) + assert await health.manager.get_cached_health_check_results() is None + assert await health.raw.pttl(RESULTS) == -2 + + +def _redis_metric(name: str) -> float: + samples: Final = chain.from_iterable(metric.samples for metric in REGISTRY.collect()) + return sum(sample.value for sample in samples if sample.name == name) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ttl", [1.9, 0.5, 0.9999999, 1.9999999, -0.5]) +async def test_fractional_ttl_matches_legacy_seconds(health: HealthRedis, ttl: float) -> None: + setattr(health.manager, "lock_ttl", ttl) + setattr(health.manager, "health_check_ttl", ttl) + assert await health.manager.acquire_health_check_lock() is (int(ttl) > 0) + await health.manager.cache_health_check_results([], []) + expected: Final = int(ttl) * 1000 if int(ttl) > 0 else -2 + assert await health.raw.pttl(LOCK) == expected + assert await health.raw.pttl(RESULTS) == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ttl", [0, -1]) +async def test_invalid_ttl_emits_failure_events(health: HealthRedis, monkeypatch: pytest.MonkeyPatch, ttl: int) -> None: + monkeypatch.setattr(litellm, "service_callback", ["prometheus_system"]) + health.manager.lock_ttl = ttl + health.manager.health_check_ttl = ttl + pending_before: Final = asyncio.all_tasks() + before: Final = _redis_metric("litellm_redis_failed_requests_total") + breaker_before: Final = _redis_metric("litellm_redis_circuit_breaker_failures_total") + assert await health.manager.acquire_health_check_lock() is False + await health.manager.cache_health_check_results([], []) + await asyncio.gather(*(asyncio.all_tasks() - pending_before)) + assert _redis_metric("litellm_redis_failed_requests_total") - before == 2 + assert _redis_metric("litellm_redis_circuit_breaker_failures_total") == breaker_before + + +@pytest.mark.asyncio +async def test_zero_ttl_outage_still_opens_breaker(health: HealthRedis) -> None: + health.server.connected = False + health.manager.lock_ttl = 0 + for _ in range(REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD): + assert await health.manager.acquire_health_check_lock() is False + health.server.connected = True + health.manager.lock_ttl = 60 + assert await health.manager.acquire_health_check_lock() is False + assert await health.raw.get(LOCK) is None From 069ded7b6778be7b51fd89dec2228ac41bc8e871 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 19:48:23 +0000 Subject: [PATCH 02/17] refactor(redis): isolate managed file metadata store --- .../proxy/hooks/managed_files.py | 354 +++---- litellm/redis/stores/managed_files.py | 65 ++ .../proxy/hooks/test_managed_files.py | 903 +++++------------- .../proxy/test_managed_files_hook.py | 197 ++-- tests/unit/redis/stores/test_managed_files.py | 264 +++++ 5 files changed, 845 insertions(+), 938 deletions(-) create mode 100644 litellm/redis/stores/managed_files.py create mode 100644 tests/unit/redis/stores/test_managed_files.py diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 5b54e7d72a0..f5ac8718be6 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -4,19 +4,18 @@ import asyncio import base64 import json +import math import time from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence +from dataclasses import dataclass from datetime import datetime, timedelta from itertools import chain from types import MappingProxyType from typing import ( TYPE_CHECKING, Any, - Dict, Final, - List, Literal, - Optional, Protocol, TypedDict, Union, @@ -36,6 +35,8 @@ from litellm import Router, verbose_logger from litellm._internal_context import with_service_target from litellm._uuid import uuid from litellm.caching.caching import DualCache +from litellm.caching.redis_batch import active_request_redis_batch +from litellm.caching.redis_cache import RedisCache from litellm.constants import ( BATCH_OUTPUT_FILE_FALLBACK_RETRY_AFTER_SECONDS, BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS, @@ -93,6 +94,8 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import ( request_tags_from_metadata, ) +from litellm.redis.stores.managed_files import ManagedFileStore, TransportManagedFileStore, file_key +from litellm.redis.transport.contract import RedisFailure from litellm.repositories.managed_file_repository import ManagedFileRepository from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ( # pyright: ignore[reportAttributeAccessIssue] @@ -158,7 +161,7 @@ def _decode_json_blob(blob: object) -> object: return json.loads(blob) if isinstance(blob, str) else blob -def _parse_managed_batch_row(row: "PrismaManagedObjectRow") -> Optional[LiteLLMBatch]: +def _parse_managed_batch_row(row: "PrismaManagedObjectRow") -> LiteLLMBatch | None: try: batch_obj: Final = LiteLLMBatch.model_validate(_decode_json_blob(row.file_object)) except Exception as e: @@ -168,7 +171,7 @@ def _parse_managed_batch_row(row: "PrismaManagedObjectRow") -> Optional[LiteLLMB return batch_obj -def _parse_managed_file_object(raw_file_object: object, unified_file_id: str) -> Optional[OpenAIFileObject]: +def _parse_managed_file_object(raw_file_object: object, unified_file_id: str) -> OpenAIFileObject | None: if raw_file_object is None: return None try: @@ -311,16 +314,16 @@ class _ManagedFileRow(Protocol): file_object: OpenAIFileObject model_mappings: Mapping[str, str] flat_model_file_ids: Sequence[str] - storage_backend: Optional[str] - storage_url: Optional[str] - created_by: Optional[str] - team_id: Optional[str] + storage_backend: str | None + storage_url: str | None + created_by: str | None + team_id: str | None def model_dump(self) -> Mapping[str, object]: ... class _ManagedFileTableActions(Protocol): - async def find_first(self, where: Mapping[str, object]) -> Optional[_ManagedFileRow]: ... + async def find_first(self, where: Mapping[str, object]) -> _ManagedFileRow | None: ... async def update_many(self, where: Mapping[str, object], data: Mapping[str, object]) -> int: ... @@ -328,24 +331,24 @@ class _ManagedFileTableActions(Protocol): self, where: Mapping[str, object], take: int = ..., - order: Union[Mapping[str, str], Sequence[Mapping[str, str]]] = ..., + order: Mapping[str, str] | Sequence[Mapping[str, str]] = ..., cursor: Mapping[str, str] = ..., skip: int = ..., ) -> Sequence[_ManagedFileRow]: ... async def upsert(self, where: Mapping[str, str], data: Mapping[str, Mapping[str, object]]) -> _ManagedFileRow: ... - async def delete(self, where: Mapping[str, str]) -> Optional[_ManagedFileRow]: ... + async def delete(self, where: Mapping[str, str]) -> _ManagedFileRow | None: ... class _ManagedObjectTableActions(Protocol): - async def find_first(self, where: Mapping[str, object]) -> "Optional[PrismaManagedObjectRow]": ... + async def find_first(self, where: Mapping[str, object]) -> "PrismaManagedObjectRow | None": ... async def find_many( self, where: Mapping[str, object], take: int, - order: Union[Mapping[str, str], Sequence[Mapping[str, str]]], + order: Mapping[str, str] | Sequence[Mapping[str, str]], cursor: Mapping[str, str] = ..., skip: int = ..., ) -> "Sequence[PrismaManagedObjectRow]": ... @@ -517,11 +520,7 @@ def _storage_metadata_of(file_object: OpenAIFileObject | None) -> Mapping[str, s "Mapping[str, object]", getattr(file_object, "_hidden_params", None) or {} ) return MappingProxyType( - { - key: value - for key in ("storage_backend", "storage_url") - if isinstance(value := hidden_params.get(key), str) - } + {key: value for key in ("storage_backend", "storage_url") if isinstance(value := hidden_params.get(key), str)} ) @@ -578,6 +577,19 @@ def _keyed_get(source: object, key: str, default: object = None) -> object: return require_attributes(source, _SupportsKeyedGet, "get").get(key, default) +class _LocalMetadataCache(Protocol): + async def async_get_cache(self, key: str, *, local_only: bool) -> object: ... + + async def async_set_cache(self, key: str, value: object, *, local_only: bool) -> None: ... + + +@dataclass(frozen=True, slots=True) +class _MetadataCaches: + local: _LocalMetadataCache + redis: RedisCache | None + in_memory_ttl: float | None + + class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Class variables or attributes def __init__( @@ -593,6 +605,56 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): self._sleep = sleep self._clock = clock + def _metadata_cache(self) -> _MetadataCaches: + cache: Final = self.internal_usage_cache + dual: Final = cache if isinstance(cache, DualCache) else getattr(cache, "dual_cache", None) + if isinstance(dual, DualCache): + return _MetadataCaches(local=dual, redis=dual.redis_cache, in_memory_ttl=dual.default_in_memory_ttl) + return _MetadataCaches( + local=require_attributes(cache, _LocalMetadataCache, "async_get_cache", "async_set_cache"), + redis=None, + in_memory_ttl=None, + ) + + async def _cache_file(self, file_id: str, record: Mapping[str, object] | None) -> None: + cache: Final = self._metadata_cache() + local_cache: Final = cache.local + await local_cache.async_set_cache(file_key(file_id), record, local_only=True) + if cache.redis is None: + return + store: Final[ManagedFileStore] = TransportManagedFileStore(cache.redis.transport()) + try: + seconds: Final = ( + int(cache.in_memory_ttl) + if cache.in_memory_ttl is not None and not math.isnan(cache.in_memory_ttl) + else cache.redis.default_ttl + ) + ttl: Final = timedelta(seconds=seconds) + except (OverflowError, ValueError): + return + if record is None: + await store.invalidate_file(file_id, ttl=ttl) + else: + await store.write_file(file_id, record, ttl=ttl) + + async def _cached_file(self, file_id: str) -> object: + cache: Final = self._metadata_cache() + local_cache: Final = cache.local + local: Final = await local_cache.async_get_cache(file_key(file_id), local_only=True) + if local is not None or cache.redis is None: + return local + request_batch: Final = active_request_redis_batch(cache.redis) + if request_batch is not None and request_batch.read_as_missing(file_key(file_id)): + return None + store: Final[ManagedFileStore] = TransportManagedFileStore(cache.redis.transport()) + result: Final = await store.read_file(file_id) + if result is None or isinstance(result, RedisFailure): + return None + value: Final = result.value + if value is not None: + await local_cache.async_set_cache(file_key(file_id), value, local_only=True) + return value + @staticmethod def _get_prometheus_logger(): """Find PrometheusLogger from litellm.callbacks, if registered.""" @@ -604,9 +666,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): async def store_unified_file_id( self, file_id: str, - file_object: Optional[OpenAIFileObject], - litellm_parent_otel_span: Optional[Span], - model_mappings: Dict[str, str], + file_object: OpenAIFileObject | None, + litellm_parent_otel_span: Span | None, + model_mappings: dict[str, str], user_api_key_dict: UserAPIKeyAuth, ) -> None: verbose_logger.info(f"Storing LiteLLM Managed File object with id={file_id} in cache") @@ -623,11 +685,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): storage_backend=storage_metadata.get("storage_backend"), storage_url=storage_metadata.get("storage_url"), ) - await self.internal_usage_cache.async_set_cache( - key=file_id, - value=litellm_managed_file_object.model_dump(), - litellm_parent_otel_span=litellm_parent_otel_span, - ) + await self._cache_file(file_id, litellm_managed_file_object.model_dump()) ## STORE MODEL MAPPINGS IN DB @@ -663,7 +721,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) verbose_logger.debug(f"LiteLLM Managed File object with id={file_id} stored in db: {result}") - async def _resolve_creator_org_id(self, user_api_key_dict: UserAPIKeyAuth) -> Optional[str]: + async def _resolve_creator_org_id(self, user_api_key_dict: UserAPIKeyAuth) -> str | None: if user_api_key_dict.org_id: return user_api_key_dict.org_id if not user_api_key_dict.team_id: @@ -689,7 +747,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): self, unified_object_id: str, file_object: Union[LiteLLMBatch, LiteLLMFineTuningJob, "ResponsesAPIResponse"], - litellm_parent_otel_span: Optional[Span], + litellm_parent_otel_span: Span | None, model_object_id: str, file_purpose: Literal["batch", "fine-tune", "response"], user_api_key_dict: UserAPIKeyAuth, @@ -817,16 +875,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): @with_service_target(_MANAGED_FILES_TARGET) async def get_unified_file_id( - self, file_id: str, litellm_parent_otel_span: Optional[Span] = None - ) -> Optional[LiteLLM_ManagedFileTable]: + self, file_id: str, litellm_parent_otel_span: Span | None = None + ) -> LiteLLM_ManagedFileTable | None: ## CHECK CACHE - result = cast( - Optional[dict], - await self.internal_usage_cache.async_get_cache( - key=file_id, - litellm_parent_otel_span=litellm_parent_otel_span, - ), - ) + result: Final = await self._cached_file(file_id) if result: return LiteLLM_ManagedFileTable.model_validate(result) @@ -840,18 +892,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): @with_service_target(_MANAGED_FILES_TARGET) async def delete_unified_file_id( - self, file_id: str, litellm_parent_otel_span: Optional[Span] = None + self, file_id: str, litellm_parent_otel_span: Span | None = None ) -> OpenAIFileObject: ## get old value initial_value = await _managed_file_table(self.prisma_client).find_first(where={"unified_file_id": file_id}) if initial_value is None: raise Exception(f"LiteLLM Managed File object with id={file_id} not found") ## delete old value - await self.internal_usage_cache.async_set_cache( - key=file_id, - value=None, - litellm_parent_otel_span=litellm_parent_otel_span, - ) + await self._cache_file(file_id, None) await _managed_file_table(self.prisma_client).delete(where={"unified_file_id": file_id}) return initial_value.file_object @@ -887,9 +935,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): detail=f"Object not found: {unified_object_id}", ) - async def enforce_batch_object_access( - self, object_id: str, user_api_key_dict: UserAPIKeyAuth - ) -> None: + async def enforce_batch_object_access(self, object_id: str, user_api_key_dict: UserAPIKeyAuth) -> None: """Deny access to a provider-format batch id owned by another caller. Ids with no ownership row (batches created before ownership tracking, @@ -913,9 +959,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): detail=f"User {user_api_key_dict.user_id} does not have access to the object {object_id}", ) - async def enforce_provider_file_access( - self, file_id: str, user_api_key_dict: UserAPIKeyAuth - ) -> None: + async def enforce_provider_file_access(self, file_id: str, user_api_key_dict: UserAPIKeyAuth) -> None: """Deny access to a provider-format file id owned by another caller. Ownership rows for provider-format ids are written when a managed @@ -940,7 +984,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) async def store_batch_output_file_ownership( - self, response: LiteLLMBatch, litellm_parent_otel_span: Optional[Span] + self, response: LiteLLMBatch, litellm_parent_otel_span: Span | None ) -> None: """Record ownership rows for a batch's provider-format output/error file ids, inherited from the owning batch row (never the caller), so @@ -957,16 +1001,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): return if self.prisma_client is None: return - batch_row = await _managed_object_table(self.prisma_client).find_first( - where={"unified_object_id": response.id} - ) - if batch_row is None or ( - batch_row.created_by is None and batch_row.team_id is None - ): + batch_row = await _managed_object_table(self.prisma_client).find_first(where={"unified_object_id": response.id}) + if batch_row is None or (batch_row.created_by is None and batch_row.team_id is None): return - owner_identity = UserAPIKeyAuth( - user_id=batch_row.created_by, team_id=batch_row.team_id - ) + owner_identity = UserAPIKeyAuth(user_id=batch_row.created_by, team_id=batch_row.team_id) for file_id in provider_file_ids: model_name = decode_model_from_file_id(file_id) raw_file_id = get_original_file_id(file_id) @@ -1084,8 +1122,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): if raise_on_failure: if isinstance(error, TimeoutError) and not str(error): raise TimeoutError( - "Provider file retrieve timed out " - f"after {BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS} seconds" + f"Provider file retrieve timed out after {BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS} seconds" ) from error raise return None, True @@ -1098,11 +1135,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): if not await ManagedFileRepository(self.prisma_client).update_file_object(stored.unified_file_id, file_object): return refreshed_row: Final = stored.model_copy(update={"file_object": file_object}) - await self.internal_usage_cache.async_set_cache( - key=stored.unified_file_id, - value=refreshed_row.model_dump(), - litellm_parent_otel_span=None, - ) + await self._cache_file(stored.unified_file_id, refreshed_row.model_dump()) async def _save_xai_refreshed_file_object( self, @@ -1117,14 +1150,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): sized_file_object, if_file_object=_stored_json_text(stored.file_object) if stored.file_object is not None else None, ): - await self.internal_usage_cache.async_set_cache( - key=stored.unified_file_id, value=None, litellm_parent_otel_span=None - ) + await self._cache_file(stored.unified_file_id, None) return None - await self.internal_usage_cache.async_set_cache( - key=stored.unified_file_id, - value=stored.model_copy(update={"file_object": sized_file_object}).model_dump(), - litellm_parent_otel_span=None, + await self._cache_file( + stored.unified_file_id, stored.model_copy(update={"file_object": sized_file_object}).model_dump() ) return sized_file_object @@ -1270,12 +1299,12 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): async def list_user_batches( self, user_api_key_dict: UserAPIKeyAuth, - limit: Optional[int] = None, - after: Optional[str] = None, - provider: Optional[str] = None, - target_model_names: Optional[str] = None, - llm_router: Optional[Router] = None, - ) -> Dict[str, object]: + limit: int | None = None, + after: str | None = None, + provider: str | None = None, + target_model_names: str | None = None, + llm_router: Router | None = None, + ) -> dict[str, object]: # Provider filtering is not supported for managed batches # This is because the encoded object ids stored in the managed objects table do not contain the provider information # To support provider filtering, we would need to store the provider information in the encoded object ids @@ -1305,7 +1334,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): if owner_filter is None: return build_list_page([]) - where_clause: Dict[str, object] = {"file_purpose": "batch", **owner_filter} + where_clause: dict[str, object] = {"file_purpose": "batch", **owner_filter} if after: cursor_row = await _managed_object_table(self.prisma_client).find_first( @@ -1329,7 +1358,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): async def _collect_listed_batches( self, where_clause: Mapping[str, object], - after: Optional[str], + after: str | None, wanted: int, user_api_key_dict: UserAPIKeyAuth, ) -> tuple[LiteLLMBatch, ...]: @@ -1341,7 +1370,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): and every chunk advances the keyset cursor, so the walk ends once the caller's rows are exhausted.""" matches: tuple[LiteLLMBatch, ...] = () # rebind-ok: accumulates survivors across chunks - cursor_id: Optional[str] = after # rebind-ok: keyset cursor advances to each chunk's last row + cursor_id: str | None = after # rebind-ok: keyset cursor advances to each chunk's last row chunk_size: int = wanted # rebind-ok: widens once a scan has to continue past the first chunk while len(matches) < wanted: cursor_args: _CursorPageArgs = {"cursor": {"unified_object_id": cursor_id}, "skip": 1} if cursor_id else {} @@ -1398,7 +1427,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): batch_obj: LiteLLMBatch, unified_id_by_raw_id: Mapping[str, str], user_api_key_dict: UserAPIKeyAuth, - ) -> Optional[LiteLLMBatch]: + ) -> LiteLLMBatch | None: apply_unified_file_ids(batch_obj, unified_id_by_raw_id) try: await ensure_batch_response_managed_file_ids( @@ -1450,8 +1479,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) async def get_user_created_file_ids( - self, user_api_key_dict: UserAPIKeyAuth, model_object_ids: List[str] - ) -> List[OpenAIFileObject]: + self, user_api_key_dict: UserAPIKeyAuth, model_object_ids: list[str] + ) -> list[OpenAIFileObject]: """ Get all file ids the caller is allowed to see for a list of model object ids. Service-account keys (no user_id) are scoped to their @@ -1476,8 +1505,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): if (parsed_file_object := _parse_managed_file_object(row.file_object, row.unified_file_id)) is not None ] - async def check_managed_file_id_access(self, data: Dict, user_api_key_dict: UserAPIKeyAuth) -> bool: - retrieve_file_id = cast(Optional[str], data.get("file_id")) + async def check_managed_file_id_access(self, data: dict, user_api_key_dict: UserAPIKeyAuth) -> bool: + retrieve_file_id = cast(str | None, data.get("file_id")) potential_file_id = _is_base64_encoded_unified_file_id(retrieve_file_id) if retrieve_file_id else False if potential_file_id and retrieve_file_id: if await self.can_user_call_unified_file_id(retrieve_file_id, user_api_key_dict): @@ -1488,9 +1517,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): detail=f"User {user_api_key_dict.user_id} does not have access to the file {retrieve_file_id}", ) if retrieve_file_id: - await self.enforce_provider_file_access( - retrieve_file_id, user_api_key_dict - ) + await self.enforce_provider_file_access(retrieve_file_id, user_api_key_dict) return False async def check_file_ids_access(self, file_ids: Sequence[str], user_api_key_dict: UserAPIKeyAuth) -> None: @@ -1518,9 +1545,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, - data: Dict, + data: dict, call_type: CallTypesLiteral, - ) -> Union[Exception, str, Dict, None]: + ) -> Exception | str | dict | None: """ - Detect litellm_proxy/ file_id - add dictionary of mappings of litellm_proxy/ file_id -> provider_file_id => {litellm_proxy/file_id: {"model_id": id, "file_id": provider_file_id}} @@ -1593,7 +1620,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): if _is_streamed_background_create(create_request): return _with_managed_background_marker(create_request) elif call_type == CallTypes.afile_content.value: - retrieve_file_id = cast(Optional[str], data.get("file_id")) + retrieve_file_id = cast(str | None, data.get("file_id")) potential_file_id = _is_base64_encoded_unified_file_id(retrieve_file_id) if retrieve_file_id else False if potential_file_id and "llm_output_file_id," in potential_file_id: model_id = self.get_model_id_from_unified_file_id(potential_file_id) @@ -1601,7 +1628,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): data["model"] = model_id data["file_id"] = self.get_output_file_id_from_unified_file_id(potential_file_id) elif call_type == CallTypes.acreate_batch.value: - input_file_id = cast(Optional[str], data.get("input_file_id")) + input_file_id = cast(str | None, data.get("input_file_id")) if input_file_id: await self.check_file_ids_access((input_file_id,), user_api_key_dict) model_file_id_mapping = await self.get_model_file_id_mapping( @@ -1615,8 +1642,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): or call_type == CallTypes.acancel_fine_tuning_job.value or call_type == CallTypes.aretrieve_fine_tuning_job.value ): - accessor_key: Optional[str] = None - retrieve_object_id: Optional[str] = None + accessor_key: str | None = None + retrieve_object_id: str | None = None if call_type == CallTypes.aretrieve_batch.value or call_type == CallTypes.acancel_batch.value: accessor_key = "batch_id" elif ( @@ -1626,7 +1653,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): accessor_key = "fine_tuning_job_id" if accessor_key: - retrieve_object_id = cast(Optional[str], data.get(accessor_key)) + retrieve_object_id = cast(str | None, data.get(accessor_key)) potential_llm_object_id = ( _is_base64_encoded_unified_file_id(retrieve_object_id) if retrieve_object_id else False @@ -1670,11 +1697,11 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): async def async_filter_deployments( self, model: str, - healthy_deployments: List, - messages: Optional[List[AllMessageValues]], - request_kwargs: Optional[Dict] = None, - parent_otel_span: Optional[Span] = None, - ) -> List[Dict]: + healthy_deployments: list, + messages: list[AllMessageValues] | None, + request_kwargs: dict | None = None, + parent_otel_span: Span | None = None, + ) -> list[dict]: if request_kwargs is None: return healthy_deployments @@ -1701,8 +1728,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): return holders if holders or "input_file_id" in request else healthy_deployments async def async_pre_call_deployment_hook( - self, kwargs: dict[str, object], call_type: Optional[CallTypes] - ) -> Optional[dict]: + self, kwargs: dict[str, object], call_type: CallTypes | None + ) -> dict | None: """ Allow modifying the request just before it's sent to the deployment. """ @@ -1710,18 +1737,18 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): if accessor_keys is None: return kwargs - model_file_id_mapping = cast(Optional[Dict[str, Dict[str, str]]], kwargs.get("model_file_id_mapping")) + model_file_id_mapping = cast(dict[str, dict[str, str]] | None, kwargs.get("model_file_id_mapping")) # model_info may be at top-level or nested under litellm_metadata # (batch/file operations use litellm_metadata) - model_id = cast(Optional[str], _keyed_get(kwargs.get("model_info", {}), "id")) + model_id = cast(str | None, _keyed_get(kwargs.get("model_info", {}), "id")) if model_id is None: model_id = cast( - Optional[str], + str | None, _keyed_get(_keyed_get(kwargs.get("litellm_metadata", {}), "model_info", {}), "id"), ) for accessor_key in accessor_keys: - input_file_id = cast(Optional[str], kwargs.get(accessor_key)) - mapped_file_id: Optional[str] = None + input_file_id = cast(str | None, kwargs.get(accessor_key)) + mapped_file_id: str | None = None if input_file_id and model_file_id_mapping and model_id: mapped_file_id = model_file_id_mapping.get(input_file_id, {}).get(model_id, None) if mapped_file_id: @@ -1729,7 +1756,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): return kwargs - def get_file_ids_from_messages(self, messages: List[AllMessageValues]) -> List[str]: + def get_file_ids_from_messages(self, messages: list[AllMessageValues]) -> list[str]: """ Gets file ids from messages """ @@ -1749,7 +1776,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): file_ids.append(file_id) return file_ids - def get_file_ids_from_responses_input(self, input: Union[str, List[Dict[str, object]]]) -> List[str]: + def get_file_ids_from_responses_input(self, input: str | list[dict[str, object]]) -> list[str]: """ Gets file ids from responses API input. @@ -1759,7 +1786,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): - type: "input_file" with file_id - content: a list that can contain items with type: "input_file" and file_id """ - file_ids: List[str] = [] + file_ids: list[str] = [] if isinstance(input, str): return file_ids @@ -1788,7 +1815,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): return file_ids - def get_file_ids_from_responses_tools(self, tools: List[Dict[str, object]]) -> List[str]: + def get_file_ids_from_responses_tools(self, tools: list[dict[str, object]]) -> list[str]: """ Gets file ids from responses API tools parameter. @@ -1800,7 +1827,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): } ] """ - file_ids: List[str] = [] + file_ids: list[str] = [] if not isinstance(tools, list): return file_ids @@ -1821,7 +1848,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): return file_ids - def get_vector_store_ids_from_file_search_tools(self, tools: List[Dict[str, object]]) -> List[str]: + def get_vector_store_ids_from_file_search_tools(self, tools: list[dict[str, object]]) -> list[str]: """ Extract unified vector_store_ids from file_search tools. @@ -1832,7 +1859,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): is_base64_encoded_unified_id, ) - vs_ids: List[str] = [] + vs_ids: list[str] = [] if not isinstance(tools, list): return vs_ids @@ -1850,7 +1877,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): async def check_vector_store_ids_access( self, - vector_store_ids: List[str], + vector_store_ids: list[str], user_api_key_dict: UserAPIKeyAuth, ) -> None: """ @@ -1876,7 +1903,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): return # Map each unified ID to its internal UUID for a single batch DB fetch - uuid_to_unified: Dict[str, str] = {} + uuid_to_unified: dict[str, str] = {} for vs_id in vector_store_ids: uuid = extract_unified_uuid_from_unified_id(vs_id) if uuid: @@ -1931,7 +1958,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): } """ - file_id_mapping: Dict[str, Dict[str, str]] = {} + file_id_mapping: dict[str, dict[str, str]] = {} litellm_managed_file_ids = [] for file_id in file_ids: @@ -1953,11 +1980,11 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): async def create_file_for_each_model( self, - llm_router: Optional[Router], + llm_router: Router | None, _create_file_request: CreateFileRequest, - target_model_names_list: List[str], + target_model_names_list: list[str], litellm_parent_otel_span: Span | None, - ) -> List[OpenAIFileObject]: + ) -> list[OpenAIFileObject]: if llm_router is None: raise Exception("LLM Router not initialized. Ensure models added to proxy.") responses = [] @@ -1971,7 +1998,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): self, create_file_request: CreateFileRequest, llm_router: Router, - target_model_names_list: List[str], + target_model_names_list: list[str], litellm_parent_otel_span: Span | None, user_api_key_dict: UserAPIKeyAuth, ) -> OpenAIFileObject: @@ -1990,7 +2017,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) ## STORE MODEL MAPPINGS IN DB - model_mappings: Dict[str, str] = {} + model_mappings: dict[str, str] = {} for file_object in responses: file_hidden_params = cast( # cast-ok: preserve mapping operations on dynamic file metadata @@ -2036,11 +2063,11 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): @staticmethod async def return_unified_file_id( - file_objects: List[OpenAIFileObject], + file_objects: list[OpenAIFileObject], create_file_request: CreateFileRequest, internal_usage_cache: InternalUsageCache, litellm_parent_otel_span: Span | None, - target_model_names_list: List[str], + target_model_names_list: list[str], ) -> OpenAIFileObject: ## GET THE FILE TYPE FROM THE CREATE FILE REQUEST _, file_type = extract_file_metadata(create_file_request["file"]) @@ -2085,7 +2112,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): unified_batch_id = SpecialEnums.LITELLM_MANAGED_BATCH_COMPLETE_STR.value.format(model_id, batch_id) return base64.urlsafe_b64encode(unified_batch_id.encode()).decode().rstrip("=") - def get_unified_output_file_id(self, output_file_id: str, model_id: str, model_name: Optional[str]) -> str: + def get_unified_output_file_id(self, output_file_id: str, model_id: str, model_name: str | None) -> str: deterministic_uuid: Final = uuid5(uuid5(NAMESPACE_URL, model_id), output_file_id) unified_output_file_id = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( "application/json", @@ -2106,7 +2133,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): return file_id.split(marker, 1)[1].split(";")[0] async def async_post_call_success_hook( - self, data: Dict, user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes + self, data: dict, user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes ) -> LLMResponseTypes: if isinstance(response, LiteLLMBatch): decoded_batch_id: Final = _is_base64_encoded_unified_file_id(response.id) @@ -2117,8 +2144,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): unified_file_id = response_hidden_params.get("unified_file_id") unified_batch_id = response_hidden_params.get("unified_batch_id") is_batch_create: Final = response_hidden_params.get(BATCH_CREATE_HIDDEN_PARAM) is True - model_id = cast(Optional[str], response_hidden_params.get("model_id")) - model_name = cast(Optional[str], response_hidden_params.get("model_name")) + model_id = cast(str | None, response_hidden_params.get("model_id")) + model_name = cast(str | None, response_hidden_params.get("model_name")) resolved_model_name = resolve_managed_output_file_model_name( unified_input_file_id=unified_file_id if isinstance(unified_file_id, str) else response.input_file_id, @@ -2206,7 +2233,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) unified_file_id = finetuning_response_hidden_params.get("unified_file_id") unified_finetuning_job_id = finetuning_response_hidden_params.get("unified_finetuning_job_id") - model_id = cast(Optional[str], finetuning_response_hidden_params.get("model_id")) + model_id = cast(str | None, finetuning_response_hidden_params.get("model_id")) original_response_id = response.id if (unified_file_id or unified_finetuning_job_id) and model_id: response.id = self.get_unified_generic_response_id(model_id=model_id, generic_response_id=response.id) @@ -2224,15 +2251,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): """ ## check if file object if hasattr(response, "data") and isinstance(response.data, list): - if all(isinstance(file_object, FileObject) for file_object in response.data): - ## Get all file id's - ## Check which file id's were created by the user - ## Filter the response to only include the files created by the user - ## Return the filtered response - file_ids = [ - file_object.id - for file_object in cast(List[FileObject], response.data) # type: ignore - ] + page_items: Final = TypeAdapter(tuple[object, ...]).validate_python(response.data) + file_objects: Final = tuple(item for item in page_items if isinstance(item, FileObject)) + if len(file_objects) == len(page_items): + file_ids = [file_object.id for file_object in file_objects] user_created_file_ids = await self.get_user_created_file_ids(user_api_key_dict, file_ids) ## Filter the response to only include the files created by the user response.data = user_created_file_ids # type: ignore @@ -2242,7 +2264,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): return response @staticmethod - def _scope_list_page_cursors(response: AsyncCursorPage, data: List[OpenAIFileObject]) -> None: + def _scope_list_page_cursors(response: AsyncCursorPage, data: list[OpenAIFileObject]) -> None: """Rebuild ``first_id`` / ``last_id`` from the caller-scoped page. The upstream cursors point at rows that were just filtered out, so @@ -2258,7 +2280,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): response.has_more = False async def afile_retrieve( - self, file_id: str, litellm_parent_otel_span: Optional[Span], llm_router: Optional[Router] = None + self, file_id: str, litellm_parent_otel_span: Span | None, llm_router: Router | None = None ) -> OpenAIFileObject: """Return public details for a managed file ID, refreshing a basic entry when possible.""" stored_file_object = await self.get_unified_file_id(file_id, litellm_parent_otel_span) @@ -2325,7 +2347,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): allow_default_provider=True, ) except Exception as e: - raise Exception(f"Failed to retrieve file {file_id} from provider: {str(e)}") from e + raise Exception(f"Failed to retrieve file {file_id} from provider: {e!s}") from e if not fetch_route_available: raise Exception( f"LiteLLM Managed File object with id={file_id} has no file_object and no provider route to fetch it" @@ -2336,14 +2358,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): async def afile_list( self, - purpose: Optional[str], - litellm_parent_otel_span: Optional[Span], + purpose: str | None, + litellm_parent_otel_span: Span | None, user_api_key_dict: UserAPIKeyAuth, - limit: Optional[int] = None, - after: Optional[str] = None, - model_scope: Optional[FileListModelScope] = None, - before: Optional[str] = None, - **data: Dict, + limit: int | None = None, + after: str | None = None, + model_scope: FileListModelScope | None = None, + before: str | None = None, + **data: dict, ) -> FileListPage: """List the managed files the caller owns, newest first. @@ -2404,7 +2426,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): return FileListPage.model_validate(build_list_page([])) page_size: Final = min(limit or MAX_FILE_LIST_LIMIT, MAX_FILE_LIST_LIMIT) - matches: Final[List[OpenAIFileObject]] = [] + matches: Final[list[OpenAIFileObject]] = [] cursor_id = cursor chunk_size = page_size + 1 direction: Final[Literal["asc", "desc"]] = "asc" if before else "desc" @@ -2471,7 +2493,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): verbose_logger.warning(f"Error checking batch polling configuration: {e}. Assuming disabled.") return False - async def _get_batches_referencing_file(self, file_id: str) -> List[Dict[str, object]]: + async def _get_batches_referencing_file(self, file_id: str) -> list[dict[str, object]]: """ Find batches that reference this file and still need cost tracking. Find batches that are in non-terminal state and have not yet been processed by CheckBatchCost. @@ -2607,9 +2629,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): async def afile_delete( self, file_id: str, - litellm_parent_otel_span: Optional[Span], + litellm_parent_otel_span: Span | None, llm_router: Router, - **data: Dict, + **data: dict, ) -> FileDeleted: # Check if file deletion should be blocked due to batch references @@ -2668,9 +2690,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): async def afile_content( self, file_id: str, - litellm_parent_otel_span: Optional[Span], + litellm_parent_otel_span: Span | None, llm_router: Router, - **data: Dict, + **data: dict, ) -> HttpxBinaryResponseContent: """ Get the content of a file from first model that has it @@ -2696,7 +2718,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # request params. credentials = llm_router.registry().credentials(model_id, form="with_provider") if credentials is not None: - data["_litellm_internal_model_credentials"] = cast(Dict, MappingProxyType(dict(credentials))) + data["_litellm_internal_model_credentials"] = cast(dict, MappingProxyType(dict(credentials))) else: data.pop("_litellm_internal_model_credentials", None) return cast( @@ -2722,9 +2744,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): async def _convert_storage_files_to_base64( self, - messages: List[AllMessageValues], - file_ids: List[str], - litellm_parent_otel_span: Optional[Span], + messages: list[AllMessageValues], + file_ids: list[str], + litellm_parent_otel_span: Span | None, ) -> None: """ Convert files stored in storage backends to base64 format for Vertex AI/Gemini. @@ -2757,9 +2779,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): try: storage_backend = get_storage_backend(storage_backend_name, prisma_client=self.prisma_client) except ValueError as e: - verbose_logger.warning( - f"Storage backend '{storage_backend_name}' error for file {file_id}: {str(e)}" - ) + verbose_logger.warning(f"Storage backend '{storage_backend_name}' error for file {file_id}: {e!s}") continue file_content = await storage_backend.download_file(storage_url) @@ -2774,7 +2794,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Update messages to use base64 instead of file_id self._update_messages_with_base64_data(messages, file_id, base64_data_uri, content_type) except Exception as e: - verbose_logger.exception(f"Error converting file {file_id} from storage backend to base64: {str(e)}") + verbose_logger.exception(f"Error converting file {file_id} from storage backend to base64: {e!s}") # Continue with other files even if one fails continue @@ -2800,7 +2820,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): def _update_messages_with_base64_data( self, - messages: List[AllMessageValues], + messages: list[AllMessageValues], file_id: str, base64_data_uri: str, content_type: str, @@ -2834,4 +2854,6 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): verbose_logger.debug( f"Converted file {file_id} from storage backend to base64 with format {content_type}" ) + + PROXY_LiteLLMManagedFiles = _PROXY_LiteLLMManagedFiles diff --git a/litellm/redis/stores/managed_files.py b/litellm/redis/stores/managed_files.py new file mode 100644 index 00000000000..d2e7d85fcc5 --- /dev/null +++ b/litellm/redis/stores/managed_files.py @@ -0,0 +1,65 @@ +import ast +import json +from collections.abc import Mapping +from dataclasses import dataclass +from datetime import timedelta +from typing import Final, Protocol + +from pydantic import TypeAdapter + +from litellm.redis.transport.contract import RedisFailure, RedisTransport + + +@dataclass(frozen=True, slots=True) +class ManagedFileValue: + value: object + + +_VALUE: Final[TypeAdapter[ManagedFileValue]] = TypeAdapter(ManagedFileValue) + + +def file_key(unified_file_id: str) -> str: + return unified_file_id + + +class ManagedFileStore(Protocol): + async def read_file(self, unified_file_id: str) -> ManagedFileValue | None | RedisFailure: ... + + async def write_file( + self, unified_file_id: str, record: Mapping[str, object], *, ttl: timedelta | None + ) -> None | RedisFailure: ... + + async def invalidate_file(self, unified_file_id: str, *, ttl: timedelta | None) -> None | RedisFailure: ... + + +@dataclass(frozen=True, slots=True) +class TransportManagedFileStore: + transport: RedisTransport + + async def read_file(self, unified_file_id: str) -> ManagedFileValue | None | RedisFailure: + raw: Final = await self.transport.get(file_key(unified_file_id)) + if raw is None or isinstance(raw, RedisFailure): + return raw + try: + return _VALUE.validate_python({"value": json.loads(raw.decode("utf-8"))}) + except (ValueError, UnicodeError, RecursionError): + try: + return _VALUE.validate_python({"value": ast.literal_eval(raw.decode("utf-8"))}) + except (ValueError, SyntaxError, UnicodeError, RecursionError) as exc: + return RedisFailure(kind="decode", message=str(exc)) + + async def write_file( + self, unified_file_id: str, record: Mapping[str, object], *, ttl: timedelta | None + ) -> None | RedisFailure: + try: + encoded: Final = json.dumps(record) + except (TypeError, ValueError) as exc: + return RedisFailure(kind="decode", message=str(exc)) + return await self._write(unified_file_id, encoded, ttl) + + async def invalidate_file(self, unified_file_id: str, *, ttl: timedelta | None) -> None | RedisFailure: + return await self._write(unified_file_id, "null", ttl) + + async def _write(self, unified_file_id: str, encoded: str, ttl: timedelta | None) -> None | RedisFailure: + result: Final = await self.transport.set_seconds(file_key(unified_file_id), encoded, ttl=ttl) + return result if isinstance(result, RedisFailure) else None diff --git a/tests/unit/enterprise/proxy/hooks/test_managed_files.py b/tests/unit/enterprise/proxy/hooks/test_managed_files.py index 1755ab97689..f24e05c6d41 100644 --- a/tests/unit/enterprise/proxy/hooks/test_managed_files.py +++ b/tests/unit/enterprise/proxy/hooks/test_managed_files.py @@ -1,8 +1,8 @@ +import asyncio import base64 import json import logging import time -import asyncio import uuid from collections.abc import Awaitable, Callable, Mapping from types import SimpleNamespace @@ -23,33 +23,27 @@ from litellm_enterprise.proxy.hooks.managed_files import ( PROXY_LiteLLMManagedFiles, _provider_file_retrieve_credentials, ) + from litellm.caching import DualCache from litellm.proxy._types import CallTypes, LiteLLM_ManagedFileTable, UserAPIKeyAuth from litellm.proxy.openai_files_endpoints.common_utils import ( BATCH_CREATE_HIDDEN_PARAM, - is_base64_encoded_unified_file_id, encode_file_id_with_model, + is_base64_encoded_unified_file_id, ) +from litellm.proxy.utils import InternalUsageCache from litellm.types.llms.openai import OpenAIFileObject from tests.unit.litellm_core_utils.attribute_forwarder import AttributeForwarder class _InMemoryManagedFileTable: def __init__(self, *rows: LiteLLM_ManagedFileTable) -> None: - self.rows: dict[str, LiteLLM_ManagedFileTable] = { - row.unified_file_id: row for row in rows - } + self.rows: dict[str, LiteLLM_ManagedFileTable] = {row.unified_file_id: row for row in rows} self.find_first_calls: list[Mapping[str, object]] = [] - self.upsert_calls: list[ - tuple[Mapping[str, object], Mapping[str, Mapping[str, object]]] - ] = [] - self.update_many_calls: list[ - tuple[Mapping[str, object], Mapping[str, object]] - ] = [] + self.upsert_calls: list[tuple[Mapping[str, object], Mapping[str, Mapping[str, object]]]] = [] + self.update_many_calls: list[tuple[Mapping[str, object], Mapping[str, object]]] = [] - async def find_first( - self, where: Mapping[str, object] - ) -> LiteLLM_ManagedFileTable | None: + async def find_first(self, where: Mapping[str, object]) -> LiteLLM_ManagedFileTable | None: self.find_first_calls.append(where) unified_file_id: Final = where.get("unified_file_id") if isinstance(unified_file_id, str): @@ -73,11 +67,7 @@ class _InMemoryManagedFileTable: unified_file_id: Final = cast(str, where["unified_file_id"]) previous_row: Final = self.rows.get(unified_file_id) values: Final = { - **( - cast(dict[str, object], previous_row.model_dump()) - if previous_row is not None - else {} - ), + **(cast(dict[str, object], previous_row.model_dump()) if previous_row is not None else {}), **data["update" if previous_row is not None else "create"], } raw_model_mappings: Final = values.get("model_mappings", "{}") @@ -89,9 +79,7 @@ class _InMemoryManagedFileTable: raw_file_object: Final = values.get("file_object") file_object: Final = cast( dict[str, object] | None, - json.loads(raw_file_object) - if isinstance(raw_file_object, str) - else raw_file_object, + json.loads(raw_file_object) if isinstance(raw_file_object, str) else raw_file_object, ) row: Final = LiteLLM_ManagedFileTable.model_validate( { @@ -118,17 +106,11 @@ class _InMemoryManagedFileTable: None if existing_row.file_object is None else existing_row.file_object.model_dump(mode="json") ): return 0 - file_object: Final = OpenAIFileObject.model_validate( - json.loads(cast(str, data["file_object"])) - ) - self.rows[unified_file_id] = existing_row.model_copy( - update={"file_object": file_object} - ) + file_object: Final = OpenAIFileObject.model_validate(json.loads(cast(str, data["file_object"]))) + self.rows[unified_file_id] = existing_row.model_copy(update={"file_object": file_object}) return 1 - async def delete( - self, where: Mapping[str, object] - ) -> LiteLLM_ManagedFileTable | None: + async def delete(self, where: Mapping[str, object]) -> LiteLLM_ManagedFileTable | None: return self.rows.pop(cast(str, where["unified_file_id"]), None) async def find_many( @@ -143,9 +125,7 @@ class _InMemoryManagedObjectTable: async def find_first(self, where: Mapping[str, object]) -> None: return None - async def update_many( - self, where: Mapping[str, object], data: Mapping[str, object] - ) -> int: + async def update_many(self, where: Mapping[str, object], data: Mapping[str, object]) -> int: return 0 async def upsert( @@ -174,9 +154,7 @@ def _managed_files_with_fake_prisma( ) -> tuple[PROXY_LiteLLMManagedFiles, _InMemoryManagedFileTable]: managed_file_table: Final = _InMemoryManagedFileTable(*rows) prisma_client: Final = _InMemoryManagedFilesPrismaClient(managed_file_table) - managed_files: Final = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=prisma_client, sleep=sleep - ) + managed_files: Final = PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=prisma_client, sleep=sleep) return managed_files, managed_file_table @@ -209,6 +187,7 @@ def _marked_fallback_file_row( def _batch_response_for_output_listing() -> "LiteLLMBatch": from openai.types.batch import BatchRequestCounts + from litellm.types.utils import LiteLLMBatch response: Final = LiteLLMBatch( @@ -252,9 +231,7 @@ def test_provider_credentials_warning_sanitizes_newlines( ) -> None: model_id: Final = "deployment-123\nforged warning" router: Final = MagicMock() - router.registry.return_value.credentials.side_effect = RuntimeError( - "credential lookup failed\nforged warning" - ) + router.registry.return_value.credentials.side_effect = RuntimeError("credential lookup failed\nforged warning") with caplog.at_level(logging.WARNING, logger="LiteLLM"): credentials: Final = _provider_file_retrieve_credentials( @@ -276,9 +253,7 @@ def test_provider_credentials_warning_sanitizes_newlines( def test_get_file_ids_from_messages(): - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=MagicMock() - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(InternalUsageCache(DualCache()), prisma_client=MagicMock()) messages = [ { "role": "user", @@ -300,9 +275,7 @@ def test_get_file_ids_from_messages(): def test_get_file_ids_from_messages_skips_bedrock_content_blocks_without_type(): - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=MagicMock() - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(InternalUsageCache(DualCache()), prisma_client=MagicMock()) messages = [ { "role": "user", @@ -340,13 +313,9 @@ async def test_async_pre_call_hook_batch_retrieve(): return_value = MagicMock() return_value.created_by = "123" prisma_client.db.litellm_managedobjecttable.find_first.return_value = return_value - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=prisma_client - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=prisma_client) data = { - "user_api_key_dict": UserAPIKeyAuth( - user_id="123", parent_otel_span=MagicMock() - ), + "user_api_key_dict": UserAPIKeyAuth(user_id="123", parent_otel_span=MagicMock()), "data": { "batch_id": "bGl0ZWxsbV9wcm94eTttb2RlbF9pZDpteS1nZW5lcmFsLWF6dXJlLWRlcGxveW1lbnQ7bGxtX2JhdGNoX2lkOmJhdGNoX2EzMjJiNmJhLWFjN2UtNDg4OC05MjljLTFhZDM0NDJmMDZlZA", }, @@ -385,9 +354,7 @@ async def test_async_pre_call_deployment_hook_resolves_model_id_from_litellm_met async_pre_call_deployment_hook must check both locations so the managed file ID is resolved to the provider-specific file ID. """ - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=MagicMock() - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(InternalUsageCache(DualCache()), prisma_client=MagicMock()) managed_file_id = "managed-file-abc" model_id = "deployment-xyz" @@ -404,13 +371,11 @@ async def test_async_pre_call_deployment_hook_resolves_model_id_from_litellm_met }, } - result = await proxy_managed_files.async_pre_call_deployment_hook( - kwargs=kwargs, call_type=CallTypes.acreate_batch - ) + result = await proxy_managed_files.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.acreate_batch) - assert ( - result["input_file_id"] == provider_file_id - ), f"Expected provider file ID '{provider_file_id}', got '{result['input_file_id']}'" + assert result["input_file_id"] == provider_file_id, ( + f"Expected provider file ID '{provider_file_id}', got '{result['input_file_id']}'" + ) @pytest.mark.asyncio @@ -419,9 +384,7 @@ async def test_async_pre_call_deployment_hook_prefers_top_level_model_info(): When model_info exists at top-level kwargs, async_pre_call_deployment_hook should use it without falling back to litellm_metadata. """ - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=MagicMock() - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(InternalUsageCache(DualCache()), prisma_client=MagicMock()) managed_file_id = "managed-file-abc" top_level_model_id = "deployment-top" @@ -443,25 +406,25 @@ async def test_async_pre_call_deployment_hook_prefers_top_level_model_info(): }, } - result = await proxy_managed_files.async_pre_call_deployment_hook( - kwargs=kwargs, call_type=CallTypes.acreate_batch - ) + result = await proxy_managed_files.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.acreate_batch) - assert ( - result["input_file_id"] == top_level_provider_file - ), "Should prefer top-level model_info over litellm_metadata" + assert result["input_file_id"] == top_level_provider_file, ( + "Should prefer top-level model_info over litellm_metadata" + ) @pytest.mark.asyncio async def test_async_pre_call_deployment_hook_reads_model_info_that_forwards_its_fields(): - proxy_managed_files: Final = PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=MagicMock()) + proxy_managed_files: Final = PROXY_LiteLLMManagedFiles(InternalUsageCache(DualCache()), prisma_client=MagicMock()) kwargs: Final = { "input_file_id": "managed-file-abc", "model_file_id_mapping": {"managed-file-abc": {"deployment-xyz": "file-provider-123"}}, "model_info": AttributeForwarder({"id": "deployment-xyz"}), } - result: Final = await proxy_managed_files.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.acreate_batch) + result: Final = await proxy_managed_files.async_pre_call_deployment_hook( + kwargs=kwargs, call_type=CallTypes.acreate_batch + ) assert result["input_file_id"] == "file-provider-123" @@ -472,9 +435,7 @@ async def test_async_pre_call_deployment_hook_no_model_info_leaves_file_id_uncha When model_info is absent from both top-level and litellm_metadata, the managed file ID should remain unchanged. """ - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=MagicMock() - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(InternalUsageCache(DualCache()), prisma_client=MagicMock()) managed_file_id = "managed-file-abc" @@ -485,13 +446,11 @@ async def test_async_pre_call_deployment_hook_no_model_info_leaves_file_id_uncha }, } - result = await proxy_managed_files.async_pre_call_deployment_hook( - kwargs=kwargs, call_type=CallTypes.acreate_batch - ) + result = await proxy_managed_files.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.acreate_batch) - assert ( - result["input_file_id"] == managed_file_id - ), "File ID should remain unchanged when model_info is not available" + assert result["input_file_id"] == managed_file_id, ( + "File ID should remain unchanged when model_info is not available" + ) # def test_list_managed_files(): @@ -624,9 +583,7 @@ async def test_async_post_call_success_hook_for_unified_finetuning_job(): "unified_file_id": unified_file_id, "model_id": "gpt-3.5-turbo-0613", } - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=AsyncMock() - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=AsyncMock()) data = { "user_api_key_dict": {"parent_otel_span": MagicMock()}, } @@ -649,13 +606,9 @@ async def test_async_pre_call_hook_for_unified_finetuning_job(): return_value = MagicMock() return_value.created_by = "123" prisma_client.db.litellm_managedobjecttable.find_first.return_value = return_value - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=prisma_client - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=prisma_client) data = { - "user_api_key_dict": UserAPIKeyAuth( - user_id="123", parent_otel_span=MagicMock() - ), + "user_api_key_dict": UserAPIKeyAuth(user_id="123", parent_otel_span=MagicMock()), "data": { "fine_tuning_job_id": "bGl0ZWxsbV9wcm94eTttb2RlbF9pZDo0OTIxODU4MWY3OGViZTllZjE4NDE0ZmE0ZjdmYjlmYTc0YzA5NWVkMTEyY2E4NDBkZDU2ZGZmZTliZDMwZGQxO2dlbmVyaWNfcmVzcG9uc2VfaWQ6ZnRqb2ItalRCeXM3YlZzYnlaRE93TDlHbHBZcVhS", }, @@ -668,9 +621,7 @@ async def test_async_pre_call_hook_for_unified_finetuning_job(): @pytest.mark.asyncio -@pytest.mark.parametrize( - "call_type", ["afile_content", "afile_delete", "afile_retrieve"] -) +@pytest.mark.parametrize("call_type", ["afile_content", "afile_delete", "afile_retrieve"]) async def test_can_user_call_unified_file_id(call_type): """ Test that on file retrieve, delete, and content we check if the user has access to the file @@ -681,16 +632,12 @@ async def test_can_user_call_unified_file_id(call_type): return_value = MagicMock() return_value.created_by = "123" prisma_client.db.litellm_managedfiletable.find_first.return_value = return_value - proxy_managed_files = PROXY_LiteLLMManagedFiles( - MagicMock(), prisma_client=prisma_client - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(MagicMock(), prisma_client=prisma_client) unified_file_id = "bGl0ZWxsbV9wcm94eTphcHBsaWNhdGlvbi9vY3RldC1zdHJlYW07dW5pZmllZF9pZCxmMTNlNDAzZS01YWM3LTRhZjktOGQzNS0wNDgwZDMxOTgyYTg7dGFyZ2V0X21vZGVsX25hbWVzLGdwdC00by1taW5pLW9wZW5haTtsbG1fb3V0cHV0X2ZpbGVfaWQsZmlsZS1Ib3UxZDFXc3c1SDNKcjFMYllpZDJiO2xsbV9vdXRwdXRfZmlsZV9tb2RlbF9pZCxmODBiNWU2NzQ1NzdkNjkyMjM4YmVhNTIxZDdiMGI5ZGYyY2FmMTEwMTU2YmU5YzBjM2NjMmNkNTBjOTM1ZDI0" with pytest.raises(HTTPException): await proxy_managed_files.async_pre_call_hook( - user_api_key_dict=UserAPIKeyAuth( - user_id="456", parent_otel_span=MagicMock() - ), + user_api_key_dict=UserAPIKeyAuth(user_id="456", parent_otel_span=MagicMock()), cache=MagicMock(), data={"file_id": unified_file_id}, call_type=call_type, @@ -708,9 +655,7 @@ async def test_router_acreate_batch_only_selects_from_file_id_mapping(monkeypatc return_value = MagicMock() return_value.created_by = "123" prisma_client.db.litellm_managedobjecttable.find_first.return_value = return_value - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=prisma_client - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=prisma_client) monkeypatch.setattr( litellm, @@ -737,9 +682,7 @@ async def test_router_acreate_batch_only_selects_from_file_id_mapping(monkeypatc model_file_id_mapping = {file_id: {"5678": "file-LLCUFB52uTMa9hNYjtew9V"}} - with patch.object( - litellm, "acreate_batch", return_value=AsyncMock() - ) as mock_acreate_batch: + with patch.object(litellm, "acreate_batch", return_value=AsyncMock()) as mock_acreate_batch: for _ in range(1000): await router.acreate_batch( model="gpt-5-mini", @@ -753,9 +696,9 @@ async def test_router_acreate_batch_only_selects_from_file_id_mapping(monkeypatc @pytest.mark.asyncio async def test_output_file_id_for_batch_retrieve(): - import litellm.proxy.proxy_server as proxy_server_module from openai.types.batch import BatchRequestCounts + import litellm.proxy.proxy_server as proxy_server_module from litellm.proxy._types import UserAPIKeyAuth from litellm.types.llms.openai import OpenAIFileObject from litellm.types.utils import LiteLLMBatch @@ -824,10 +767,7 @@ async def test_output_file_id_for_batch_retrieve(): stored_file_object = managed_file_table.rows[output_file_id].file_object assert stored_file_object is not None assert stored_file_object.id == output_file_id - assert ( - managed_file_table.upsert_calls[0][1]["create"]["file_object"] - == stored_file_object.model_dump_json() - ) + assert managed_file_table.upsert_calls[0][1]["create"]["file_object"] == stored_file_object.model_dump_json() @pytest.mark.asyncio @@ -837,9 +777,9 @@ async def test_output_file_id_preserves_target_model_names_when_model_name_missi (e.g. Vertex batch retrieve), unified output_file_id should still include target_model_names from the managed input file ID. """ - import litellm.proxy.proxy_server as proxy_server_module from openai.types.batch import BatchRequestCounts + import litellm.proxy.proxy_server as proxy_server_module from litellm.proxy._types import UserAPIKeyAuth from litellm.types.llms.openai import OpenAIFileObject from litellm.types.utils import LiteLLMBatch @@ -865,11 +805,7 @@ async def test_output_file_id_preserves_target_model_names_when_model_name_missi "llm_output_file_id,file-input-1;" "llm_output_file_model_id,model-id-1" ) - managed_input_file_id = ( - base64.urlsafe_b64encode(managed_input_file_payload.encode()) - .decode() - .rstrip("=") - ) + managed_input_file_id = base64.urlsafe_b64encode(managed_input_file_payload.encode()).decode().rstrip("=") batch._hidden_params = { "model_id": "model-id-1", @@ -904,9 +840,7 @@ async def test_output_file_id_preserves_target_model_names_when_model_name_missi response=batch, ) - decoded_output_file_id = is_base64_encoded_unified_file_id( - cast(LiteLLMBatch, response).output_file_id - ) + decoded_output_file_id = is_base64_encoded_unified_file_id(cast(LiteLLMBatch, response).output_file_id) assert decoded_output_file_id assert "target_model_names,gemini-2.5-pro" in cast(str, decoded_output_file_id) mock_retrieve.assert_awaited_once() @@ -914,10 +848,7 @@ async def test_output_file_id_preserves_target_model_names_when_model_name_missi stored_file_object = managed_file_table.rows[output_file_id].file_object assert stored_file_object is not None assert stored_file_object.id == output_file_id - assert ( - managed_file_table.upsert_calls[0][1]["create"]["file_object"] - == stored_file_object.model_dump_json() - ) + assert managed_file_table.upsert_calls[0][1]["create"]["file_object"] == stored_file_object.model_dump_json() @pytest.mark.asyncio @@ -925,9 +856,9 @@ async def test_error_file_id_for_failed_batch(): """ Test that the error_file_id is properly managed when a batch fails """ - import litellm.proxy.proxy_server as proxy_server_module from openai.types.batch import BatchRequestCounts + import litellm.proxy.proxy_server as proxy_server_module from litellm.proxy._types import UserAPIKeyAuth from litellm.types.llms.openai import OpenAIFileObject from litellm.types.utils import LiteLLMBatch @@ -991,7 +922,6 @@ async def test_error_file_id_for_failed_batch(): return_value=error_file_object, ) as mock_retrieve, ): - user_api_key_dict = UserAPIKeyAuth(user_id="test-user-123") response = await proxy_managed_files.async_post_call_success_hook( @@ -1010,10 +940,7 @@ async def test_error_file_id_for_failed_batch(): stored_file_object = managed_file_table.rows[error_file_id].file_object assert stored_file_object is not None assert stored_file_object.id == error_file_id - assert ( - managed_file_table.upsert_calls[0][1]["create"]["file_object"] - == stored_file_object.model_dump_json() - ) + assert managed_file_table.upsert_calls[0][1]["create"]["file_object"] == stored_file_object.model_dump_json() @pytest.mark.asyncio @@ -1048,17 +975,13 @@ async def test_async_post_call_success_hook_twice_assert_no_unique_violation(): "unified_batch_id": "litellm_proxy;model_id:12345679;llm_batch_id:batch_685c5e5d63988190b85bdb2147ba131d", } - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=prisma_client - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=prisma_client) # first retrieve batch tasks = [] first_create_task = asyncio.create_task with patch("asyncio.create_task") as mock_create_task: - mock_create_task.side_effect = ( - lambda coro: tasks.append(first_create_task(coro)) or tasks[-1] - ) + mock_create_task.side_effect = lambda coro: tasks.append(first_create_task(coro)) or tasks[-1] response = await proxy_managed_files.async_post_call_success_hook( data={}, @@ -1080,9 +1003,7 @@ async def test_async_post_call_success_hook_twice_assert_no_unique_violation(): tasks = [] second_create_task = asyncio.create_task with patch("asyncio.create_task") as mock_create_task: - mock_create_task.side_effect = ( - lambda coro: tasks.append(second_create_task(coro)) or tasks[-1] - ) + mock_create_task.side_effect = lambda coro: tasks.append(second_create_task(coro)) or tasks[-1] await proxy_managed_files.async_post_call_success_hook( data={}, @@ -1134,10 +1055,7 @@ def test_update_responses_input_with_unified_file_id(): assert updated_input[0]["content"][0]["type"] == "input_file" assert updated_input[0]["content"][0]["file_id"] == "file-ECBPW7ML9g7XHdwGgUPZaM" assert updated_input[0]["content"][1]["type"] == "input_text" - assert ( - updated_input[0]["content"][1]["text"] - == "What is the first dragon in the book?" - ) + assert updated_input[0]["content"][1]["text"] == "What is the first dragon in the book?" def test_update_responses_input_with_regular_file_id(): @@ -1424,9 +1342,7 @@ def test_get_file_ids_from_responses_tools(): Test that get_file_ids_from_responses_tools correctly extracts file IDs from the tools parameter. """ - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=MagicMock() - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(InternalUsageCache(DualCache()), prisma_client=MagicMock()) tools = [ { @@ -1447,9 +1363,7 @@ def test_get_file_ids_from_responses_tools_multiple_tools(): """ Test that get_file_ids_from_responses_tools handles multiple tools. """ - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=MagicMock() - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(InternalUsageCache(DualCache()), prisma_client=MagicMock()) tools = [ { @@ -1481,9 +1395,7 @@ def test_get_file_ids_from_responses_tools_empty(): """ Test that get_file_ids_from_responses_tools handles empty or None tools. """ - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=MagicMock() - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(InternalUsageCache(DualCache()), prisma_client=MagicMock()) # Test with None file_ids = proxy_managed_files.get_file_ids_from_responses_tools(None) @@ -1534,9 +1446,7 @@ async def test_check_file_ids_access_with_unified_file_ids(): ) # Verify can_user_call_unified_file_id was called for the unified file ID - proxy_managed_files.can_user_call_unified_file_id.assert_called_once_with( - unified_file_id, user_api_key_dict - ) + proxy_managed_files.can_user_call_unified_file_id.assert_called_once_with(unified_file_id, user_api_key_dict) @pytest.mark.asyncio @@ -1745,9 +1655,7 @@ async def test_store_unified_file_id_with_none_file_object(): from litellm.proxy._types import UserAPIKeyAuth prisma_client = AsyncMock() - prisma_client.db.litellm_managedfiletable.upsert = AsyncMock( - return_value=MagicMock() - ) + prisma_client.db.litellm_managedfiletable.upsert = AsyncMock(return_value=MagicMock()) internal_usage_cache = MagicMock() internal_usage_cache.async_set_cache = AsyncMock() @@ -1780,11 +1688,8 @@ async def test_store_unified_file_id_updates_file_metadata_on_existing_row(): from litellm.types.llms.openai import OpenAIFileObject prisma_client = AsyncMock() - prisma_client.db.litellm_managedfiletable.upsert = AsyncMock( - return_value=MagicMock() - ) - internal_usage_cache = MagicMock() - internal_usage_cache.async_set_cache = AsyncMock() + prisma_client.db.litellm_managedfiletable.upsert = AsyncMock(return_value=MagicMock()) + internal_usage_cache = DualCache() proxy_managed_files = PROXY_LiteLLMManagedFiles( internal_usage_cache=internal_usage_cache, @@ -1822,12 +1727,8 @@ async def test_store_unified_file_id_updates_file_metadata_on_existing_row(): user_api_key_dict=user_api_key_dict, ) - first_update = prisma_client.db.litellm_managedfiletable.upsert.await_args_list[ - 0 - ].kwargs["data"]["update"] - second_update = prisma_client.db.litellm_managedfiletable.upsert.await_args_list[ - 1 - ].kwargs["data"]["update"] + first_update = prisma_client.db.litellm_managedfiletable.upsert.await_args_list[0].kwargs["data"]["update"] + second_update = prisma_client.db.litellm_managedfiletable.upsert.await_args_list[1].kwargs["data"]["update"] assert "file_object" not in first_update assert second_update["file_object"] == file_object.model_dump_json() assert second_update["storage_backend"] == "s3" @@ -1854,9 +1755,7 @@ async def test_store_batch_output_file_skips_existing_file_object(): flat_model_file_ids=["provider-output"], created_by="user-123", ) - proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma( - existing_file_row - ) + proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma(existing_file_row) with patch("litellm.afile_retrieve", new_callable=AsyncMock) as retrieve: await proxy_managed_files.store_batch_output_file( @@ -1868,9 +1767,7 @@ async def test_store_batch_output_file_skips_existing_file_object(): ) retrieve.assert_not_called() - assert managed_file_table.find_first_calls == [ - {"unified_file_id": "unified-output"} - ] + assert managed_file_table.find_first_calls == [{"unified_file_id": "unified-output"}] assert managed_file_table.upsert_calls == [] @@ -1895,9 +1792,7 @@ async def test_store_batch_output_file_stores_provider_object_with_unified_id(): async def record_sleep(delay: float) -> None: sleep_calls.append(delay) - proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma( - sleep=record_sleep - ) + proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma(sleep=record_sleep) with ( patch.object(proxy_server_module, "llm_router", router), @@ -1926,13 +1821,8 @@ async def test_store_batch_output_file_stores_provider_object_with_unified_id(): assert stored_object is not None assert stored_object.id == "unified-output" assert stored_object.bytes == 123 - assert ( - managed_file_table.upsert_calls[0][1]["create"]["file_object"] - == stored_object.model_dump_json() - ) - assert managed_file_table.rows["unified-output"].model_mappings == { - "model-123": "provider-output" - } + assert managed_file_table.upsert_calls[0][1]["create"]["file_object"] == stored_object.model_dump_json() + assert managed_file_table.rows["unified-output"].model_mappings == {"model-123": "provider-output"} @pytest.mark.asyncio @@ -1953,9 +1843,7 @@ async def test_store_batch_output_file_retries_model_name_provider_fetch(): async def record_sleep(delay: float) -> None: sleep_calls.append(delay) - proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma( - sleep=record_sleep - ) + proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma(sleep=record_sleep) with patch( "litellm.afile_retrieve", new_callable=AsyncMock, @@ -1998,9 +1886,7 @@ async def test_store_batch_output_file_marks_model_name_fallback_after_four_tran async def record_sleep(delay: float) -> None: sleep_calls.append(delay) - proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma( - sleep=record_sleep - ) + proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma(sleep=record_sleep) with patch( "litellm.afile_retrieve", new_callable=AsyncMock, @@ -2063,13 +1949,8 @@ async def test_store_batch_output_file_falls_back_when_provider_retrieve_raises( assert stored_object.filename == "output.jsonl" assert stored_object.bytes == 321 assert stored_object.purpose == "batch_output" - assert ( - managed_file_table.upsert_calls[0][1]["create"]["file_object"] - == stored_object.model_dump_json() - ) - assert managed_file_table.rows["unified-output"].model_mappings == { - "model-123": "s3://bucket/output.jsonl" - } + assert managed_file_table.upsert_calls[0][1]["create"]["file_object"] == stored_object.model_dump_json() + assert managed_file_table.rows["unified-output"].model_mappings == {"model-123": "s3://bucket/output.jsonl"} @pytest.mark.asyncio @@ -2163,10 +2044,7 @@ async def test_store_batch_output_file_falls_back_without_model_id(): assert stored_object.filename == "error.jsonl" assert stored_object.bytes == 0 assert stored_object.purpose == "batch_output" - assert ( - managed_file_table.upsert_calls[0][1]["create"]["file_object"] - == stored_object.model_dump_json() - ) + assert managed_file_table.upsert_calls[0][1]["create"]["file_object"] == stored_object.model_dump_json() assert managed_file_table.rows["unified-error"].model_mappings == {} assert stored_object.litellm_details_fallback is None assert "litellm_details_fallback" not in stored_object.model_dump() @@ -2182,9 +2060,7 @@ async def test_listed_batch_saves_basic_output_without_provider_fetch(): patch.object(proxy_server_module, "llm_router", None), patch("litellm.afile_retrieve", new_callable=AsyncMock) as retrieve, ): - resolved: Final = await _resolve_batch_for_output_listing( - proxy_managed_files, response - ) + resolved: Final = await _resolve_batch_for_output_listing(proxy_managed_files, response) assert resolved is response retrieve.assert_not_called() @@ -2206,18 +2082,14 @@ async def test_listed_batch_leaves_existing_marked_output_untouched(): model_id="model-123", model_name="bedrock/model-x", ) - row: Final = _marked_fallback_file_row().model_copy( - update={"unified_file_id": unified_file_id} - ) + row: Final = _marked_fallback_file_row().model_copy(update={"unified_file_id": unified_file_id}) proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma(row) response: Final = _batch_response_for_output_listing() with ( patch.object(proxy_server_module, "llm_router", None), patch("litellm.afile_retrieve", new_callable=AsyncMock) as retrieve, ): - resolved: Final = await _resolve_batch_for_output_listing( - proxy_managed_files, response - ) + resolved: Final = await _resolve_batch_for_output_listing(proxy_managed_files, response) await proxy_managed_files.store_batch_output_file( unified_file_id=unified_file_id, provider_file_id="provider-output", @@ -2325,9 +2197,7 @@ async def test_batch_file_detail_retry_classification( async def record_sleep(delay: float) -> None: sleep_calls.append(delay) - proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma( - sleep=record_sleep - ) + proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma(sleep=record_sleep) router: Final = MagicMock() router.registry.return_value.credentials.return_value = {"api_key": "key"} provider_file_object: Final = OpenAIFileObject( @@ -2339,11 +2209,7 @@ async def test_batch_file_detail_retry_classification( purpose="batch_output", status="processed", ) - retrieve_side_effect: Final = ( - [error_factory(), provider_file_object] - if expected_attempts == 2 - else error_factory() - ) + retrieve_side_effect: Final = [error_factory(), provider_file_object] if expected_attempts == 2 else error_factory() with ( patch.object(proxy_server_module, "llm_router", router), patch( @@ -2361,9 +2227,7 @@ async def test_batch_file_detail_retry_classification( ) assert retrieve.await_count == expected_attempts - assert [call.kwargs["max_retries"] for call in retrieve.call_args_list] == [ - 0 - ] * expected_attempts + assert [call.kwargs["max_retries"] for call in retrieve.call_args_list] == [0] * expected_attempts assert sleep_calls == ([0.5] if expected_attempts == 2 else []) stored_object: Final = managed_file_table.rows["unified-output"].file_object assert stored_object is not None @@ -2391,9 +2255,7 @@ async def test_store_batch_output_file_ignores_unmarked_row_with_provider_route( flat_model_file_ids=["provider-output"], created_by="user-123", ) - proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma( - existing_file_row - ) + proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma(existing_file_row) router: Final = MagicMock() router.registry.return_value.credentials.return_value = {"api_key": "key"} with ( @@ -2665,9 +2527,7 @@ class _RefreshLandsBeforeTheFirstWriteTable(_InMemoryManagedFileTable): async def update_many(self, where: Mapping[str, object], data: Mapping[str, object]) -> int: if self.pending_refreshes: refreshed_row: Final = self.rows["unified-output"] - self.rows["unified-output"] = refreshed_row.model_copy( - update={"file_object": self.pending_refreshes.pop()} - ) + self.rows["unified-output"] = refreshed_row.model_copy(update={"file_object": self.pending_refreshes.pop()}) return await super().update_many(where, data) @@ -2936,9 +2796,7 @@ _NON_XAI_BATCH_OUTPUTS: Final = pytest.mark.parametrize( ) -def _non_xai_batch_output_row( - provider_file_id: str, *, size_bytes: int, fallback: bool -) -> LiteLLM_ManagedFileTable: +def _non_xai_batch_output_row(provider_file_id: str, *, size_bytes: int, fallback: bool) -> LiteLLM_ManagedFileTable: return _batch_output_row("xai", size_bytes=size_bytes, fallback=fallback).model_copy( update={"model_mappings": {"model-123": provider_file_id}, "flat_model_file_ids": [provider_file_id]} ) @@ -3117,9 +2975,7 @@ async def test_store_batch_output_file_marks_fallback_after_four_transient_failu async def record_sleep(delay: float) -> None: sleep_calls.append(delay) - proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma( - sleep=record_sleep - ) + proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma(sleep=record_sleep) with ( patch.object(proxy_server_module, "llm_router", router), patch( @@ -3177,9 +3033,7 @@ async def test_store_batch_output_file_skips_lookup_for_fallback_written_moments assert stored_object is not None assert (stored_object.bytes, stored_object.litellm_details_fallback) == (836, True) assert managed_file_table.upsert_calls == [] - assert managed_file_table.update_many_calls[0][0] == { - "unified_file_id": "unified-output" - } + assert managed_file_table.update_many_calls[0][0] == {"unified_file_id": "unified-output"} assert set(managed_file_table.update_many_calls[0][1]) == {"file_object"} @@ -3195,9 +3049,7 @@ async def test_store_batch_output_file_does_not_retry_not_found(): async def record_sleep(delay: float) -> None: sleep_calls.append(delay) - proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma( - sleep=record_sleep - ) + proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma(sleep=record_sleep) with ( patch.object(proxy_server_module, "llm_router", router), patch( @@ -3278,9 +3130,7 @@ async def test_store_batch_output_file_does_not_recreate_deleted_row_after_refre purpose="batch_output", status="processed", ) - proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma( - _marked_fallback_file_row() - ) + proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma(_marked_fallback_file_row()) router: Final = MagicMock() router.registry.return_value.credentials.return_value = {"api_key": "key"} @@ -3307,9 +3157,7 @@ async def test_store_batch_output_file_does_not_recreate_deleted_row_after_refre assert "unified-output" not in managed_file_table.rows assert managed_file_table.upsert_calls == [] assert len(managed_file_table.update_many_calls) == 1 - saved_file_object_json: Final = cast( - str, managed_file_table.update_many_calls[0][1]["file_object"] - ) + saved_file_object_json: Final = cast(str, managed_file_table.update_many_calls[0][1]["file_object"]) saved_file_object: Final = json.loads(saved_file_object_json) assert saved_file_object["id"] == "unified-output" assert ( @@ -3328,9 +3176,7 @@ async def test_store_batch_output_file_keeps_marked_fallback_when_retryable_deta router = MagicMock() router.registry.return_value.credentials.return_value = {"api_key": "key"} - proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma( - _marked_fallback_file_row() - ) + proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma(_marked_fallback_file_row()) with ( patch.object(proxy_server_module, "llm_router", router), patch( @@ -3359,9 +3205,7 @@ async def test_store_batch_output_file_updates_only_size_for_failed_marked_fallb router = MagicMock() router.registry.return_value.credentials.return_value = {"api_key": "key"} - proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma( - _marked_fallback_file_row() - ) + proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma(_marked_fallback_file_row()) with ( patch.object(proxy_server_module, "llm_router", router), patch( @@ -3399,23 +3243,21 @@ async def test_afile_delete_returns_provider_response_when_stored_file_object_no prisma_client = AsyncMock() db_record = MagicMock() db_record.model_mappings = '{"model-123": "file-provider-xyz"}' - prisma_client.db.litellm_managedfiletable.find_first = AsyncMock( - return_value=db_record - ) + prisma_client.db.litellm_managedfiletable.find_first = AsyncMock(return_value=db_record) prisma_client.db.litellm_managedfiletable.delete = AsyncMock() - internal_usage_cache = MagicMock() - internal_usage_cache.async_get_cache = AsyncMock( - return_value={ + internal_usage_cache = DualCache() + await internal_usage_cache.async_set_cache( + key=unified_file_id, + value={ "unified_file_id": unified_file_id, "model_mappings": {"model-123": "file-provider-xyz"}, "flat_model_file_ids": ["file-provider-xyz"], "file_object": None, "created_by": "test-user", "updated_by": "test-user", - } + }, ) - internal_usage_cache.async_set_cache = AsyncMock() proxy_managed_files = PROXY_LiteLLMManagedFiles( internal_usage_cache=internal_usage_cache, @@ -3469,9 +3311,7 @@ async def test_afile_retrieve_fetches_from_provider_when_file_object_none(): async def record_sleep(delay: float) -> None: sleep_calls.append(delay) - proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma( - stored_file, sleep=record_sleep - ) + proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma(stored_file, sleep=record_sleep) # Mock the router and provider response provider_file_response = OpenAIFileObject( @@ -3531,7 +3371,7 @@ async def test_afile_retrieve_raises_error_when_no_router_and_file_object_none() unified_file_id = "test-unified-file-id" - with pytest.raises(Exception, match='LiteLLM Managed File object with id=test-unified-file-id') as exc_info: + with pytest.raises(Exception, match="LiteLLM Managed File object with id=test-unified-file-id") as exc_info: await proxy_managed_files.afile_retrieve( file_id=unified_file_id, litellm_parent_otel_span=None, @@ -3610,9 +3450,7 @@ async def test_afile_retrieve_refreshes_marked_fallback_and_preserves_ownership( assert updated_row.file_object == response assert updated_row.created_by == "user-123" assert updated_row.team_id == "team-123" - assert managed_file_table.update_many_calls[0][1] == { - "file_object": response.model_dump_json() - } + assert managed_file_table.update_many_calls[0][1] == {"file_object": response.model_dump_json()} assert managed_file_table.upsert_calls == [] cached_row = await proxy_managed_files.internal_usage_cache.async_get_cache( key="unified-output", @@ -3693,9 +3531,7 @@ async def test_afile_retrieve_uses_model_name_when_router_does_not_know_deployme assert response.bytes == 836 retrieve.assert_awaited_once() - router.registry.return_value.credentials.assert_called_once_with( - "bedrock/model-x", form="with_provider" - ) + router.registry.return_value.credentials.assert_called_once_with("bedrock/model-x", form="with_provider") assert retrieve.await_args.kwargs["custom_llm_provider"] == "bedrock" assert retrieve.await_args.kwargs["max_retries"] == 0 assert "_litellm_internal_model_credentials" not in retrieve.await_args.kwargs @@ -3717,13 +3553,9 @@ async def test_afile_retrieve_does_not_recreate_deleted_row_after_refresh(): router.registry.return_value.credentials.return_value = {"api_key": "key"} proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma(row) - async def delete_row_then_return_file( - *, file_id: str, **_options: object - ) -> OpenAIFileObject: + async def delete_row_then_return_file(*, file_id: str, **_options: object) -> OpenAIFileObject: assert file_id == "provider-output" - deleted_row: Final = await managed_file_table.delete( - where={"unified_file_id": "unified-output"} - ) + deleted_row: Final = await managed_file_table.delete(where={"unified_file_id": "unified-output"}) assert deleted_row is not None return provider_object @@ -3782,10 +3614,7 @@ async def test_afile_retrieve_case3_includes_provider_error_text(): ) assert type(error.value) is Exception - assert str(error.value) == ( - "Failed to retrieve file unified-output from provider: " - "404: provider file is missing" - ) + assert str(error.value) == ("Failed to retrieve file unified-output from provider: 404: provider file is missing") @pytest.mark.asyncio @@ -3902,9 +3731,7 @@ async def test_afile_retrieve_case3_without_route_raises_accurate_error(): async def test_afile_retrieve_returns_marked_fallback_without_marker_when_refresh_fails(): router = MagicMock() router.registry.return_value.credentials.return_value = {"api_key": "key"} - proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma( - _marked_fallback_file_row() - ) + proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma(_marked_fallback_file_row()) with patch( "litellm.afile_retrieve", @@ -3932,7 +3759,7 @@ async def test_afile_retrieve_raises_error_for_non_managed_file(): """ proxy_managed_files, _ = _managed_files_with_fake_prisma() - with pytest.raises(Exception, match='LiteLLM Managed File object with id=non-existent-file-id') as exc_info: + with pytest.raises(Exception, match="LiteLLM Managed File object with id=non-existent-file-id") as exc_info: await proxy_managed_files.afile_retrieve( file_id="non-existent-file-id", litellm_parent_otel_span=None, @@ -3943,7 +3770,6 @@ async def test_afile_retrieve_raises_error_for_non_managed_file(): @pytest.mark.asyncio async def test_list_batches_from_managed_objects_table(): - from openai.types.batch import BatchRequestCounts from litellm.proxy._types import UserAPIKeyAuth @@ -3984,9 +3810,7 @@ async def test_list_batches_from_managed_objects_table(): batch_record_2, ] - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=prisma_client - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=prisma_client) result = await proxy_managed_files.list_user_batches( user_api_key_dict=UserAPIKeyAuth(user_id="test-user"), @@ -4015,9 +3839,7 @@ async def test_list_batches_from_managed_objects_table_empty_list(): prisma_client = AsyncMock() prisma_client.db.litellm_managedobjecttable.find_many.return_value = [] - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=prisma_client - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=prisma_client) result = await proxy_managed_files.list_user_batches( user_api_key_dict=UserAPIKeyAuth(user_id="test-user"), @@ -4094,23 +3916,17 @@ async def test_list_batches_registers_and_returns_unified_output_file_ids(): prisma_client = AsyncMock() prisma_client.db.litellm_managedobjecttable.find_many.return_value = [ - _terminal_batch_record( - unified_batch_uid, raw_input_file_id, raw_output_file_id, raw_error_file_id - ) + _terminal_batch_record(unified_batch_uid, raw_input_file_id, raw_output_file_id, raw_error_file_id) ] input_file_row = MagicMock() input_file_row.unified_file_id = unified_input_file_id input_file_row.flat_model_file_ids = [raw_input_file_id] - prisma_client.db.litellm_managedfiletable.find_many = AsyncMock( - return_value=[input_file_row] - ) + prisma_client.db.litellm_managedfiletable.find_many = AsyncMock(return_value=[input_file_row]) prisma_client.db.litellm_managedfiletable.find_first = AsyncMock(return_value=None) - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=prisma_client - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=prisma_client) result = await proxy_managed_files.list_user_batches( user_api_key_dict=UserAPIKeyAuth(user_id="owner-user"), @@ -4138,9 +3954,7 @@ async def test_list_batches_registers_and_returns_unified_output_file_ids(): assert f"llm_output_file_id,{raw_error_file_id}" in decoded_error upsert_calls = prisma_client.db.litellm_managedfiletable.upsert.await_args_list - stored_raw_ids = { - c.kwargs["data"]["create"]["flat_model_file_ids"][0] for c in upsert_calls - } + stored_raw_ids = {c.kwargs["data"]["create"]["flat_model_file_ids"][0] for c in upsert_calls} assert stored_raw_ids == {raw_output_file_id, raw_error_file_id} for c in upsert_calls: assert c.kwargs["data"]["create"]["created_by"] == "owner-user" @@ -4151,9 +3965,7 @@ async def test_list_batches_registers_and_returns_unified_output_file_ids(): async def test_afile_list_returns_persisted_batch_output_file(): from litellm.proxy._types import UserAPIKeyAuth - proxy_managed_files, _ = _managed_files_with_fake_prisma( - _marked_fallback_file_row() - ) + proxy_managed_files, _ = _managed_files_with_fake_prisma(_marked_fallback_file_row()) result = await proxy_managed_files.afile_list( purpose="batch_output", @@ -4174,9 +3986,7 @@ async def test_afile_list_returns_persisted_batch_output_file(): async def test_get_user_created_file_ids_strips_fallback_marker(): from litellm.proxy._types import UserAPIKeyAuth - proxy_managed_files, _ = _managed_files_with_fake_prisma( - _marked_fallback_file_row() - ) + proxy_managed_files, _ = _managed_files_with_fake_prisma(_marked_fallback_file_row()) files = await proxy_managed_files.get_user_created_file_ids( UserAPIKeyAuth(user_id="user-123"), @@ -4224,14 +4034,10 @@ async def test_list_batches_resolves_existing_managed_rows_without_minting(): for raw_id, unified_id in zip(raw_output_file_ids, existing_unified_output_ids) ] - prisma_client.db.litellm_managedfiletable.find_many = AsyncMock( - return_value=existing_rows - ) + prisma_client.db.litellm_managedfiletable.find_many = AsyncMock(return_value=existing_rows) prisma_client.db.litellm_managedfiletable.find_first = AsyncMock() - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=prisma_client - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=prisma_client) result = await proxy_managed_files.list_user_batches( user_api_key_dict=UserAPIKeyAuth(user_id="owner-user"), @@ -4254,19 +4060,14 @@ async def test_list_batches_caps_page_size_at_100(): prisma_client = AsyncMock() prisma_client.db.litellm_managedobjecttable.find_many.return_value = [] - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=prisma_client - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=prisma_client) result = await proxy_managed_files.list_user_batches( user_api_key_dict=UserAPIKeyAuth(user_id="owner-user"), limit=100000, ) - assert ( - prisma_client.db.litellm_managedobjecttable.find_many.await_args.kwargs["take"] - == 101 - ) + assert prisma_client.db.litellm_managedobjecttable.find_many.await_args.kwargs["take"] == 101 assert result["data"] == [] @@ -4276,9 +4077,7 @@ async def test_list_batches_from_managed_objects_table_provider_filter_raises_ex prisma_client = AsyncMock() - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=prisma_client - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=prisma_client) # Filtering by provider should raise Exception with pytest.raises(Exception, match="Filtering by 'provider' is not supported when using managed") as exc_info: @@ -4288,9 +4087,7 @@ async def test_list_batches_from_managed_objects_table_provider_filter_raises_ex provider="openai", ) - assert str(exc_info.value) == ( - "Filtering by 'provider' is not supported when using managed batches." - ) + assert str(exc_info.value) == ("Filtering by 'provider' is not supported when using managed batches.") # Verify find_many was NOT called since exception is raised before database query prisma_client.db.litellm_managedobjecttable.find_many.assert_not_called() @@ -4302,9 +4099,7 @@ async def test_list_batches_from_managed_objects_table_target_model_name_filter_ prisma_client = AsyncMock() - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=prisma_client - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=prisma_client) # Filtering by provider should raise Exception with pytest.raises(Exception, match="Filtering by 'target_model_names' is not supported when") as exc_info: @@ -4314,9 +4109,7 @@ async def test_list_batches_from_managed_objects_table_target_model_name_filter_ target_model_names="gpt-5.5,gpt-3.5", ) - assert str(exc_info.value) == ( - "Filtering by 'target_model_names' is not supported when using managed batches." - ) + assert str(exc_info.value) == ("Filtering by 'target_model_names' is not supported when using managed batches.") # Verify find_many was NOT called since exception is raised before database query prisma_client.db.litellm_managedobjecttable.find_many.assert_not_called() @@ -4360,9 +4153,7 @@ async def test_list_batches_from_managed_objects_table_filters_by_created_by(): } ) - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=prisma_client - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=prisma_client) # Query with user1's API key - should only return user1's batch prisma_client.db.litellm_managedobjecttable.find_many.return_value = [batch_user1] @@ -4410,9 +4201,7 @@ async def test_list_batches_pagination_uses_unified_object_id_cursor(): prisma_client.db.litellm_managedobjecttable.find_first.return_value = MagicMock() prisma_client.db.litellm_managedobjecttable.find_many.return_value = [] - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=prisma_client - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=prisma_client) await proxy_managed_files.list_user_batches( user_api_key_dict=UserAPIKeyAuth(user_id="test-user"), @@ -4478,19 +4267,17 @@ async def test_list_batches_pagination_walks_all_pages_without_loops_or_gaps(): result = [r for r in result if r.id > id_filter["gt"]] order_keys = order if isinstance(order, list) else [order] for clause in reversed(order_keys): - (order_field, direction), = clause.items() - result.sort( - key=lambda r: getattr(r, order_field), reverse=(direction == "desc") - ) + ((order_field, direction),) = clause.items() + result.sort(key=lambda r: getattr(r, order_field), reverse=(direction == "desc")) if cursor is not None: - (cur_field, cur_val), = cursor.items() + ((cur_field, cur_val),) = cursor.items() idx = next( (i for i, r in enumerate(result) if getattr(r, cur_field) == cur_val), None, ) if idx is None: return [] - result = result[idx + skip:] + result = result[idx + skip :] return result[:take] async def fake_find_first(where): @@ -4500,24 +4287,16 @@ async def test_list_batches_pagination_walks_all_pages_without_loops_or_gaps(): ) prisma_client = AsyncMock() - prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - side_effect=fake_find_many - ) - prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock( - side_effect=fake_find_first - ) + prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(side_effect=fake_find_many) + prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock(side_effect=fake_find_first) - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=prisma_client - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=prisma_client) user = UserAPIKeyAuth(user_id="test-user") seen: list = [] after = None for _ in range(total + 5): - resp = await proxy_managed_files.list_user_batches( - user_api_key_dict=user, limit=3, after=after - ) + resp = await proxy_managed_files.list_user_batches(user_api_key_dict=user, limit=3, after=after) page_ids = [b.id for b in resp["data"]] seen.extend(page_ids) if not resp["has_more"]: @@ -4582,9 +4361,7 @@ async def test_list_batches_pagination_stable_when_created_at_ties(): result = list(rows) for clause in reversed(order_keys): field, direction = next(iter(clause.items())) - result.sort( - key=lambda r: getattr(r, field), reverse=(direction == "desc") - ) + result.sort(key=lambda r: getattr(r, field), reverse=(direction == "desc")) def order_key(r): return tuple(getattr(r, f) for f in fields) @@ -4611,7 +4388,7 @@ async def test_list_batches_pagination_stable_when_created_at_ties(): ) if idx is None: return [] - result = result[idx + skip:] + result = result[idx + skip :] return result[:take] async def fake_find_first(where): @@ -4621,24 +4398,16 @@ async def test_list_batches_pagination_stable_when_created_at_ties(): ) prisma_client = AsyncMock() - prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - side_effect=fake_find_many - ) - prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock( - side_effect=fake_find_first - ) + prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(side_effect=fake_find_many) + prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock(side_effect=fake_find_first) - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=prisma_client - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=prisma_client) user = UserAPIKeyAuth(user_id="test-user") seen: list = [] after = None for _ in range(total + 5): - resp = await proxy_managed_files.list_user_batches( - user_api_key_dict=user, limit=2, after=after - ) + resp = await proxy_managed_files.list_user_batches(user_api_key_dict=user, limit=2, after=after) page_ids = [b.id for b in resp["data"]] seen.extend(page_ids) if not resp["has_more"]: @@ -4678,18 +4447,16 @@ def _managed_batch_row(index, file_object=None): def _fake_managed_object_table(rows): async def find_many(where, take, order, cursor=None, skip=0): - result = sorted( - rows, key=lambda r: (r.created_at, r.unified_object_id), reverse=True - ) + result = sorted(rows, key=lambda r: (r.created_at, r.unified_object_id), reverse=True) if cursor is not None: - (cur_field, cur_val), = cursor.items() + ((cur_field, cur_val),) = cursor.items() idx = next( (i for i, r in enumerate(result) if getattr(r, cur_field) == cur_val), None, ) if idx is None: return [] - result = result[idx + skip:] + result = result[idx + skip :] return result[:take] async def find_first(where): @@ -4699,12 +4466,8 @@ def _fake_managed_object_table(rows): ) prisma_client = AsyncMock() - prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - side_effect=find_many - ) - prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock( - side_effect=find_first - ) + prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(side_effect=find_many) + prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock(side_effect=find_first) return prisma_client @@ -4712,9 +4475,7 @@ async def _walk_batch_pages(proxy_managed_files, user, limit, max_pages=20): pages = [] after = None for _ in range(max_pages): - resp = await proxy_managed_files.list_user_batches( - user_api_key_dict=user, limit=limit, after=after - ) + resp = await proxy_managed_files.list_user_batches(user_api_key_dict=user, limit=limit, after=after) pages.append(resp) if not resp["has_more"]: break @@ -4738,14 +4499,10 @@ async def test_list_batches_rejects_unknown_after_cursor(): from litellm.proxy._types import UserAPIKeyAuth prisma_client = AsyncMock() - prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock( - return_value=None - ) + prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=None) prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[]) - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=prisma_client - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=prisma_client) with pytest.raises(HTTPException) as exc_info: await proxy_managed_files.list_user_batches( @@ -4778,9 +4535,7 @@ async def test_list_batches_treats_empty_after_as_no_cursor(): rows = [_managed_batch_row(i) for i in range(2)] prisma_client = _fake_managed_object_table(rows) - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=prisma_client - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=prisma_client) page = await proxy_managed_files.list_user_batches( user_api_key_dict=UserAPIKeyAuth(user_id="test-user"), limit=2, after="" @@ -4815,14 +4570,10 @@ async def test_list_batches_rejects_after_cursor_owned_by_another_user(): return other_users_batch prisma_client = AsyncMock() - prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock( - side_effect=find_first - ) + prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock(side_effect=find_first) prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[]) - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=prisma_client - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=prisma_client) with pytest.raises(HTTPException) as exc_info: await proxy_managed_files.list_user_batches( @@ -4848,13 +4599,9 @@ async def test_list_batches_has_more_false_on_exactly_full_final_page(): rows = [_managed_batch_row(i) for i in range(4)] prisma_client = _fake_managed_object_table(rows) - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=prisma_client - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=prisma_client) - pages = await _walk_batch_pages( - proxy_managed_files, UserAPIKeyAuth(user_id="test-user"), limit=2 - ) + pages = await _walk_batch_pages(proxy_managed_files, UserAPIKeyAuth(user_id="test-user"), limit=2) seen = [batch.id for page in pages for batch in page["data"]] assert seen == [r.unified_object_id for r in reversed(rows)] @@ -4876,13 +4623,9 @@ async def test_list_batches_unparseable_row_does_not_truncate_pagination(): rows[2].file_object = "{ not valid json" prisma_client = _fake_managed_object_table(rows) - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=prisma_client - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=prisma_client) - pages = await _walk_batch_pages( - proxy_managed_files, UserAPIKeyAuth(user_id="test-user"), limit=2 - ) + pages = await _walk_batch_pages(proxy_managed_files, UserAPIKeyAuth(user_id="test-user"), limit=2) seen = [batch.id for page in pages for batch in page["data"]] assert seen == [rows[3].unified_object_id, rows[1].unified_object_id, rows[0].unified_object_id] @@ -4905,13 +4648,9 @@ async def test_list_batches_fills_a_page_past_a_full_page_of_unparseable_rows(): corrupt_row.file_object = "{ not valid json" prisma_client = _fake_managed_object_table(rows) - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=prisma_client - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=prisma_client) - pages = await _walk_batch_pages( - proxy_managed_files, UserAPIKeyAuth(user_id="test-user"), limit=1 - ) + pages = await _walk_batch_pages(proxy_managed_files, UserAPIKeyAuth(user_id="test-user"), limit=1) assert [[batch.id for batch in page["data"]] for page in pages] == [ [rows[4].unified_object_id], @@ -4931,25 +4670,17 @@ async def test_list_batches_bounds_the_queries_a_deep_unparseable_run_costs(): from litellm.proxy._types import UserAPIKeyAuth rows = [_managed_batch_row(0)] + [ - _managed_batch_row(index, file_object="{ not valid json") - for index in range(1, _DEEP_BATCH_SCAN_ROW_COUNT + 1) + _managed_batch_row(index, file_object="{ not valid json") for index in range(1, _DEEP_BATCH_SCAN_ROW_COUNT + 1) ] prisma_client = _fake_managed_object_table(rows) - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=prisma_client - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=prisma_client) - page = await proxy_managed_files.list_user_batches( - user_api_key_dict=UserAPIKeyAuth(user_id="test-user"), limit=1 - ) + page = await proxy_managed_files.list_user_batches(user_api_key_dict=UserAPIKeyAuth(user_id="test-user"), limit=1) assert [batch.id for batch in page["data"]] == [rows[0].unified_object_id] assert page["has_more"] is False - assert ( - prisma_client.db.litellm_managedobjecttable.find_many.call_count - <= _DEEP_BATCH_SCAN_QUERY_BUDGET - ) + assert prisma_client.db.litellm_managedobjecttable.find_many.call_count <= _DEEP_BATCH_SCAN_QUERY_BUDGET @pytest.mark.asyncio @@ -4960,13 +4691,9 @@ async def test_list_batches_reads_one_chunk_when_the_first_one_fills_the_page(): rows = [_managed_batch_row(index) for index in range(_DEEP_BATCH_SCAN_ROW_COUNT)] prisma_client = _fake_managed_object_table(rows) - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=prisma_client - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=prisma_client) - page = await proxy_managed_files.list_user_batches( - user_api_key_dict=UserAPIKeyAuth(user_id="test-user"), limit=2 - ) + page = await proxy_managed_files.list_user_batches(user_api_key_dict=UserAPIKeyAuth(user_id="test-user"), limit=2) assert [batch.id for batch in page["data"]] == [ rows[-1].unified_object_id, @@ -5041,20 +4768,14 @@ async def test_user_b_cannot_retrieve_user_a_batch(): batch_record.created_by = "user_a_id" prisma_client.db.litellm_managedobjecttable.find_first.return_value = batch_record - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=prisma_client - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=prisma_client) # User B tries to retrieve User A's batch - unified_batch_id = ( - "bGl0ZWxsbV9wcm94eTttb2RlbF9pZDpteS1tb2RlbDtsbG1fYmF0Y2hfaWQ6YmF0Y2hfYWJjMTIz" - ) + unified_batch_id = "bGl0ZWxsbV9wcm94eTttb2RlbF9pZDpteS1tb2RlbDtsbG1fYmF0Y2hfaWQ6YmF0Y2hfYWJjMTIz" with pytest.raises(HTTPException) as exc_info: await proxy_managed_files.async_pre_call_hook( - user_api_key_dict=UserAPIKeyAuth( - user_id="user_b_id", parent_otel_span=MagicMock() - ), + user_api_key_dict=UserAPIKeyAuth(user_id="user_b_id", parent_otel_span=MagicMock()), cache=MagicMock(), data={"batch_id": unified_batch_id}, call_type="aretrieve_batch", @@ -5078,20 +4799,14 @@ async def test_user_b_cannot_cancel_user_a_batch(): batch_record.created_by = "user_a_id" prisma_client.db.litellm_managedobjecttable.find_first.return_value = batch_record - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=prisma_client - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=prisma_client) # User B tries to cancel User A's batch - unified_batch_id = ( - "bGl0ZWxsbV9wcm94eTttb2RlbF9pZDpteS1tb2RlbDtsbG1fYmF0Y2hfaWQ6YmF0Y2hfYWJjMTIz" - ) + unified_batch_id = "bGl0ZWxsbV9wcm94eTttb2RlbF9pZDpteS1tb2RlbDtsbG1fYmF0Y2hfaWQ6YmF0Y2hfYWJjMTIz" with pytest.raises(HTTPException) as exc_info: await proxy_managed_files.async_pre_call_hook( - user_api_key_dict=UserAPIKeyAuth( - user_id="user_b_id", parent_otel_span=MagicMock() - ), + user_api_key_dict=UserAPIKeyAuth(user_id="user_b_id", parent_otel_span=MagicMock()), cache=MagicMock(), data={"batch_id": unified_batch_id}, call_type="acancel_batch", @@ -5118,20 +4833,14 @@ async def test_user_a_can_retrieve_own_batch(): batch_record.created_by = "user_a_id" prisma_client.db.litellm_managedobjecttable.find_first.return_value = batch_record - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=prisma_client - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=prisma_client) # User A retrieves their own batch - unified_batch_id = ( - "bGl0ZWxsbV9wcm94eTttb2RlbF9pZDpteS1tb2RlbDtsbG1fYmF0Y2hfaWQ6YmF0Y2hfYWJjMTIz" - ) + unified_batch_id = "bGl0ZWxsbV9wcm94eTttb2RlbF9pZDpteS1tb2RlbDtsbG1fYmF0Y2hfaWQ6YmF0Y2hfYWJjMTIz" # Should not raise an exception result = await proxy_managed_files.async_pre_call_hook( - user_api_key_dict=UserAPIKeyAuth( - user_id="user_a_id", parent_otel_span=MagicMock() - ), + user_api_key_dict=UserAPIKeyAuth(user_id="user_a_id", parent_otel_span=MagicMock()), cache=MagicMock(), data={"batch_id": unified_batch_id}, call_type="aretrieve_batch", @@ -5156,20 +4865,14 @@ async def test_user_b_cannot_retrieve_user_a_file(): file_record.created_by = "user_a_id" prisma_client.db.litellm_managedfiletable.find_first.return_value = file_record - proxy_managed_files = PROXY_LiteLLMManagedFiles( - MagicMock(), prisma_client=prisma_client - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(MagicMock(), prisma_client=prisma_client) # User B tries to retrieve User A's file - unified_file_id = ( - "bGl0ZWxsbV9wcm94eTphcHBsaWNhdGlvbi9qc29uO3VuaWZpZWRfaWQsZmlsZS1hYmMxMjM" - ) + unified_file_id = "bGl0ZWxsbV9wcm94eTphcHBsaWNhdGlvbi9qc29uO3VuaWZpZWRfaWQsZmlsZS1hYmMxMjM" with pytest.raises(HTTPException) as exc_info: await proxy_managed_files.async_pre_call_hook( - user_api_key_dict=UserAPIKeyAuth( - user_id="user_b_id", parent_otel_span=MagicMock() - ), + user_api_key_dict=UserAPIKeyAuth(user_id="user_b_id", parent_otel_span=MagicMock()), cache=MagicMock(), data={"file_id": unified_file_id}, call_type="afile_retrieve", @@ -5193,20 +4896,14 @@ async def test_user_b_cannot_download_user_a_file_content(): file_record.created_by = "user_a_id" prisma_client.db.litellm_managedfiletable.find_first.return_value = file_record - proxy_managed_files = PROXY_LiteLLMManagedFiles( - MagicMock(), prisma_client=prisma_client - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(MagicMock(), prisma_client=prisma_client) # User B tries to download User A's file content - unified_file_id = ( - "bGl0ZWxsbV9wcm94eTphcHBsaWNhdGlvbi9qc29uO3VuaWZpZWRfaWQsZmlsZS1hYmMxMjM" - ) + unified_file_id = "bGl0ZWxsbV9wcm94eTphcHBsaWNhdGlvbi9qc29uO3VuaWZpZWRfaWQsZmlsZS1hYmMxMjM" with pytest.raises(HTTPException) as exc_info: await proxy_managed_files.async_pre_call_hook( - user_api_key_dict=UserAPIKeyAuth( - user_id="user_b_id", parent_otel_span=MagicMock() - ), + user_api_key_dict=UserAPIKeyAuth(user_id="user_b_id", parent_otel_span=MagicMock()), cache=MagicMock(), data={"file_id": unified_file_id}, call_type="afile_content", @@ -5230,20 +4927,14 @@ async def test_user_b_cannot_delete_user_a_file(): file_record.created_by = "user_a_id" prisma_client.db.litellm_managedfiletable.find_first.return_value = file_record - proxy_managed_files = PROXY_LiteLLMManagedFiles( - MagicMock(), prisma_client=prisma_client - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(MagicMock(), prisma_client=prisma_client) # User B tries to delete User A's file - unified_file_id = ( - "bGl0ZWxsbV9wcm94eTphcHBsaWNhdGlvbi9qc29uO3VuaWZpZWRfaWQsZmlsZS1hYmMxMjM" - ) + unified_file_id = "bGl0ZWxsbV9wcm94eTphcHBsaWNhdGlvbi9qc29uO3VuaWZpZWRfaWQsZmlsZS1hYmMxMjM" with pytest.raises(HTTPException) as exc_info: await proxy_managed_files.async_pre_call_hook( - user_api_key_dict=UserAPIKeyAuth( - user_id="user_b_id", parent_otel_span=MagicMock() - ), + user_api_key_dict=UserAPIKeyAuth(user_id="user_b_id", parent_otel_span=MagicMock()), cache=MagicMock(), data={"file_id": unified_file_id}, call_type="afile_delete", @@ -5280,20 +4971,14 @@ async def test_user_a_can_retrieve_own_file(): ) prisma_client.db.litellm_managedfiletable.find_first.return_value = file_record - proxy_managed_files = PROXY_LiteLLMManagedFiles( - MagicMock(), prisma_client=prisma_client - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(MagicMock(), prisma_client=prisma_client) # User A retrieves their own file - unified_file_id = ( - "bGl0ZWxsbV9wcm94eTphcHBsaWNhdGlvbi9qc29uO3VuaWZpZWRfaWQsZmlsZS1hYmMxMjM" - ) + unified_file_id = "bGl0ZWxsbV9wcm94eTphcHBsaWNhdGlvbi9qc29uO3VuaWZpZWRfaWQsZmlsZS1hYmMxMjM" # Should not raise an exception result = await proxy_managed_files.async_pre_call_hook( - user_api_key_dict=UserAPIKeyAuth( - user_id="user_a_id", parent_otel_span=MagicMock() - ), + user_api_key_dict=UserAPIKeyAuth(user_id="user_a_id", parent_otel_span=MagicMock()), cache=MagicMock(), data={"file_id": unified_file_id}, call_type="afile_retrieve", @@ -5333,9 +5018,7 @@ async def test_list_batches_only_returns_user_own_batches(): # Mock database to only return User A's batches prisma_client.db.litellm_managedobjecttable.find_many.return_value = [batch_user_a] - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=prisma_client - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=prisma_client) # User A requests their batches result = await proxy_managed_files.list_user_batches( @@ -5372,19 +5055,13 @@ async def test_same_user_different_keys_can_access_batch(): batch_record.created_by = "user_a_id" prisma_client.db.litellm_managedobjecttable.find_first.return_value = batch_record - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=prisma_client - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=prisma_client) - unified_batch_id = ( - "bGl0ZWxsbV9wcm94eTttb2RlbF9pZDpteS1tb2RlbDtsbG1fYmF0Y2hfaWQ6YmF0Y2hfYWJjMTIz" - ) + unified_batch_id = "bGl0ZWxsbV9wcm94eTttb2RlbF9pZDpteS1tb2RlbDtsbG1fYmF0Y2hfaWQ6YmF0Y2hfYWJjMTIz" # First API key for User A retrieves the batch result1 = await proxy_managed_files.async_pre_call_hook( - user_api_key_dict=UserAPIKeyAuth( - user_id="user_a_id", api_key="key-1", parent_otel_span=MagicMock() - ), + user_api_key_dict=UserAPIKeyAuth(user_id="user_a_id", api_key="key-1", parent_otel_span=MagicMock()), cache=MagicMock(), data={"batch_id": unified_batch_id}, call_type="aretrieve_batch", @@ -5394,9 +5071,7 @@ async def test_same_user_different_keys_can_access_batch(): # Second API key for the same User A retrieves the batch result2 = await proxy_managed_files.async_pre_call_hook( - user_api_key_dict=UserAPIKeyAuth( - user_id="user_a_id", api_key="key-2", parent_otel_span=MagicMock() - ), + user_api_key_dict=UserAPIKeyAuth(user_id="user_a_id", api_key="key-2", parent_otel_span=MagicMock()), cache=MagicMock(), data={"batch_id": unified_batch_id}, call_type="aretrieve_batch", @@ -5407,12 +5082,8 @@ async def test_same_user_different_keys_can_access_batch(): assert result1["batch_id"] == result2["batch_id"] -MODEL_ENCODED_BATCH_ID = encode_file_id_with_model( - "batch_provider123", "gpt-4o-team-alias", id_type="batch" -) -MODEL_ENCODED_OUTPUT_FILE_ID = encode_file_id_with_model( - "file-output456", "gpt-4o-team-alias", id_type="file" -) +MODEL_ENCODED_BATCH_ID = encode_file_id_with_model("batch_provider123", "gpt-4o-team-alias", id_type="batch") +MODEL_ENCODED_OUTPUT_FILE_ID = encode_file_id_with_model("file-output456", "gpt-4o-team-alias", id_type="file") RAW_PROVIDER_BATCH_ID = "batch_provider123" RAW_PROVIDER_FILE_ID = "file-output456" @@ -5444,12 +5115,8 @@ def _batch_response(batch_id, output_file_id=None, is_create=False): @pytest.mark.asyncio @pytest.mark.parametrize("call_type", ["aretrieve_batch", "acancel_batch"]) -@pytest.mark.parametrize( - "batch_id", [MODEL_ENCODED_BATCH_ID, RAW_PROVIDER_BATCH_ID] -) -async def test_team_b_cannot_access_team_a_provider_format_batch( - call_type, batch_id -): +@pytest.mark.parametrize("batch_id", [MODEL_ENCODED_BATCH_ID, RAW_PROVIDER_BATCH_ID]) +async def test_team_b_cannot_access_team_a_provider_format_batch(call_type, batch_id): """ Cross-team retrieve/cancel of a model-encoded or raw provider batch id must 403 when an ownership row exists for another team. @@ -5461,18 +5128,14 @@ async def test_team_b_cannot_access_team_a_provider_format_batch( from litellm.proxy._types import UserAPIKeyAuth prisma_client = AsyncMock() - prisma_client.db.litellm_managedobjecttable.find_first.return_value = ( - _owned_record(created_by="user_a", team_id="team_a") - ) - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=prisma_client + prisma_client.db.litellm_managedobjecttable.find_first.return_value = _owned_record( + created_by="user_a", team_id="team_a" ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=prisma_client) with pytest.raises(HTTPException) as exc_info: await proxy_managed_files.async_pre_call_hook( - user_api_key_dict=UserAPIKeyAuth( - user_id="user_b", team_id="team_b", parent_otel_span=MagicMock() - ), + user_api_key_dict=UserAPIKeyAuth(user_id="user_b", team_id="team_b", parent_otel_span=MagicMock()), cache=MagicMock(), data={"batch_id": batch_id}, call_type=call_type, @@ -5502,17 +5165,13 @@ async def test_authorized_callers_can_access_provider_format_batch(caller_kwargs from litellm.proxy._types import UserAPIKeyAuth prisma_client = AsyncMock() - prisma_client.db.litellm_managedobjecttable.find_first.return_value = ( - _owned_record(created_by="user_a", team_id="team_a") - ) - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=prisma_client + prisma_client.db.litellm_managedobjecttable.find_first.return_value = _owned_record( + created_by="user_a", team_id="team_a" ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=prisma_client) result = await proxy_managed_files.async_pre_call_hook( - user_api_key_dict=UserAPIKeyAuth( - parent_otel_span=MagicMock(), **caller_kwargs - ), + user_api_key_dict=UserAPIKeyAuth(parent_otel_span=MagicMock(), **caller_kwargs), cache=MagicMock(), data={"batch_id": MODEL_ENCODED_BATCH_ID}, call_type="aretrieve_batch", @@ -5532,14 +5191,10 @@ async def test_provider_format_batch_without_ownership_row_stays_accessible(): prisma_client = AsyncMock() prisma_client.db.litellm_managedobjecttable.find_first.return_value = None - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=prisma_client - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=prisma_client) result = await proxy_managed_files.async_pre_call_hook( - user_api_key_dict=UserAPIKeyAuth( - user_id="user_b", team_id="team_b", parent_otel_span=MagicMock() - ), + user_api_key_dict=UserAPIKeyAuth(user_id="user_b", team_id="team_b", parent_otel_span=MagicMock()), cache=MagicMock(), data={"batch_id": RAW_PROVIDER_BATCH_ID}, call_type="aretrieve_batch", @@ -5557,17 +5212,13 @@ async def test_fine_tuning_provider_format_id_not_enforced(): from litellm.proxy._types import UserAPIKeyAuth prisma_client = AsyncMock() - prisma_client.db.litellm_managedobjecttable.find_first.return_value = ( - _owned_record(created_by="user_a", team_id="team_a") - ) - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=prisma_client + prisma_client.db.litellm_managedobjecttable.find_first.return_value = _owned_record( + created_by="user_a", team_id="team_a" ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=prisma_client) result = await proxy_managed_files.async_pre_call_hook( - user_api_key_dict=UserAPIKeyAuth( - user_id="user_b", team_id="team_b", parent_otel_span=MagicMock() - ), + user_api_key_dict=UserAPIKeyAuth(user_id="user_b", team_id="team_b", parent_otel_span=MagicMock()), cache=MagicMock(), data={"fine_tuning_job_id": "ftjob-abc123"}, call_type="aretrieve_fine_tuning_job", @@ -5578,15 +5229,9 @@ async def test_fine_tuning_provider_format_id_not_enforced(): @pytest.mark.asyncio -@pytest.mark.parametrize( - "call_type", ["afile_content", "afile_retrieve", "afile_delete"] -) -@pytest.mark.parametrize( - "file_id", [MODEL_ENCODED_OUTPUT_FILE_ID, RAW_PROVIDER_FILE_ID] -) -async def test_team_b_cannot_access_team_a_provider_format_file( - call_type, file_id -): +@pytest.mark.parametrize("call_type", ["afile_content", "afile_retrieve", "afile_delete"]) +@pytest.mark.parametrize("file_id", [MODEL_ENCODED_OUTPUT_FILE_ID, RAW_PROVIDER_FILE_ID]) +async def test_team_b_cannot_access_team_a_provider_format_file(call_type, file_id): """ Cross-team content/retrieve/delete of a model-encoded or raw provider file id must 403 when an ownership row exists for another team. @@ -5594,18 +5239,14 @@ async def test_team_b_cannot_access_team_a_provider_format_file( from litellm.proxy._types import UserAPIKeyAuth prisma_client = AsyncMock() - prisma_client.db.litellm_managedfiletable.find_first.return_value = ( - _owned_record(created_by="user_a", team_id="team_a") - ) - proxy_managed_files = PROXY_LiteLLMManagedFiles( - MagicMock(), prisma_client=prisma_client + prisma_client.db.litellm_managedfiletable.find_first.return_value = _owned_record( + created_by="user_a", team_id="team_a" ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(MagicMock(), prisma_client=prisma_client) with pytest.raises(HTTPException) as exc_info: await proxy_managed_files.async_pre_call_hook( - user_api_key_dict=UserAPIKeyAuth( - user_id="user_b", team_id="team_b", parent_otel_span=MagicMock() - ), + user_api_key_dict=UserAPIKeyAuth(user_id="user_b", team_id="team_b", parent_otel_span=MagicMock()), cache=MagicMock(), data={"file_id": file_id}, call_type=call_type, @@ -5622,17 +5263,13 @@ async def test_same_team_can_access_provider_format_file(): from litellm.proxy._types import UserAPIKeyAuth prisma_client = AsyncMock() - prisma_client.db.litellm_managedfiletable.find_first.return_value = ( - _owned_record(created_by="user_a", team_id="team_a") - ) - proxy_managed_files = PROXY_LiteLLMManagedFiles( - MagicMock(), prisma_client=prisma_client + prisma_client.db.litellm_managedfiletable.find_first.return_value = _owned_record( + created_by="user_a", team_id="team_a" ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(MagicMock(), prisma_client=prisma_client) result = await proxy_managed_files.async_pre_call_hook( - user_api_key_dict=UserAPIKeyAuth( - user_id="teammate_of_a", team_id="team_a", parent_otel_span=MagicMock() - ), + user_api_key_dict=UserAPIKeyAuth(user_id="teammate_of_a", team_id="team_a", parent_otel_span=MagicMock()), cache=MagicMock(), data={"file_id": MODEL_ENCODED_OUTPUT_FILE_ID}, call_type="afile_content", @@ -5647,14 +5284,10 @@ async def test_provider_format_file_without_ownership_row_stays_accessible(): prisma_client = AsyncMock() prisma_client.db.litellm_managedfiletable.find_first.return_value = None - proxy_managed_files = PROXY_LiteLLMManagedFiles( - MagicMock(), prisma_client=prisma_client - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(MagicMock(), prisma_client=prisma_client) result = await proxy_managed_files.async_pre_call_hook( - user_api_key_dict=UserAPIKeyAuth( - user_id="user_b", team_id="team_b", parent_otel_span=MagicMock() - ), + user_api_key_dict=UserAPIKeyAuth(user_id="user_b", team_id="team_b", parent_otel_span=MagicMock()), cache=MagicMock(), data={"file_id": RAW_PROVIDER_FILE_ID}, call_type="afile_content", @@ -5669,9 +5302,7 @@ async def test_post_call_batch_create_stores_ownership_row(batch_id): from litellm.proxy._types import UserAPIKeyAuth prisma_client = AsyncMock() - proxy_managed_files = PROXY_LiteLLMManagedFiles( - MagicMock(async_set_cache=AsyncMock()), prisma_client=prisma_client - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(InternalUsageCache(DualCache()), prisma_client=prisma_client) await proxy_managed_files.async_post_call_success_hook( data={ @@ -5679,9 +5310,7 @@ async def test_post_call_batch_create_stores_ownership_row(batch_id): "endpoint": "/v1/chat/completions", "completion_window": "24h", }, - user_api_key_dict=UserAPIKeyAuth( - user_id="user_a", team_id="team_a", parent_otel_span=MagicMock() - ), + user_api_key_dict=UserAPIKeyAuth(user_id="user_a", team_id="team_a", parent_otel_span=MagicMock()), response=_batch_response(batch_id, is_create=True), ) @@ -5704,17 +5333,12 @@ async def test_post_call_batch_sync_does_not_claim_ownership(): prisma_client = AsyncMock() prisma_client.db.litellm_managedobjecttable.update_many.return_value = 0 - proxy_managed_files = PROXY_LiteLLMManagedFiles( - MagicMock(async_get_cache=AsyncMock(return_value=None), async_set_cache=AsyncMock()), - prisma_client=prisma_client, - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(InternalUsageCache(DualCache()), prisma_client=prisma_client) prisma_client.db.litellm_managedfiletable.find_first.return_value = None await proxy_managed_files.async_post_call_success_hook( data={"batch_id": MODEL_ENCODED_BATCH_ID}, - user_api_key_dict=UserAPIKeyAuth( - user_id="user_b", team_id="team_b", parent_otel_span=MagicMock() - ), + user_api_key_dict=UserAPIKeyAuth(user_id="user_b", team_id="team_b", parent_otel_span=MagicMock()), response=_batch_response(MODEL_ENCODED_BATCH_ID), ) @@ -5728,10 +5352,10 @@ async def test_post_call_batch_sync_updates_existing_row(): prisma_client = AsyncMock() prisma_client.db.litellm_managedobjecttable.update_many.return_value = 1 - prisma_client.db.litellm_managedobjecttable.find_first.return_value = ( - _owned_record(created_by="user_a", team_id="team_a") + prisma_client.db.litellm_managedobjecttable.find_first.return_value = _owned_record( + created_by="user_a", team_id="team_a" ) - proxy_managed_files = PROXY_LiteLLMManagedFiles(MagicMock(async_set_cache=AsyncMock()), prisma_client=prisma_client) + proxy_managed_files = PROXY_LiteLLMManagedFiles(InternalUsageCache(DualCache()), prisma_client=prisma_client) await proxy_managed_files.async_post_call_success_hook( data={"batch_id": MODEL_ENCODED_BATCH_ID}, @@ -5756,13 +5380,10 @@ async def test_post_call_batch_sync_stores_output_file_ownership_from_batch_row( prisma_client = AsyncMock() prisma_client.db.litellm_managedobjecttable.update_many.return_value = 1 - prisma_client.db.litellm_managedobjecttable.find_first.return_value = ( - _owned_record(created_by="user_a", team_id="team_a") - ) - proxy_managed_files = PROXY_LiteLLMManagedFiles( - MagicMock(async_get_cache=AsyncMock(return_value=None), async_set_cache=AsyncMock()), - prisma_client=prisma_client, + prisma_client.db.litellm_managedobjecttable.find_first.return_value = _owned_record( + created_by="user_a", team_id="team_a" ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(InternalUsageCache(DualCache()), prisma_client=prisma_client) prisma_client.db.litellm_managedfiletable.find_first.return_value = None await proxy_managed_files.async_post_call_success_hook( @@ -5772,15 +5393,11 @@ async def test_post_call_batch_sync_stores_output_file_ownership_from_batch_row( user_role="proxy_admin", parent_otel_span=MagicMock(), ), - response=_batch_response( - MODEL_ENCODED_BATCH_ID, output_file_id=MODEL_ENCODED_OUTPUT_FILE_ID - ), + response=_batch_response(MODEL_ENCODED_BATCH_ID, output_file_id=MODEL_ENCODED_OUTPUT_FILE_ID), ) file_upsert = prisma_client.db.litellm_managedfiletable.upsert.await_args - assert file_upsert.kwargs["where"] == { - "unified_file_id": MODEL_ENCODED_OUTPUT_FILE_ID - } + assert file_upsert.kwargs["where"] == {"unified_file_id": MODEL_ENCODED_OUTPUT_FILE_ID} create_data = file_upsert.kwargs["data"]["create"] assert create_data["created_by"] == "user_a" assert create_data["team_id"] == "team_a" @@ -5792,9 +5409,7 @@ async def test_post_call_batch_create_does_not_store_output_file_ownership(): from litellm.proxy._types import UserAPIKeyAuth prisma_client = AsyncMock() - proxy_managed_files = PROXY_LiteLLMManagedFiles( - MagicMock(async_set_cache=AsyncMock()), prisma_client=prisma_client - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(InternalUsageCache(DualCache()), prisma_client=prisma_client) await proxy_managed_files.async_post_call_success_hook( data={ @@ -5802,9 +5417,7 @@ async def test_post_call_batch_create_does_not_store_output_file_ownership(): "endpoint": "/v1/chat/completions", "completion_window": "24h", }, - user_api_key_dict=UserAPIKeyAuth( - user_id="user_a", team_id="team_a", parent_otel_span=MagicMock() - ), + user_api_key_dict=UserAPIKeyAuth(user_id="user_a", team_id="team_a", parent_otel_span=MagicMock()), response=_batch_response( MODEL_ENCODED_BATCH_ID, output_file_id=MODEL_ENCODED_OUTPUT_FILE_ID, @@ -5842,15 +5455,11 @@ async def test_file_list_cursors_are_scoped_to_the_caller(): prisma_client = AsyncMock() prisma_client.db.litellm_managedfiletable.find_many.return_value = [] - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=prisma_client - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=prisma_client) response = await proxy_managed_files.async_post_call_success_hook( data={}, - user_api_key_dict=UserAPIKeyAuth( - user_id="other-user", team_id="other-team", parent_otel_span=MagicMock() - ), + user_api_key_dict=UserAPIKeyAuth(user_id="other-user", team_id="other-team", parent_otel_span=MagicMock()), response=upstream_page, ) @@ -5899,15 +5508,11 @@ async def test_file_list_cursors_follow_the_owner_scoped_page(): } prisma_client = AsyncMock() prisma_client.db.litellm_managedfiletable.find_many.return_value = [managed_row] - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=prisma_client - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=prisma_client) response = await proxy_managed_files.async_post_call_success_hook( data={}, - user_api_key_dict=UserAPIKeyAuth( - user_id="mine-user", parent_otel_span=MagicMock() - ), + user_api_key_dict=UserAPIKeyAuth(user_id="mine-user", parent_otel_span=MagicMock()), response=upstream_page, ) @@ -5921,9 +5526,7 @@ async def test_file_list_cursors_follow_the_owner_scoped_page(): async def test_list_user_batches_provider_filter_rejected_with_400(): from litellm.proxy._types import ProxyException, UserAPIKeyAuth - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=MagicMock() - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(InternalUsageCache(DualCache()), prisma_client=MagicMock()) with pytest.raises(ProxyException) as exc: await proxy_managed_files.list_user_batches( @@ -5941,9 +5544,7 @@ async def test_list_user_batches_provider_filter_rejected_with_400(): async def test_list_user_batches_target_model_names_filter_rejected_with_400(): from litellm.proxy._types import ProxyException, UserAPIKeyAuth - proxy_managed_files = PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=MagicMock() - ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(InternalUsageCache(DualCache()), prisma_client=MagicMock()) with pytest.raises(ProxyException) as exc: await proxy_managed_files.list_user_batches( @@ -6262,7 +5863,9 @@ async def test_a_read_of_a_proxy_issued_response_id_is_marked_managed( monkeypatch.setenv("LITELLM_SALT_KEY", _MANAGED_READ_SALT_KEY) addressed_id: Final = _proxy_issued_response_id(encrypted) - result: Final = await PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=MagicMock()).async_pre_call_hook( + result: Final = await PROXY_LiteLLMManagedFiles( + InternalUsageCache(DualCache()), prisma_client=MagicMock() + ).async_pre_call_hook( user_api_key_dict=UserAPIKeyAuth(user_id="alice"), cache=MagicMock(), data={"response_id": addressed_id, "litellm_metadata": {"user_api_key_user_id": "alice"}}, @@ -6282,7 +5885,9 @@ async def test_a_read_of_a_raw_provider_response_id_is_not_marked_managed() -> N from litellm.constants import MANAGED_BACKGROUND_RESPONSE_METADATA_KEY from litellm.proxy._types import UserAPIKeyAuth - result: Final = await PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=MagicMock()).async_pre_call_hook( + result: Final = await PROXY_LiteLLMManagedFiles( + InternalUsageCache(DualCache()), prisma_client=MagicMock() + ).async_pre_call_hook( user_api_key_dict=UserAPIKeyAuth(user_id="alice"), cache=MagicMock(), data={"response_id": "resp_raw_provider_job", "metadata": {}}, @@ -6313,10 +5918,18 @@ async def test_only_a_streamed_background_create_is_marked_for_the_managed_bill( from litellm.constants import MANAGED_BACKGROUND_RESPONSE_METADATA_KEY from litellm.proxy._types import UserAPIKeyAuth - result: Final = await PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=MagicMock()).async_pre_call_hook( + result: Final = await PROXY_LiteLLMManagedFiles( + InternalUsageCache(DualCache()), prisma_client=MagicMock() + ).async_pre_call_hook( user_api_key_dict=UserAPIKeyAuth(user_id="alice"), cache=MagicMock(), - data={"model": "gpt-5.6", "input": "hello", "metadata": {"run": "nightly"}, "litellm_metadata": {}, **create_fields}, + data={ + "model": "gpt-5.6", + "input": "hello", + "metadata": {"run": "nightly"}, + "litellm_metadata": {}, + **create_fields, + }, call_type="aresponses", ) diff --git a/tests/unit/enterprise/proxy/test_managed_files_hook.py b/tests/unit/enterprise/proxy/test_managed_files_hook.py index f2df14cc8a3..c0f6018469f 100644 --- a/tests/unit/enterprise/proxy/test_managed_files_hook.py +++ b/tests/unit/enterprise/proxy/test_managed_files_hook.py @@ -10,14 +10,15 @@ import base64 import json import logging from types import MappingProxyType +from typing import Final +from unittest.mock import AsyncMock, MagicMock, patch import pytest -from typing import Final, Optional -from unittest.mock import AsyncMock, MagicMock, patch from litellm.llms.base_llm.files.transformation import FileListModelScope from litellm.proxy._types import LitellmUserRoles, ProxyException, UserAPIKeyAuth from litellm.proxy.openai_files_endpoints.common_utils import BATCH_CREATE_HIDDEN_PARAM +from litellm.proxy.utils import InternalUsageCache from litellm.types.llms.openai import FileListPage, OpenAIFileObject from litellm.types.utils import LiteLLMBatch @@ -36,7 +37,7 @@ def _make_file_object(file_id: str = "file-output-abc") -> OpenAIFileObject: def _make_batch_response( batch_id: str = "batch-123", - output_file_id: Optional[str] = "file-output-abc", + output_file_id: str | None = "file-output-abc", status: str = "completed", model_id: str = "model-deploy-xyz", model_name: str = "azure/gpt-4", @@ -100,11 +101,9 @@ def _make_managed_file_row( unified_file_id: str, purpose: str = "batch_output", created_by: str = "test-user", - team_id: Optional[str] = None, + team_id: str | None = None, ) -> MagicMock: - file_object = _make_file_object(f"file-provider-{unified_file_id}").model_copy( - update={"purpose": purpose} - ) + file_object = _make_file_object(f"file-provider-{unified_file_id}").model_copy(update={"purpose": purpose}) return MagicMock( unified_file_id=unified_file_id, file_object=file_object.model_dump(), @@ -116,7 +115,7 @@ def _make_managed_file_row( def _make_unparseable_managed_file_row( unified_file_id: str, created_by: str = "test-user", - team_id: Optional[str] = None, + team_id: str | None = None, ) -> MagicMock: """A row whose stored blob cannot be parsed back into a file object.""" return MagicMock( @@ -156,16 +155,10 @@ class _FakeManagedFileTable: return next(iter(self._owned_rows(where)), None) async def find_many(self, where, take=None, order=None, cursor=None, skip=0): - self.find_many_calls.append( - {"where": where, "take": take, "order": order, "cursor": cursor, "skip": skip} - ) + self.find_many_calls.append({"where": where, "take": take, "order": order, "cursor": cursor, "skip": skip}) rows = self._owned_rows(where) if cursor is not None: - start = next( - index - for index, row in enumerate(rows) - if row.unified_file_id == cursor["unified_file_id"] - ) + start = next(index for index, row in enumerate(rows) if row.unified_file_id == cursor["unified_file_id"]) rows = rows[start + skip :] return rows if take is None else rows[:take] @@ -183,13 +176,14 @@ def _make_managed_files_instance(): PROXY_LiteLLMManagedFiles, ) - mock_cache = MagicMock() - mock_cache.async_get_cache = AsyncMock(return_value=None) + from litellm.caching.dual_cache import DualCache + + cache: Final = DualCache() mock_prisma = MagicMock() mock_prisma.db.litellm_managedfiletable.find_first = AsyncMock(return_value=None) instance = PROXY_LiteLLMManagedFiles( - internal_usage_cache=mock_cache, + internal_usage_cache=cache, prisma_client=mock_prisma, ) instance.store_unified_file_id = AsyncMock() @@ -222,9 +216,7 @@ async def test_should_pass_credentials_to_afile_retrieve(): } mock_router = MagicMock() - mock_router.registry.return_value.credentials = MagicMock( - return_value=mock_credentials - ) + mock_router.registry.return_value.credentials = MagicMock(return_value=mock_credentials) mock_afile_retrieve = AsyncMock(return_value=_make_file_object("file-output-abc")) @@ -242,14 +234,10 @@ async def test_should_pass_credentials_to_afile_retrieve(): call_kwargs = mock_afile_retrieve.call_args assert call_kwargs.kwargs.get("api_key") == "test-azure-key", ( - f"afile_retrieve must receive api_key from router credentials. " - f"Got kwargs: {call_kwargs.kwargs}" + f"afile_retrieve must receive api_key from router credentials. Got kwargs: {call_kwargs.kwargs}" ) - assert ( - call_kwargs.kwargs.get("api_base") == "https://my-azure.openai.azure.com/" - ), ( - f"afile_retrieve must receive api_base from router credentials. " - f"Got kwargs: {call_kwargs.kwargs}" + assert call_kwargs.kwargs.get("api_base") == "https://my-azure.openai.azure.com/", ( + f"afile_retrieve must receive api_base from router credentials. Got kwargs: {call_kwargs.kwargs}" ) @@ -266,9 +254,7 @@ async def test_get_user_created_file_ids_skips_rows_without_file_object(): ] ) - files = await managed_files.get_user_created_file_ids( - _make_user_api_key_dict(), ["file-output-abc"] - ) + files = await managed_files.get_user_created_file_ids(_make_user_api_key_dict(), ["file-output-abc"]) assert [file.id for file in files] == ["unified-id-1"] @@ -294,9 +280,7 @@ async def test_get_user_created_file_ids_remaps_stored_raw_provider_id_to_unifie ] ) - files = await managed_files.get_user_created_file_ids( - _make_user_api_key_dict(), ["file-raw-provider-123"] - ) + files = await managed_files.get_user_created_file_ids(_make_user_api_key_dict(), ["file-raw-provider-123"]) assert [file.id for file in files] == [unified_id] assert files[0].filename == raw_provider_object.filename @@ -304,9 +288,7 @@ async def test_get_user_created_file_ids_remaps_stored_raw_provider_id_to_unifie @pytest.mark.asyncio -async def test_provider_file_id_resolver_returns_owned_mappings_with_owner_scoped_filter() -> ( - None -): +async def test_provider_file_id_resolver_returns_owned_mappings_with_owner_scoped_filter() -> None: managed_files: Final = _make_managed_files_instance() managed_row: Final = MagicMock( unified_file_id="unified-file-id", @@ -315,16 +297,14 @@ async def test_provider_file_id_resolver_returns_owned_mappings_with_owner_scope find_many: Final = AsyncMock(return_value=[managed_row]) managed_files.prisma_client.db.litellm_managedfiletable.find_many = find_many - unified_file_ids: Final = ( - await managed_files.get_unified_file_ids_for_provider_file_ids( - provider_file_ids=( - "file-provider-1", - "file-provider-2", - "file-unmanaged-2", - "file-provider-1", - ), - user_api_key_dict=_make_team_member_api_key_dict(), - ) + unified_file_ids: Final = await managed_files.get_unified_file_ids_for_provider_file_ids( + provider_file_ids=( + "file-provider-1", + "file-provider-2", + "file-unmanaged-2", + "file-provider-1", + ), + user_api_key_dict=_make_team_member_api_key_dict(), ) assert unified_file_ids == { @@ -343,9 +323,7 @@ async def test_provider_file_id_resolver_returns_owned_mappings_with_owner_scope @pytest.mark.asyncio -async def test_provider_file_id_resolver_denies_unowned_callers_without_database_query() -> ( - None -): +async def test_provider_file_id_resolver_denies_unowned_callers_without_database_query() -> None: managed_files: Final = _make_managed_files_instance() find_many: Final = AsyncMock() managed_files.prisma_client.db.litellm_managedfiletable.find_many = find_many @@ -357,11 +335,9 @@ async def test_provider_file_id_resolver_denies_unowned_callers_without_database parent_otel_span=None, ) - unified_file_ids: Final = ( - await managed_files.get_unified_file_ids_for_provider_file_ids( - provider_file_ids=("file-provider-1",), - user_api_key_dict=no_owner, - ) + unified_file_ids: Final = await managed_files.get_unified_file_ids_for_provider_file_ids( + provider_file_ids=("file-provider-1",), + user_api_key_dict=no_owner, ) assert unified_file_ids == {} @@ -375,11 +351,9 @@ async def test_provider_file_id_resolver_skips_database_query_for_empty_input() find_many: Final = AsyncMock() managed_files.prisma_client.db.litellm_managedfiletable.find_many = find_many - unified_file_ids: Final = ( - await managed_files.get_unified_file_ids_for_provider_file_ids( - provider_file_ids=(), - user_api_key_dict=_make_user_api_key_dict(), - ) + unified_file_ids: Final = await managed_files.get_unified_file_ids_for_provider_file_ids( + provider_file_ids=(), + user_api_key_dict=_make_user_api_key_dict(), ) assert unified_file_ids == {} @@ -397,9 +371,9 @@ async def test_afile_list_returns_owner_scoped_managed_files(): unified_file_id="unified-file-id", ), MagicMock( - file_object=_make_file_object("file-other-purpose").model_copy( - update={"purpose": "batch"} - ).model_dump(), + file_object=_make_file_object("file-other-purpose") + .model_copy(update={"purpose": "batch"}) + .model_dump(), unified_file_id="unified-other-purpose", ), ] @@ -519,9 +493,7 @@ async def test_afile_list_returns_own_and_team_files_for_a_key_carrying_both_ids ) assert [file.id for file in response.data] == ["unified-mine", "unified-teammates"] - assert table.find_many_calls[0]["where"] == { - "OR": [{"created_by": "test-user"}, {"team_id": "test-team"}] - } + assert table.find_many_calls[0]["where"] == {"OR": [{"created_by": "test-user"}, {"team_id": "test-team"}]} @pytest.mark.asyncio @@ -848,9 +820,7 @@ async def test_afile_list_reads_one_chunk_when_the_first_one_fills_the_page(): @pytest.mark.asyncio async def test_afile_list_reports_no_more_pages_when_nothing_matches(): - managed_files, _ = _make_managed_files_over_rows( - [_make_managed_file_row(f"unified-{index}") for index in range(5)] - ) + managed_files, _ = _make_managed_files_over_rows([_make_managed_file_row(f"unified-{index}") for index in range(5)]) page = await managed_files.afile_list( purpose="batch", @@ -1029,9 +999,7 @@ async def test_get_user_created_file_ids_skips_unparseable_rows(): ] ) - files = await managed_files.get_user_created_file_ids( - _make_user_api_key_dict(), ["file-output-abc"] - ) + files = await managed_files.get_user_created_file_ids(_make_user_api_key_dict(), ["file-output-abc"]) assert [file.id for file in files] == ["unified-valid"] @@ -1096,9 +1064,7 @@ async def test_should_not_double_wrap_already_unified_output_file_id(): "custom_llm_provider": "openai", } mock_router = MagicMock() - mock_router.registry.return_value.credentials = MagicMock( - return_value=mock_credentials - ) + mock_router.registry.return_value.credentials = MagicMock(return_value=mock_credentials) mock_afile_retrieve = AsyncMock(return_value=_make_file_object(provider_file_id)) with ( @@ -1115,13 +1081,9 @@ async def test_should_not_double_wrap_already_unified_output_file_id(): mock_afile_retrieve.assert_called_once() assert mock_afile_retrieve.call_args.kwargs["file_id"] == provider_file_id managed_files.store_unified_file_id.assert_awaited_once() - assert managed_files.store_unified_file_id.await_args.kwargs["model_mappings"] == { - model_id: provider_file_id - } + assert managed_files.store_unified_file_id.await_args.kwargs["model_mappings"] == {model_id: provider_file_id} - decoded = base64.urlsafe_b64decode( - already_unified + "=" * (-len(already_unified) % 4) - ).decode() + decoded = base64.urlsafe_b64decode(already_unified + "=" * (-len(already_unified) % 4)).decode() assert decoded.count(f"llm_output_file_id,{provider_file_id}") == 1 @@ -1132,9 +1094,7 @@ async def test_should_skip_non_file_unified_id_on_output_file_id(): managed_files = _make_managed_files_instance() batch_unified = ( - base64.urlsafe_b64encode( - b"litellm_proxy;model_id:openai/openai/gpt-5.5-batch;llm_batch_id:batch_abc" - ) + base64.urlsafe_b64encode(b"litellm_proxy;model_id:openai/openai/gpt-5.5-batch;llm_batch_id:batch_abc") .decode() .rstrip("=") ) @@ -1172,9 +1132,7 @@ async def test_afile_content_passes_trusted_model_credentials_to_router(): unified_file_id = "unified-file-id" s3_uri = "s3://my-bucket/litellm-batch-outputs/job-123/input.jsonl.out" managed_files.get_unified_file_id = AsyncMock(return_value=None) - managed_files.get_model_file_id_mapping = AsyncMock( - return_value={unified_file_id: {"model-123": s3_uri}} - ) + managed_files.get_model_file_id_mapping = AsyncMock(return_value={unified_file_id: {"model-123": s3_uri}}) mock_router = MagicMock() mock_router.registry.return_value.credentials = MagicMock( @@ -1210,9 +1168,10 @@ def _managed_deletion_file_id(provider_file_id): def _managed_files_with_deletion_row(unified_file_id, provider_file_id, file_object): + from litellm_enterprise.proxy.hooks.managed_files import PROXY_LiteLLMManagedFiles + from litellm.caching import DualCache from litellm.models.managed_files import LiteLLM_ManagedFileTable - from litellm_enterprise.proxy.hooks.managed_files import PROXY_LiteLLMManagedFiles row = LiteLLM_ManagedFileTable( unified_file_id=unified_file_id, @@ -1225,7 +1184,7 @@ def _managed_files_with_deletion_row(unified_file_id, provider_file_id, file_obj delete=AsyncMock(), ) return PROXY_LiteLLMManagedFiles( - internal_usage_cache=DualCache(), + internal_usage_cache=InternalUsageCache(DualCache()), prisma_client=MagicMock(db=MagicMock(litellm_managedfiletable=table)), ), table @@ -1260,9 +1219,9 @@ async def test_afile_delete_bedrock_uses_deployment_bucket_and_signed_s3_delete( unified_file_id = _managed_deletion_file_id(s3_uri) managed_files, table = _managed_files_with_deletion_row(unified_file_id, s3_uri, None) with respx.mock: - route = respx.delete( - "https://s3.us-west-2.amazonaws.com/my-bucket/litellm-bedrock-files/input.jsonl" - ).mock(return_value=httpx.Response(204)) + route = respx.delete("https://s3.us-west-2.amazonaws.com/my-bucket/litellm-bedrock-files/input.jsonl").mock( + return_value=httpx.Response(204) + ) response = await managed_files.afile_delete( file_id=unified_file_id, litellm_parent_otel_span=None, @@ -1344,15 +1303,11 @@ async def test_afile_content_bedrock_unified_id_end_to_end(monkeypatch): unified_file_id = "unified-file-id" s3_uri = "s3://my-bucket/litellm-batch-outputs/job-123/input.jsonl.out" managed_files.get_unified_file_id = AsyncMock(return_value=None) - managed_files.get_model_file_id_mapping = AsyncMock( - return_value={unified_file_id: {"model-123": s3_uri}} - ) + managed_files.get_model_file_id_mapping = AsyncMock(return_value={unified_file_id: {"model-123": s3_uri}}) expected_url = "https://s3.us-west-2.amazonaws.com/my-bucket/litellm-batch-outputs/job-123/input.jsonl.out" with respx.mock: - route = respx.get(expected_url).mock( - return_value=httpx.Response(200, content=b'{"recordId": "x"}') - ) + route = respx.get(expected_url).mock(return_value=httpx.Response(200, content=b'{"recordId": "x"}')) response = await managed_files.afile_content( file_id=unified_file_id, @@ -1361,9 +1316,7 @@ async def test_afile_content_bedrock_unified_id_end_to_end(monkeypatch): ) assert route.called - assert ( - route.calls[0].request.headers["Authorization"].startswith("AWS4-HMAC-SHA256") - ) + assert route.calls[0].request.headers["Authorization"].startswith("AWS4-HMAC-SHA256") assert response.content == b'{"recordId": "x"}' @@ -1375,15 +1328,13 @@ async def test_afile_content_error_reports_unified_id_not_provider_uri(): unified_file_id = "litellm_proxy_unified_id_abc" s3_uri = "s3://my-bucket/litellm-batch-outputs/job-123/input.jsonl.out" managed_files.get_unified_file_id = AsyncMock(return_value=None) - managed_files.get_model_file_id_mapping = AsyncMock( - return_value={unified_file_id: {"model-123": s3_uri}} - ) + managed_files.get_model_file_id_mapping = AsyncMock(return_value={unified_file_id: {"model-123": s3_uri}}) mock_router = MagicMock() mock_router.registry.return_value.credentials = MagicMock(return_value=None) mock_router.afile_content = AsyncMock(side_effect=Exception("deployment failed")) - with pytest.raises(Exception, match='LiteLLM Managed File object with') as exc_info: + with pytest.raises(Exception, match="LiteLLM Managed File object with") as exc_info: await managed_files.afile_content( file_id=unified_file_id, litellm_parent_otel_span=None, @@ -1402,17 +1353,15 @@ def _make_real_managed_files_instance(): PROXY_LiteLLMManagedFiles, ) - mock_cache = MagicMock() - mock_cache.async_get_cache = AsyncMock(return_value=None) - mock_cache.async_set_cache = AsyncMock() + from litellm.caching.dual_cache import DualCache + + mock_cache = DualCache() mock_prisma = MagicMock() mock_prisma.db.litellm_managedfiletable.find_first = AsyncMock(return_value=None) mock_prisma.db.litellm_managedfiletable.upsert = AsyncMock() mock_prisma.db.litellm_managedfiletable.create = AsyncMock( - side_effect=AssertionError( - "store_unified_file_id must upsert, not create, on the retrieve path" - ) + side_effect=AssertionError("store_unified_file_id must upsert, not create, on the retrieve path") ) return ( @@ -1454,9 +1403,7 @@ async def test_poll_refreshes_batch_state_without_claiming_the_row(): create-time snapshot forever -- while writing none of the attribution columns and never creating a row it would then own.""" managed_files, mock_prisma = _make_object_store_instance() - poller = UserAPIKeyAuth( - api_key="sk-the-poller", user_id="bob", team_id="team-bravo", parent_otel_span=None - ) + poller = UserAPIKeyAuth(api_key="sk-the-poller", user_id="bob", team_id="team-bravo", parent_otel_span=None) await managed_files.store_unified_object_id( unified_object_id="uoi-1", @@ -1490,9 +1437,7 @@ async def test_create_still_upserts_and_claims_attribution(): """The create is the one caller that can speak for the batch, so it keeps the upsert (creating the row when absent) and writes the attribution columns.""" managed_files, mock_prisma = _make_object_store_instance() - creator = UserAPIKeyAuth( - api_key="sk-the-creator", user_id="alice", team_id="team-alpha", parent_otel_span=None - ) + creator = UserAPIKeyAuth(api_key="sk-the-creator", user_id="alice", team_id="team-alpha", parent_otel_span=None) await managed_files.store_unified_object_id( unified_object_id="uoi-2", @@ -1672,8 +1617,7 @@ async def test_concurrent_first_registrations_converge_on_one_row(): def _b64_unified_input_file_id(target_model_names: str) -> str: unified_input_file_id = ( - "litellm_proxy:application/octet-stream;unified_id,input-uuid;" - f"target_model_names,{target_model_names}" + f"litellm_proxy:application/octet-stream;unified_id,input-uuid;target_model_names,{target_model_names}" ) return base64.urlsafe_b64encode(unified_input_file_id.encode()).decode().rstrip("=") @@ -1682,9 +1626,7 @@ def _b64_unified_input_file_id(target_model_names: str) -> str: async def test_hook_mint_prefers_input_file_target_model_names(): managed_files = _make_managed_files_instance() batch_response = _make_batch_response(model_name="model-a") - batch_response._hidden_params["unified_file_id"] = _b64_unified_input_file_id( - "model-a,model-b" - ) + batch_response._hidden_params["unified_file_id"] = _b64_unified_input_file_id("model-a,model-b") mock_router = MagicMock() mock_router.registry.return_value.credentials = MagicMock(return_value={}) @@ -1738,10 +1680,11 @@ async def test_hook_mint_falls_back_to_response_input_file_id_target_models(): @pytest.mark.asyncio async def test_cost_job_and_retrieve_paths_mint_identical_unified_output_file_ids(): + from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost + from litellm.proxy.openai_files_endpoints.common_utils import ( ensure_batch_response_managed_file_ids, ) - from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost managed_files, mock_prisma = _make_real_managed_files_instance() mock_prisma.db.litellm_managedfiletable.find_first = AsyncMock(return_value=None) @@ -2040,7 +1983,7 @@ async def test_afile_delete_storage_backed_row_deletes_stored_content_not_provid file_table = MagicMock(find_first=AsyncMock(return_value=row), delete=AsyncMock()) content_table = MagicMock(delete=AsyncMock()) managed_files = PROXY_LiteLLMManagedFiles( - internal_usage_cache=DualCache(), + internal_usage_cache=InternalUsageCache(DualCache()), prisma_client=MagicMock( db=MagicMock(litellm_managedfiletable=file_table, litellm_managedfilecontenttable=content_table) ), @@ -2084,7 +2027,7 @@ async def test_afile_content_storage_backed_row_returns_stored_bytes_not_provide file_table = MagicMock(find_first=AsyncMock(return_value=row)) content_table = MagicMock(find_unique=AsyncMock(return_value=MagicMock(content=Base64.encode(stored_bytes)))) managed_files = PROXY_LiteLLMManagedFiles( - internal_usage_cache=DualCache(), + internal_usage_cache=InternalUsageCache(DualCache()), prisma_client=MagicMock( db=MagicMock(litellm_managedfiletable=file_table, litellm_managedfilecontenttable=content_table) ), @@ -2142,7 +2085,7 @@ async def test_store_unified_file_id_caches_the_storage_location_the_db_row_gets file_table = MagicMock(upsert=AsyncMock(), find_first=AsyncMock(side_effect=AssertionError("cache miss"))) managed_files = PROXY_LiteLLMManagedFiles( - internal_usage_cache=DualCache(), + internal_usage_cache=InternalUsageCache(DualCache()), prisma_client=MagicMock(db=MagicMock(litellm_managedfiletable=file_table)), ) stored = _make_file_object("file-kept").model_copy(update={"purpose": "batch"}) diff --git a/tests/unit/redis/stores/test_managed_files.py b/tests/unit/redis/stores/test_managed_files.py new file mode 100644 index 00000000000..516ad0745c5 --- /dev/null +++ b/tests/unit/redis/stores/test_managed_files.py @@ -0,0 +1,264 @@ +import json +import sys +import time +from collections.abc import Mapping +from datetime import timedelta +from typing import Final, cast +from unittest.mock import patch + +import fakeredis +import pytest +from litellm_enterprise.proxy.hooks.managed_files import PROXY_LiteLLMManagedFiles +from pydantic import ValidationError +from redis.asyncio import Redis +from redis.exceptions import TimeoutError as RedisTimeoutError + +from litellm.caching.dual_cache import DualCache +from litellm.models.managed_files import LiteLLM_ManagedFileTable +from litellm.proxy.utils import InternalUsageCache, PrismaClient +from litellm.types.llms.openai import OpenAIFileObject +from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth +from tests._support.fake_redis import fake_redis_cache +from tests.unit.enterprise.proxy.hooks.test_managed_files import ( + _InMemoryManagedFilesPrismaClient, # pyright: ignore[reportPrivateUsage] # in-memory doubles shared with the hook suite + _InMemoryManagedFileTable, # pyright: ignore[reportPrivateUsage] # in-memory doubles shared with the hook suite +) + + +async def _stored(raw: "Redis[bytes]", key: str) -> bytes | None: + return await raw.get(key) + + +def _row(file_id: str = "unified-é") -> LiteLLM_ManagedFileTable: + return LiteLLM_ManagedFileTable( + unified_file_id=file_id, + file_object=OpenAIFileObject( + id=file_id, object="file", bytes=123, created_at=456, filename="文.jsonl", purpose="batch" + ), + model_mappings={"model": "provider-file"}, + flat_model_file_ids=["provider-file"], + created_by="owner", + updated_by="owner", + ) + + +def _hook(cache: DualCache, table: _InMemoryManagedFileTable) -> PROXY_LiteLLMManagedFiles: + return PROXY_LiteLLMManagedFiles( + InternalUsageCache(cache), cast(PrismaClient, _InMemoryManagedFilesPrismaClient(table)) + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("file_id", ["unified-é", "tenant:unified-é"]) +@pytest.mark.parametrize("ttl", [1, 17.9, None, 0.5, 0, -1]) +async def test_file_metadata_matches_legacy_wire_and_ttl( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis, file_id: str, ttl: float | None +) -> None: + cache: Final = fake_redis_cache(server, namespace="tenant") + dual: Final = DualCache(redis_cache=cache, default_in_memory_ttl=ttl, default_redis_ttl=999) + dual.update_cache_ttl(ttl, None) + hook: Final = _hook(dual, _InMemoryManagedFileTable()) + row: Final = _row(file_id) + key: Final = file_id if file_id.startswith("tenant:") else f"tenant:{file_id}" + with patch("time.time", return_value=time.time()): + await hook.store_unified_file_id( + file_id, row.file_object, None, row.model_mappings, UserAPIKeyAuth(user_id="owner") + ) + encoded: Final = await _stored(raw, key) + expiry: Final = await raw.pttl(key) + await cache.async_set_cache("legacy", row.model_dump(), **({"ttl": ttl} if ttl is not None else {})) + assert encoded == await raw.get("tenant:legacy") + assert expiry == await raw.pttl("tenant:legacy") + assert encoded == (None if ttl is not None and ttl < 1 else json.dumps(row.model_dump()).encode()) + assert await cache.async_get_cache(file_id) == (None if ttl is not None and ttl < 1 else row.model_dump()) + assert await raw.get(f"tenant:{key}") is None + + +@pytest.mark.asyncio +async def test_old_writer_new_reader_backfills_only_redis_hits( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + cache: Final = fake_redis_cache(server, namespace="tenant") + dual: Final = DualCache(redis_cache=cache, default_in_memory_ttl=1) + row: Final = _row() + table: Final = _InMemoryManagedFileTable(row) + hook: Final = _hook(dual, table) + await cache.async_set_cache(row.unified_file_id, row.model_dump(), ttl=60) + assert await hook.get_unified_file_id(row.unified_file_id) == row + assert table.find_first_calls == [] + server.connected = False + assert await hook.get_unified_file_id(row.unified_file_id) == row + server.connected = True + dual.in_memory_cache.flush_cache() + await raw.delete(f"tenant:{row.unified_file_id}") + assert await hook.get_unified_file_id(row.unified_file_id) == row + assert await raw.get(f"tenant:{row.unified_file_id}") is None + assert len(table.find_first_calls) == 1 + + +@pytest.mark.asyncio +async def test_file_expiry_and_legacy_literal_record( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + row: Final = _row() + with patch("time.time", return_value=1_800_000_000): + cache: Final = fake_redis_cache(server) + table: Final = _InMemoryManagedFileTable(row) + hook: Final = _hook(DualCache(redis_cache=cache, default_in_memory_ttl=1), table) + await raw.set(row.unified_file_id, repr(row.model_dump()), ex=1) + assert await hook.get_unified_file_id(row.unified_file_id) == row + assert table.find_first_calls == [] + with patch("time.time", return_value=1_800_000_002): + assert await hook.get_unified_file_id(row.unified_file_id) == row + assert len(table.find_first_calls) == 1 + assert await raw.get(row.unified_file_id) is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("encoded", [b"null", b"{}", b"broken", b"\xff"]) +async def test_file_misses_and_malformed_bytes_fall_back_to_database( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis, encoded: bytes +) -> None: + row: Final = _row() + table: Final = _InMemoryManagedFileTable(row) + hook: Final = _hook(DualCache(redis_cache=fake_redis_cache(server)), table) + await raw.set(row.unified_file_id, encoded) + assert await hook.get_unified_file_id(row.unified_file_id) == row + assert len(table.find_first_calls) == 1 + assert await raw.get(row.unified_file_id) == encoded + + +@pytest.mark.asyncio +async def test_recursive_file_metadata_retains_legacy_database_fallback( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + row: Final = _row() + table: Final = _InMemoryManagedFileTable(row) + cache: Final = fake_redis_cache(server) + hook: Final = _hook(DualCache(redis_cache=cache), table) + await raw.set(row.unified_file_id, b"[" * 10000 + b"0" + b"]" * 10000) + limit: Final = sys.getrecursionlimit() + try: + sys.setrecursionlimit(1000) + assert await cache.async_get_cache(row.unified_file_id) is None + assert await hook.get_unified_file_id(row.unified_file_id) == row + assert len(table.find_first_calls) == 1 + finally: + sys.setrecursionlimit(limit) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("encoded", [b'{"invalid": true}', b"(1, 2)", b"{'item'}", b"{1: 'value'}"]) +async def test_truthy_invalid_file_record_still_raises_validation_error( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis, encoded: bytes +) -> None: + table: Final = _InMemoryManagedFileTable(_row()) + hook: Final = _hook(DualCache(redis_cache=fake_redis_cache(server)), table) + await raw.set("unified-é", encoded) + with pytest.raises(ValidationError): + await hook.get_unified_file_id("unified-é") + assert table.find_first_calls == [] + + +class _DeleteObservingTable(_InMemoryManagedFileTable): + def __init__(self, raw: fakeredis.FakeAsyncRedis, row: LiteLLM_ManagedFileTable) -> None: + super().__init__(row) + self.raw = raw + + async def delete(self, where: Mapping[str, object]) -> LiteLLM_ManagedFileTable | None: + assert await self.raw.get(f"tenant:{where['unified_file_id']}") == b"null" + assert await self.raw.pttl(f"tenant:{where['unified_file_id']}") == 1000 + return await super().delete(where) + + +@pytest.mark.asyncio +async def test_file_delete_writes_null_before_database_delete( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + row: Final = _row() + table: Final = _DeleteObservingTable(raw, row) + hook: Final = _hook( + DualCache(redis_cache=fake_redis_cache(server, namespace="tenant"), default_in_memory_ttl=1), table + ) + with patch("time.time", return_value=time.time()): + assert await hook.delete_unified_file_id(row.unified_file_id) == row.file_object + assert table.rows == {} + assert await hook.get_unified_file_id(row.unified_file_id) is None + assert len(table.find_first_calls) == 2 + + +@pytest.mark.asyncio +async def test_file_outage_keeps_local_write_and_database_policy(server: fakeredis.FakeServer) -> None: + row: Final = _row() + cache: Final = fake_redis_cache(server) + table: Final = _InMemoryManagedFileTable() + dual: Final = DualCache(redis_cache=cache, default_in_memory_ttl=1) + hook: Final = _hook(dual, table) + server.connected = False + await hook.store_unified_file_id( + row.unified_file_id, row.file_object, None, row.model_mappings, UserAPIKeyAuth(user_id="owner") + ) + assert await hook.get_unified_file_id(row.unified_file_id) == row + assert table.rows[row.unified_file_id] == row + dual.in_memory_cache.flush_cache() + for _ in range(8): + assert await hook.get_unified_file_id(row.unified_file_id) == row + assert cache._circuit_breaker.state == "open" # pyright: ignore[reportPrivateUsage] # asserts the breaker opened after the outage + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ttl", [float("nan"), float("inf"), float("-inf")]) +async def test_invalid_file_cache_ttl_does_not_block_database_write( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis, ttl: float +) -> None: + row: Final = _row() + table: Final = _InMemoryManagedFileTable() + dual: Final = DualCache(redis_cache=fake_redis_cache(server), default_in_memory_ttl=ttl) + hook: Final = _hook(dual, table) + with patch("time.time", return_value=1_800_000_000): + await dual.async_set_cache("legacy", row.model_dump()) + await hook.store_unified_file_id( + row.unified_file_id, row.file_object, None, row.model_mappings, UserAPIKeyAuth(user_id="owner") + ) + assert table.rows[row.unified_file_id] == row + assert await raw.get(row.unified_file_id) == await raw.get("legacy") + assert await raw.pttl(row.unified_file_id) == await raw.pttl("legacy") + + +@pytest.mark.asyncio +async def test_unknown_applied_file_write_is_not_rolled_back( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + cache: Final = fake_redis_cache(server, namespace="tenant") + row: Final = _row() + table: Final = _InMemoryManagedFileTable() + hook: Final = _hook(DualCache(redis_cache=cache, default_in_memory_ttl=1), table) + + async def applied_then_timeout(name: str, value: str, *, ex: timedelta | None, nx: bool) -> bool: + await raw.set(name, value, ex=ex, nx=nx) + raise RedisTimeoutError("reply lost after apply") + + with patch.object(cache.init_async_client(), "set", side_effect=applied_then_timeout): + await hook.store_unified_file_id( + row.unified_file_id, row.file_object, None, row.model_mappings, UserAPIKeyAuth(user_id="owner") + ) + assert await raw.get(f"tenant:{row.unified_file_id}") == json.dumps(row.model_dump()).encode() + assert table.rows[row.unified_file_id] == row + + +@pytest.mark.asyncio +async def test_file_refresh_and_lost_compare_and_swap_preserve_cache_order( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + row: Final = _row() + table: Final = _InMemoryManagedFileTable(row) + hook: Final = _hook(DualCache(redis_cache=fake_redis_cache(server), default_in_memory_ttl=1), table) + assert row.file_object is not None + updated: Final = row.file_object.model_copy(update={"bytes": 987}) + await hook._save_refreshed_file_object(row, updated) # pyright: ignore[reportPrivateUsage] # exercises the refresh write path + assert await raw.get(row.unified_file_id) == json.dumps(table.rows[row.unified_file_id].model_dump()).encode() + assert await hook._save_xai_refreshed_file_object(row, updated) is None # pyright: ignore[reportPrivateUsage] # exercises the refresh write path + assert await raw.get(row.unified_file_id) == b"null" + assert table.rows[row.unified_file_id].file_object == updated + assert await hook._save_xai_refreshed_file_object(table.rows[row.unified_file_id], updated) == updated # pyright: ignore[reportPrivateUsage] # exercises the refresh write path + assert await raw.get(row.unified_file_id) == json.dumps(table.rows[row.unified_file_id].model_dump()).encode() From d6e9763b39e3268f7d0bcc2115adacc86eb29b21 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 19:49:22 +0000 Subject: [PATCH 03/17] refactor(redis): isolate managed vector resource metadata store --- .../base_managed_resource.py | 80 ++++- .../redis/stores/managed_vector_resources.py | 71 +++++ .../base_llm/test_base_managed_resource.py | 36 +-- .../stores/test_managed_vector_resources.py | 286 ++++++++++++++++++ 4 files changed, 436 insertions(+), 37 deletions(-) create mode 100644 litellm/redis/stores/managed_vector_resources.py create mode 100644 tests/unit/redis/stores/test_managed_vector_resources.py diff --git a/litellm/llms/base_llm/managed_resources/base_managed_resource.py b/litellm/llms/base_llm/managed_resources/base_managed_resource.py index 2663a3c42e7..fc988a8ec3a 100644 --- a/litellm/llms/base_llm/managed_resources/base_managed_resource.py +++ b/litellm/llms/base_llm/managed_resources/base_managed_resource.py @@ -4,18 +4,29 @@ import base64 import json +import math from abc import ABC, abstractmethod from collections.abc import Mapping +from datetime import timedelta from typing import TYPE_CHECKING, Any, Final, Generic, Protocol, TypeVar, cast, runtime_checkable from litellm import verbose_logger from litellm._internal_context import with_service_target +from litellm.caching.dual_cache import DualCache +from litellm.caching.redis_batch import active_request_redis_batch +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.llms.base_llm.managed_resources.isolation import ( build_list_page, build_owner_filter, can_access_resource, resolve_resource_owner_id, ) +from litellm.redis.stores.managed_vector_resources import ( + ManagedVectorResourceStore, + TransportManagedVectorResourceStore, + vector_resource_key, +) +from litellm.redis.transport.contract import RedisFailure from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth from litellm.types.utils import SpecialEnums @@ -66,6 +77,12 @@ class _ManagedResourceTable(Protocol[ResourceObjectType]): async def delete(self, *, where: Mapping[str, object]) -> object: ... +class _LocalVectorResourceCache(Protocol): + async def async_get_cache(self, key: str, *, local_only: bool) -> object: ... + + async def async_set_cache(self, key: str, value: object, *, local_only: bool) -> None: ... + + class BaseManagedResource(ABC, Generic[ResourceObjectType]): """ Base class for managing resources with target_model_names support. @@ -95,6 +112,50 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): def _resource_table(self) -> _ManagedResourceTable[ResourceObjectType]: return getattr(self.prisma_client.db, self.table_name) + def _metadata_cache(self) -> DualCache: + if isinstance(self.internal_usage_cache, DualCache): + return self.internal_usage_cache + return self.internal_usage_cache.dual_cache + + async def _cache_vector_resource(self, unified_resource_id: str, record: Mapping[str, object] | None) -> None: + cache: Final = self._metadata_cache() + local_cache: Final = require_attributes(cache, _LocalVectorResourceCache, "async_get_cache", "async_set_cache") + await local_cache.async_set_cache(vector_resource_key(unified_resource_id), record, local_only=True) + if cache.redis_cache is None: + return + store: Final[ManagedVectorResourceStore] = TransportManagedVectorResourceStore(cache.redis_cache.transport()) + try: + seconds: Final = ( + int(cache.default_in_memory_ttl) + if cache.default_in_memory_ttl is not None and not math.isnan(cache.default_in_memory_ttl) + else cache.redis_cache.default_ttl + ) + ttl: Final = timedelta(seconds=seconds) + except (OverflowError, ValueError): + return + if record is None: + await store.invalidate_vector_resource(unified_resource_id, ttl=ttl) + else: + await store.write_vector_resource(unified_resource_id, record, ttl=ttl) + + async def _cached_vector_resource(self, unified_resource_id: str) -> object: + cache: Final = self._metadata_cache() + local_cache: Final = require_attributes(cache, _LocalVectorResourceCache, "async_get_cache", "async_set_cache") + local: Final = await local_cache.async_get_cache(vector_resource_key(unified_resource_id), local_only=True) + if local is not None or cache.redis_cache is None: + return local + request_batch: Final = active_request_redis_batch(cache.redis_cache) + if request_batch is not None and request_batch.read_as_missing(vector_resource_key(unified_resource_id)): + return None + store: Final[ManagedVectorResourceStore] = TransportManagedVectorResourceStore(cache.redis_cache.transport()) + result: Final = await store.read_vector_resource(unified_resource_id) + if result is None or isinstance(result, RedisFailure): + return None + value: Final = result.value + if value is not None: + await local_cache.async_set_cache(vector_resource_key(unified_resource_id), value, local_only=True) + return value + # ============================================================================ # ABSTRACT METHODS # ============================================================================ @@ -201,11 +262,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): # Store in cache if resource_object is not None: - await self.internal_usage_cache.async_set_cache( - key=unified_resource_id, - value=cache_data, - litellm_parent_otel_span=litellm_parent_otel_span, - ) + await self._cache_vector_resource(unified_resource_id, cache_data) # Prepare database data db_data: Final[dict[str, object]] = { @@ -263,10 +320,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): # Check cache first result: Final = cast( dict | None, - await self.internal_usage_cache.async_get_cache( - key=unified_resource_id, - litellm_parent_otel_span=litellm_parent_otel_span, - ), + await self._cached_vector_resource(unified_resource_id), ) if result: @@ -305,11 +359,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): raise Exception(f"LiteLLM Managed {self.resource_type} with id={unified_resource_id} not found") # Delete from cache - await self.internal_usage_cache.async_set_cache( - key=unified_resource_id, - value=None, - litellm_parent_otel_span=litellm_parent_otel_span, - ) + await self._cache_vector_resource(unified_resource_id, None) # Delete from database await table.delete(where={"unified_resource_id": unified_resource_id}) @@ -355,7 +405,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): async def get_model_resource_id_mapping( self, resource_ids: list[str], - litellm_parent_otel_span: Span, + litellm_parent_otel_span: Span | None, ) -> dict[str, dict[str, str]]: """ Get model-specific resource IDs for a list of unified resource IDs. diff --git a/litellm/redis/stores/managed_vector_resources.py b/litellm/redis/stores/managed_vector_resources.py new file mode 100644 index 00000000000..732174dfb2e --- /dev/null +++ b/litellm/redis/stores/managed_vector_resources.py @@ -0,0 +1,71 @@ +import ast +import json +from collections.abc import Mapping +from dataclasses import dataclass +from datetime import timedelta +from typing import Final, Protocol + +from pydantic import TypeAdapter + +from litellm.redis.transport.contract import RedisFailure, RedisTransport + + +@dataclass(frozen=True, slots=True) +class ManagedVectorResourceValue: + value: object + + +_VALUE: Final[TypeAdapter[ManagedVectorResourceValue]] = TypeAdapter(ManagedVectorResourceValue) + + +def vector_resource_key(unified_resource_id: str) -> str: + return unified_resource_id + + +class ManagedVectorResourceStore(Protocol): + async def read_vector_resource( + self, unified_resource_id: str + ) -> ManagedVectorResourceValue | None | RedisFailure: ... + + async def write_vector_resource( + self, unified_resource_id: str, record: Mapping[str, object], *, ttl: timedelta | None + ) -> None | RedisFailure: ... + + async def invalidate_vector_resource( + self, unified_resource_id: str, *, ttl: timedelta | None + ) -> None | RedisFailure: ... + + +@dataclass(frozen=True, slots=True) +class TransportManagedVectorResourceStore: + transport: RedisTransport + + async def read_vector_resource(self, unified_resource_id: str) -> ManagedVectorResourceValue | None | RedisFailure: + raw: Final = await self.transport.get(vector_resource_key(unified_resource_id)) + if raw is None or isinstance(raw, RedisFailure): + return raw + try: + return _VALUE.validate_python({"value": json.loads(raw.decode("utf-8"))}) + except (ValueError, UnicodeError, RecursionError): + try: + return _VALUE.validate_python({"value": ast.literal_eval(raw.decode("utf-8"))}) + except (ValueError, SyntaxError, UnicodeError, RecursionError) as exc: + return RedisFailure(kind="decode", message=str(exc)) + + async def write_vector_resource( + self, unified_resource_id: str, record: Mapping[str, object], *, ttl: timedelta | None + ) -> None | RedisFailure: + try: + encoded: Final = json.dumps(record) + except (TypeError, ValueError) as exc: + return RedisFailure(kind="decode", message=str(exc)) + return await self._write(unified_resource_id, encoded, ttl) + + async def invalidate_vector_resource( + self, unified_resource_id: str, *, ttl: timedelta | None + ) -> None | RedisFailure: + return await self._write(unified_resource_id, "null", ttl) + + async def _write(self, unified_resource_id: str, encoded: str, ttl: timedelta | None) -> None | RedisFailure: + result: Final = await self.transport.set_seconds(vector_resource_key(unified_resource_id), encoded, ttl=ttl) + return result if isinstance(result, RedisFailure) else None diff --git a/tests/unit/llms/base_llm/test_base_managed_resource.py b/tests/unit/llms/base_llm/test_base_managed_resource.py index a6c518616f4..ebe93afe322 100644 --- a/tests/unit/llms/base_llm/test_base_managed_resource.py +++ b/tests/unit/llms/base_llm/test_base_managed_resource.py @@ -2,15 +2,16 @@ Integration tests for BaseManagedResource listing and access control. """ -from typing import List from unittest.mock import AsyncMock, MagicMock import pytest +from litellm.caching.dual_cache import DualCache from litellm.llms.base_llm.managed_resources.base_managed_resource import ( BaseManagedResource, ) from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.utils import InternalUsageCache class _StubResource(BaseManagedResource): @@ -24,18 +25,14 @@ class _StubResource(BaseManagedResource): def table_name(self) -> str: return "litellm_test_resource_table" - def get_unified_resource_id_format( - self, resource_object, target_model_names_list - ) -> str: + def get_unified_resource_id_format(self, resource_object, target_model_names_list) -> str: return "test" - async def create_resource_for_model( - self, llm_router, model, request_data, litellm_parent_otel_span - ): + async def create_resource_for_model(self, llm_router, model, request_data, litellm_parent_otel_span): return {"id": "test"} -def _make_resource(records: List = None) -> _StubResource: +def _make_resource(records: list[object] | None = None) -> _StubResource: cache = MagicMock() cache.async_get_cache = AsyncMock(return_value=None) @@ -68,9 +65,7 @@ async def test_list_user_filters_by_user_id(): await resource.list_user_resources(user_api_key_dict=user) - where = resource.prisma_client.db.litellm_test_resource_table.find_many.await_args.kwargs[ - "where" - ] + where = resource.prisma_client.db.litellm_test_resource_table.find_many.await_args.kwargs["where"] assert where["created_by"] == "alice" assert "team_id" not in where @@ -82,9 +77,7 @@ async def test_list_service_account_filters_by_team_id(): await resource.list_user_resources(user_api_key_dict=service_account) - where = resource.prisma_client.db.litellm_test_resource_table.find_many.await_args.kwargs[ - "where" - ] + where = resource.prisma_client.db.litellm_test_resource_table.find_many.await_args.kwargs["where"] assert where["team_id"] == "team-eng" assert "created_by" not in where @@ -114,18 +107,17 @@ async def test_list_identity_less_caller_returns_empty_without_query(): [("team-eng", True), ("team-sales", False), (None, False)], ) async def test_can_access_uses_team_id_for_service_account(caller_team_id, expected): - cache = MagicMock() - cache.async_get_cache = AsyncMock( - return_value={ + cache = DualCache() + await cache.async_set_cache( + key="rid", + value={ "created_by": None, "team_id": "team-eng", - } + }, ) prisma = MagicMock() - resource = _StubResource(internal_usage_cache=cache, prisma_client=prisma) + resource = _StubResource(internal_usage_cache=InternalUsageCache(cache), prisma_client=prisma) - caller = ( - UserAPIKeyAuth(team_id=caller_team_id) if caller_team_id else UserAPIKeyAuth() - ) + caller = UserAPIKeyAuth(team_id=caller_team_id) if caller_team_id else UserAPIKeyAuth() assert await resource.can_user_access_unified_resource_id("rid", caller) is expected diff --git a/tests/unit/redis/stores/test_managed_vector_resources.py b/tests/unit/redis/stores/test_managed_vector_resources.py new file mode 100644 index 00000000000..624ab3bf21c --- /dev/null +++ b/tests/unit/redis/stores/test_managed_vector_resources.py @@ -0,0 +1,286 @@ +import base64 +import json +import sys +import time +from collections.abc import Mapping +from dataclasses import dataclass +from datetime import timedelta +from types import SimpleNamespace +from typing import Final, cast +from unittest.mock import patch + +import fakeredis +import pytest +from litellm_enterprise.proxy.hooks.managed_vector_stores import PROXY_LiteLLMManagedVectorStores +from redis.exceptions import TimeoutError as RedisTimeoutError + +from litellm.caching.dual_cache import DualCache +from litellm.proxy.utils import InternalUsageCache, PrismaClient +from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth +from litellm.types.vector_stores import VectorStoreCreateResponse +from tests._support.fake_redis import fake_redis_cache + + +@dataclass(frozen=True) +class _Row: + data: Mapping[str, object] + + @property + def resource_object(self) -> object: + return self.data.get("resource_object") + + def model_dump(self) -> dict[str, object]: + return dict(self.data) + + +class _Table: + def __init__(self, row: _Row | None = None, raw: fakeredis.FakeAsyncRedis | None = None) -> None: + self.row = row + self.reads = 0 + self.raw = raw + + async def create(self, data: Mapping[str, object]) -> _Row: + self.row = _Row(data) + return self.row + + async def find_first(self, where: Mapping[str, object]) -> _Row | None: + self.reads += 1 + return self.row if self.row and self.row.data["unified_resource_id"] == where["unified_resource_id"] else None + + async def delete(self, where: Mapping[str, object]) -> None: + if self.raw is not None: + assert await self.raw.get(f"tenant:{where['unified_resource_id']}") == b"null" + self.row = None + + +def _resource() -> VectorStoreCreateResponse: + return cast(VectorStoreCreateResponse, {"id": "vs-é", "object": "vector_store", "name": "文", "usage_bytes": 2**70}) + + +def _hook(cache: DualCache, table: _Table) -> PROXY_LiteLLMManagedVectorStores: + return PROXY_LiteLLMManagedVectorStores( + InternalUsageCache(cache), + cast(PrismaClient, SimpleNamespace(db=SimpleNamespace(litellm_managedvectorstoretable=table))), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ttl", [1, 12.9, None, 0.5, 0, -1]) +async def test_vector_wire_matches_legacy_with_namespace_and_effective_ttl( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis, ttl: float | None +) -> None: + cache: Final = fake_redis_cache(server, namespace="tenant") + dual: Final = DualCache(redis_cache=cache, default_in_memory_ttl=ttl, default_redis_ttl=999) + dual.update_cache_ttl(ttl, None) + table: Final = _Table() + hook: Final = _hook(dual, table) + resource: Final = _resource() + unified_id: Final = hook.generate_unified_resource_id([resource], ["model-a"]) + assert ( + base64.urlsafe_b64decode(unified_id + "=" * (-len(unified_id) % 4)) + .decode() + .startswith("litellm_proxy:vector_store;unified_id,") + ) + record: Final = { + "unified_resource_id": unified_id, + "resource_object": resource, + "model_mappings": {"model-a": "vs-é"}, + "flat_model_resource_ids": ["vs-é"], + "created_by": "owner", + "team_id": "team-a", + "updated_by": "owner", + "extra": "文", + } + with patch("time.time", return_value=time.time()): + await hook.store_unified_resource_id( + unified_id, + resource, + None, + {"model-a": "vs-é"}, + UserAPIKeyAuth(user_id="owner", team_id="team-a"), + {"extra": "文"}, + ) + await cache.async_set_cache("legacy", record, **({"ttl": ttl} if ttl is not None else {})) + assert await raw.get(f"tenant:{unified_id}") == await raw.get("tenant:legacy") + assert await raw.pttl(f"tenant:{unified_id}") == await raw.pttl("tenant:legacy") + assert await raw.get(f"tenant:{unified_id}") == ( + None if ttl is not None and ttl < 1 else json.dumps(record).encode() + ) + assert await cache.async_get_cache(unified_id) == (None if ttl is not None and ttl < 1 else record) + assert await raw.get(f"tenant:tenant:{unified_id}") is None + assert table.row is not None + assert table.row.data == { + **record, + "resource_object": json.dumps(resource), + "model_mappings": json.dumps(record["model_mappings"]), + } + + +@pytest.mark.asyncio +async def test_old_vector_writer_new_reader_preserves_access_mapping_and_expiry( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + resource_id: Final = ( + base64.urlsafe_b64encode( + b"litellm_proxy:vector_store;unified_id,uuid;target_model_names,model;resource_id,vs-1;model_id,model" + ) + .decode() + .rstrip("=") + ) + record: Final = { + "unified_resource_id": resource_id, + "resource_object": _resource(), + "model_mappings": {"model": "vs-é"}, + "created_by": "owner", + "team_id": "team-a", + } + with patch("time.time", return_value=1_800_000_000): + cache: Final = fake_redis_cache(server, namespace="tenant") + dual: Final = DualCache(redis_cache=cache, default_in_memory_ttl=1) + table: Final = _Table(_Row({**record, "model_mappings": json.dumps(record["model_mappings"])})) + hook: Final = _hook(dual, table) + await cache.async_set_cache(resource_id, record, ttl=1) + assert await hook.get_unified_resource_id(resource_id) == record + assert await hook.can_user_access_unified_resource_id(resource_id, UserAPIKeyAuth(user_id="owner")) is True + assert await hook.can_user_access_unified_resource_id(resource_id, UserAPIKeyAuth(team_id="team-a")) is True + assert await hook.can_user_access_unified_resource_id(resource_id, UserAPIKeyAuth(team_id="team-b")) is False + assert await hook.get_model_resource_id_mapping([resource_id], None) == {resource_id: {"model": "vs-é"}} + assert table.reads == 0 + with patch("time.time", return_value=1_800_000_002): + assert await hook.get_model_resource_id_mapping([resource_id], None) == {resource_id: {"model": "vs-é"}} + assert table.reads == 1 + assert await raw.get(f"tenant:{resource_id}") is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("encoded", [b"null", b"{}", b"broken", b"\xff"]) +async def test_vector_null_or_malformed_value_falls_back_without_backfill( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis, encoded: bytes +) -> None: + row: Final = _Row({"unified_resource_id": "rid", "resource_object": _resource()}) + table: Final = _Table(row) + hook: Final = _hook(DualCache(redis_cache=fake_redis_cache(server)), table) + await raw.set("rid", encoded) + assert await hook.get_unified_resource_id("rid") == row.model_dump() + assert await raw.get("rid") == encoded + assert table.reads == 1 + + +@pytest.mark.asyncio +async def test_recursive_vector_metadata_retains_legacy_database_fallback( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + row: Final = _Row({"unified_resource_id": "rid", "resource_object": _resource()}) + table: Final = _Table(row) + cache: Final = fake_redis_cache(server) + hook: Final = _hook(DualCache(redis_cache=cache), table) + await raw.set("rid", b"[" * 10000 + b"0" + b"]" * 10000) + limit: Final = sys.getrecursionlimit() + try: + sys.setrecursionlimit(1000) + assert await cache.async_get_cache("rid") is None + assert await hook.get_unified_resource_id("rid") == row.model_dump() + assert table.reads == 1 + finally: + sys.setrecursionlimit(limit) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("encoded", [b"(1, 2)", b"{'item'}", b"{1: 'value'}"]) +async def test_python_literal_vector_value_retains_legacy_read_outcome( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis, encoded: bytes +) -> None: + cache: Final = fake_redis_cache(server) + table: Final = _Table(_Row({"unified_resource_id": "rid", "resource_object": _resource()})) + hook: Final = _hook(DualCache(redis_cache=cache), table) + await raw.set("rid", encoded) + assert await hook.get_unified_resource_id("rid") == await cache.async_get_cache("rid") + assert table.reads == 0 + + +@pytest.mark.asyncio +async def test_vector_invalidation_precedes_database_delete_and_keeps_null_ttl( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + row: Final = _Row({"unified_resource_id": "rid", "resource_object": _resource()}) + table: Final = _Table(row, raw) + hook: Final = _hook( + DualCache(redis_cache=fake_redis_cache(server, namespace="tenant"), default_in_memory_ttl=1), table + ) + with patch("time.time", return_value=time.time()): + assert await hook.delete_unified_resource_id("rid") == _resource() + assert await raw.get("tenant:rid") == b"null" + assert await raw.pttl("tenant:rid") == 1000 + assert await hook.get_unified_resource_id("rid") is None + assert table.reads == 2 + + +@pytest.mark.asyncio +async def test_vector_outage_keeps_local_and_database_writes(server: fakeredis.FakeServer) -> None: + cache: Final = fake_redis_cache(server) + dual: Final = DualCache(redis_cache=cache) + table: Final = _Table() + hook: Final = _hook(dual, table) + server.connected = False + await hook.store_unified_resource_id("rid", _resource(), None, {"model": "vs-é"}, UserAPIKeyAuth(user_id="owner")) + cached: Final = await hook.get_unified_resource_id("rid") + assert cached is not None + assert cached["resource_object"] == _resource() + assert table.reads == 0 + dual.in_memory_cache.flush_cache() + assert table.row is not None + for _ in range(8): + assert await hook.get_unified_resource_id("rid") == table.row.model_dump() + assert cache._circuit_breaker.state == "open" # pyright: ignore[reportPrivateUsage] # asserts the breaker opened after the outage + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ttl", [float("nan"), float("inf"), float("-inf")]) +async def test_invalid_vector_cache_ttl_does_not_block_database_write( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis, ttl: float +) -> None: + table: Final = _Table() + dual: Final = DualCache(redis_cache=fake_redis_cache(server), default_in_memory_ttl=ttl) + hook: Final = _hook(dual, table) + with patch("time.time", return_value=1_800_000_000): + await dual.async_set_cache("legacy", _resource()) + await hook.store_unified_resource_id( + "rid", _resource(), None, {"model": "vs-é"}, UserAPIKeyAuth(user_id="owner") + ) + assert table.row is not None + assert table.row.data["unified_resource_id"] == "rid" + assert (await raw.get("rid") is None) == (await raw.get("legacy") is None) + assert await raw.pttl("rid") == await raw.pttl("legacy") + + +@pytest.mark.asyncio +async def test_vector_unknown_applied_write_does_not_undo_cache_or_database( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + cache: Final = fake_redis_cache(server) + table: Final = _Table() + hook: Final = _hook(DualCache(redis_cache=cache, default_in_memory_ttl=1), table) + + async def applied_then_timeout(name: str, value: str, *, ex: timedelta | None, nx: bool) -> bool: + await raw.set(name, value, ex=ex, nx=nx) + raise RedisTimeoutError("reply lost after apply") + + with patch.object(cache.init_async_client(), "set", side_effect=applied_then_timeout): + await hook.store_unified_resource_id( + "rid", _resource(), None, {"model": "vs-é"}, UserAPIKeyAuth(user_id="owner") + ) + assert await raw.get("rid") == json.dumps(await hook.get_unified_resource_id("rid")).encode() + assert table.row is not None + + +@pytest.mark.asyncio +async def test_absent_vector_object_does_not_overwrite_cache( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + table: Final = _Table() + hook: Final = _hook(DualCache(redis_cache=fake_redis_cache(server)), table) + await raw.set("rid", b"previous") + await hook.store_unified_resource_id("rid", None, None, {"model": "vs-é"}, UserAPIKeyAuth()) + assert await raw.get("rid") == b"previous" + assert table.row is not None + assert "resource_object" not in table.row.data From 3e6c55da51096946b783d557df3aa7700aeb12d2 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 19:49:22 +0000 Subject: [PATCH 04/17] refactor(redis): isolate managed object metadata store --- .../proxy/hooks/managed_files.py | 25 ++- litellm/redis/stores/managed_objects.py | 32 +++ .../proxy/test_managed_files_hook.py | 11 +- .../unit/redis/stores/test_managed_objects.py | 209 ++++++++++++++++++ 4 files changed, 268 insertions(+), 9 deletions(-) create mode 100644 litellm/redis/stores/managed_objects.py create mode 100644 tests/unit/redis/stores/test_managed_objects.py diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index f5ac8718be6..d7f075bc9ae 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -95,6 +95,7 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attributio request_tags_from_metadata, ) from litellm.redis.stores.managed_files import ManagedFileStore, TransportManagedFileStore, file_key +from litellm.redis.stores.managed_objects import ManagedObjectStore, TransportManagedObjectStore, object_key from litellm.redis.transport.contract import RedisFailure from litellm.repositories.managed_file_repository import ManagedFileRepository from litellm.responses.utils import ResponsesAPIRequestUtils @@ -655,6 +656,24 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): await local_cache.async_set_cache(file_key(file_id), value, local_only=True) return value + async def _cache_object(self, unified_object_id: str, record: Mapping[str, object]) -> None: + cache: Final = self._metadata_cache() + local_cache: Final = cache.local + await local_cache.async_set_cache(object_key(unified_object_id), record, local_only=True) + if cache.redis is None: + return + store: Final[ManagedObjectStore] = TransportManagedObjectStore(cache.redis.transport()) + try: + seconds: Final = ( + int(cache.in_memory_ttl) + if cache.in_memory_ttl is not None and not math.isnan(cache.in_memory_ttl) + else cache.redis.default_ttl + ) + ttl: Final = timedelta(seconds=seconds) + except (OverflowError, ValueError): + return + await store.write_object(unified_object_id, record, ttl=ttl) + @staticmethod def _get_prometheus_logger(): """Find PrometheusLogger from litellm.callbacks, if registered.""" @@ -785,11 +804,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): file_purpose=file_purpose, file_object=stored_object, ) - await self.internal_usage_cache.async_set_cache( - key=unified_object_id, - value=litellm_managed_object.model_dump(), - litellm_parent_otel_span=litellm_parent_otel_span, - ) + await self._cache_object(unified_object_id, litellm_managed_object.model_dump()) from prisma import Json diff --git a/litellm/redis/stores/managed_objects.py b/litellm/redis/stores/managed_objects.py new file mode 100644 index 00000000000..f1b4c2d7579 --- /dev/null +++ b/litellm/redis/stores/managed_objects.py @@ -0,0 +1,32 @@ +import json +from collections.abc import Mapping +from dataclasses import dataclass +from datetime import timedelta +from typing import Final, Protocol + +from litellm.redis.transport.contract import RedisFailure, RedisTransport + + +def object_key(unified_object_id: str) -> str: + return unified_object_id + + +class ManagedObjectStore(Protocol): + async def write_object( + self, unified_object_id: str, record: Mapping[str, object], *, ttl: timedelta | None + ) -> None | RedisFailure: ... + + +@dataclass(frozen=True, slots=True) +class TransportManagedObjectStore: + transport: RedisTransport + + async def write_object( + self, unified_object_id: str, record: Mapping[str, object], *, ttl: timedelta | None + ) -> None | RedisFailure: + try: + encoded: Final = json.dumps(record) + except (TypeError, ValueError) as exc: + return RedisFailure(kind="decode", message=str(exc)) + result: Final = await self.transport.set_seconds(object_key(unified_object_id), encoded, ttl=ttl) + return result if isinstance(result, RedisFailure) else None diff --git a/tests/unit/enterprise/proxy/test_managed_files_hook.py b/tests/unit/enterprise/proxy/test_managed_files_hook.py index c0f6018469f..055583322a1 100644 --- a/tests/unit/enterprise/proxy/test_managed_files_hook.py +++ b/tests/unit/enterprise/proxy/test_managed_files_hook.py @@ -1380,8 +1380,9 @@ def _make_object_store_instance(): PROXY_LiteLLMManagedFiles, ) - mock_cache = MagicMock() - mock_cache.async_set_cache = AsyncMock() + from litellm.caching.dual_cache import DualCache + + cache = DualCache() mock_prisma = MagicMock() mock_prisma.db.litellm_managedobjecttable.upsert = AsyncMock() @@ -1389,7 +1390,7 @@ def _make_object_store_instance(): return ( PROXY_LiteLLMManagedFiles( - internal_usage_cache=mock_cache, + internal_usage_cache=cache, prisma_client=mock_prisma, ), mock_prisma, @@ -1498,7 +1499,9 @@ async def test_a_stored_background_response_names_its_creators_end_user_and_hold ) written = mock_prisma.db.litellm_managedobjecttable.upsert.await_args.kwargs["data"] - cached = managed_files.internal_usage_cache.async_set_cache.await_args.kwargs["value"] + cached: object = await managed_files.internal_usage_cache.async_get_cache( + "resp_unified", litellm_parent_otel_span=None + ) assert generated_text not in json.dumps(written, default=str) assert generated_text not in json.dumps(cached, default=str) assert json.loads(written["create"]["file_object"])["status"] == "completed" diff --git a/tests/unit/redis/stores/test_managed_objects.py b/tests/unit/redis/stores/test_managed_objects.py new file mode 100644 index 00000000000..05cd5e5d7b2 --- /dev/null +++ b/tests/unit/redis/stores/test_managed_objects.py @@ -0,0 +1,209 @@ +import json +import time +from collections.abc import Mapping +from datetime import timedelta +from types import SimpleNamespace +from typing import Final, Literal, cast +from unittest.mock import patch + +import fakeredis +import pytest +from litellm_enterprise.proxy.hooks.managed_files import PROXY_LiteLLMManagedFiles +from redis.exceptions import TimeoutError as RedisTimeoutError + +from litellm.caching.dual_cache import DualCache +from litellm.models.managed_files import LiteLLM_ManagedObjectTable +from litellm.proxy.utils import InternalUsageCache, PrismaClient +from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth +from litellm.types.utils import LiteLLMBatch, LiteLLMFineTuningJob +from tests._support.fake_redis import fake_redis_cache + + +class _Table: + def __init__(self, raw: fakeredis.FakeAsyncRedis | None = None, expected: bytes | None = None) -> None: + self.raw = raw + self.expected = expected + self.created: Mapping[str, object] | None = None + self.updated: Mapping[str, object] | None = None + + async def upsert(self, where: Mapping[str, object], data: Mapping[str, object]) -> None: + if self.raw is not None: + assert await self.raw.get(f"tenant:{where['unified_object_id']}") == self.expected + self.created = data + + async def update_many(self, where: Mapping[str, object], data: Mapping[str, object]) -> int: + if self.raw is not None: + assert await self.raw.get(f"tenant:{where['unified_object_id']}") == self.expected + self.updated = data + return 0 + + +def _hook(cache: DualCache, table: _Table) -> PROXY_LiteLLMManagedFiles: + db: Final = SimpleNamespace(litellm_managedobjecttable=table) + return PROXY_LiteLLMManagedFiles( + InternalUsageCache(cache), cast(PrismaClient, SimpleNamespace(db=db, writer_db=db)) + ) + + +def _object( + purpose: Literal["batch", "fine-tune", "response"], +) -> LiteLLMBatch | LiteLLMFineTuningJob | ResponsesAPIResponse: + if purpose == "batch": + return LiteLLMBatch( + id="batch-é", + completion_window="24h", + created_at=123, + endpoint="/v1/chat/completions", + input_file_id="input-文", + object="batch", + status="completed", + ) + if purpose == "fine-tune": + return LiteLLMFineTuningJob( + id="ft-é", + object="fine_tuning.job", + model="model", + created_at=123, + organization_id="org", + result_files=[], + status="validating_files", + training_file="file-文", + hyperparameters={"n_epochs": 8}, + seed=0, + ) + return ResponsesAPIResponse( + id="resp-é", + created_at=123, + status="completed", + output=[ + { + "type": "message", + "id": "msg", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "private generated output", "annotations": []}], + } + ], + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("purpose", ["batch", "fine-tune", "response"]) +@pytest.mark.parametrize("ttl", [1, 21.9, None, 0.5, 0, -1]) +async def test_object_wire_and_ttl_match_old_reader_and_write_precedes_database( + server: fakeredis.FakeServer, + raw: fakeredis.FakeAsyncRedis, + purpose: Literal["batch", "fine-tune", "response"], + ttl: float | None, +) -> None: + value: Final = _object(purpose) + sanitized: Final = value.model_copy(update={"output": []}) if isinstance(value, ResponsesAPIResponse) else value + record: Final = LiteLLM_ManagedObjectTable( + unified_object_id="object-é", + model_object_id=value.id, + file_purpose=purpose, + file_object=sanitized, + ).model_dump() + expected: Final = None if ttl is not None and ttl < 1 else json.dumps(record).encode() + cache: Final = fake_redis_cache(server, namespace="tenant") + table: Final = _Table(raw, expected) + dual: Final = DualCache(redis_cache=cache, default_in_memory_ttl=ttl, default_redis_ttl=999) + dual.update_cache_ttl(ttl, None) + hook: Final = _hook(dual, table) + with patch("time.time", return_value=time.time()): + await hook.store_unified_object_id("object-é", value, None, value.id, purpose, UserAPIKeyAuth(user_id="owner")) + await cache.async_set_cache("legacy", record, **({"ttl": ttl} if ttl is not None else {})) + assert await raw.get("tenant:object-é") == await raw.get("tenant:legacy") == expected + assert await raw.pttl("tenant:object-é") == await raw.pttl("tenant:legacy") + assert await cache.async_get_cache("object-é") == (None if ttl is not None and ttl < 1 else record) + assert await raw.get("tenant:tenant:object-é") is None + assert await dual.async_get_cache("object-é", local_only=True) == ( + None if ttl is not None and ttl < 0 else record + ) + assert table.created is not None + assert "private generated output" not in json.dumps(table.created, default=str) + if isinstance(value, ResponsesAPIResponse): + assert len(value.output) == 1 # pyright: ignore[reportUnknownArgumentType] # output carries a bare dict member + + +@pytest.mark.asyncio +async def test_object_update_only_still_writes_and_expires_cache( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + value: Final = _object("batch") + with patch("time.time", return_value=1_800_000_000): + cache: Final = fake_redis_cache(server) + dual: Final = DualCache(redis_cache=cache, default_in_memory_ttl=1) + table: Final = _Table() + hook: Final = _hook(dual, table) + await hook.store_unified_object_id( + "object", value, None, value.id, "batch", UserAPIKeyAuth(user_id="poller"), create_if_missing=False + ) + assert await raw.pttl("object") == 1000 + assert table.created is None + assert table.updated is not None + assert "created_by" not in table.updated + with patch("time.time", return_value=1_800_000_002): + assert await raw.get("object") is None + assert await dual.async_get_cache("object", local_only=True) is None + + +@pytest.mark.asyncio +async def test_object_outage_keeps_local_and_database_write(server: fakeredis.FakeServer) -> None: + cache: Final = fake_redis_cache(server) + dual: Final = DualCache(redis_cache=cache) + table: Final = _Table() + hook: Final = _hook(dual, table) + value: Final = _object("batch") + server.connected = False + for _ in range(8): + await hook.store_unified_object_id("object", value, None, value.id, "batch", UserAPIKeyAuth(user_id="owner")) + assert table.created is not None + assert ( + await dual.async_get_cache("object", local_only=True) + == LiteLLM_ManagedObjectTable( + unified_object_id="object", + model_object_id=value.id, + file_purpose="batch", + file_object=value, + ).model_dump() + ) + assert cache._circuit_breaker.state == "open" # pyright: ignore[reportPrivateUsage] # asserts the breaker opened after the outage + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ttl", [float("nan"), float("inf"), float("-inf")]) +async def test_invalid_object_cache_ttl_does_not_block_database_write( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis, ttl: float +) -> None: + table: Final = _Table() + dual: Final = DualCache(redis_cache=fake_redis_cache(server), default_in_memory_ttl=ttl) + hook: Final = _hook(dual, table) + value: Final = _object("batch") + with patch("time.time", return_value=1_800_000_000): + await dual.async_set_cache("legacy", value.model_dump()) + await hook.store_unified_object_id("object", value, None, value.id, "batch", UserAPIKeyAuth(user_id="owner")) + assert table.created is not None + assert (await raw.get("object") is None) == (await raw.get("legacy") is None) + assert await raw.pttl("object") == await raw.pttl("legacy") + + +@pytest.mark.asyncio +async def test_object_unknown_applied_write_is_not_undone( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + cache: Final = fake_redis_cache(server) + dual: Final = DualCache(redis_cache=cache, default_in_memory_ttl=1) + table: Final = _Table() + hook: Final = _hook(dual, table) + value: Final = _object("batch") + + async def applied_then_timeout(name: str, value: str, *, ex: timedelta | None, nx: bool) -> bool: + await raw.set(name, value, ex=ex, nx=nx) + raise RedisTimeoutError("reply lost after apply") + + with patch.object(cache.init_async_client(), "set", side_effect=applied_then_timeout): + await hook.store_unified_object_id("object", value, None, value.id, "batch", UserAPIKeyAuth(user_id="owner")) + assert await raw.get("object") == json.dumps(await dual.async_get_cache("object", local_only=True)).encode() + assert table.created is not None From 8c5f8d5db20658276902e60dff8d7264adaa74c4 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 20:21:18 +0000 Subject: [PATCH 05/17] refactor(redis): migrate user MCP permission links to typed store --- .../proxy/common_utils/user_api_key_cache.py | 79 +++++-- .../internal_user_endpoints.py | 6 +- .../mcp_server/auth/user_api_key_auth_mcp.py | 13 +- .../redis/stores/mcp_user_permission_links.py | 64 ++++++ .../test_internal_user_endpoints.py | 180 ++++++++-------- .../auth/test_user_api_key_auth_mcp.py | 49 ++++- .../stores/test_mcp_user_permission_links.py | 198 ++++++++++++++++++ 7 files changed, 484 insertions(+), 105 deletions(-) create mode 100644 litellm/redis/stores/mcp_user_permission_links.py create mode 100644 tests/unit/redis/stores/test_mcp_user_permission_links.py diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index e64fd755ef9..f86920ebd92 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -3,16 +3,28 @@ from __future__ import annotations import asyncio import re from collections.abc import Sequence +from contextlib import suppress +from datetime import timedelta from typing import TYPE_CHECKING, Any, Final, TypeVar, cast, overload -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.redis_batch import active_request_redis_batch from litellm.caching.redis_cache import RedisCache from litellm.constants import DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec +from litellm.redis.stores.mcp_user_permission_links import ( + USER_NO_MCP_PERMISSION_SENTINEL as _USER_NO_PERMISSION, +) +from litellm.redis.stores.mcp_user_permission_links import ( + McpUserPermissionLinkStore, + TransportMcpUserPermissionLinkStore, + user_permission_link_key, +) +from litellm.redis.transport.contract import RedisFailure if TYPE_CHECKING: from opentelemetry.trace import Span @@ -291,21 +303,64 @@ class UserApiKeyCache(DualCache): await super().async_set_cache_pipeline(cache_list=other_entries, local_only=local_only, **kwargs) -#: Value cached under ``user_object_permission_id_cache_key`` when the user links no permission row, -#: so a human without an entitlement costs no DB read per request. Lives beside the key builder -#: because it is part of the same cache protocol: a reader that knows the key must know this value. -USER_NO_MCP_PERMISSION_SENTINEL: Final = "__user_no_mcp_permission__" +_PERMISSION_LINK_VALUE: Final = TypeAdapter(object) +USER_NO_MCP_PERMISSION_SENTINEL: Final = _USER_NO_PERMISSION def user_object_permission_id_cache_key(user_id: str) -> str: - """Cache key for the ``user_id -> object_permission_id`` link. + return user_permission_link_key(user_id) - Lives here rather than next to either user because two modules own the two halves: the MCP auth - resolver writes it on read, and ``/user/update`` deletes it after changing the link. A key format - duplicated across those two drifts silently, and the failure is an entitlement change that never - takes effect. - """ - return f"user_object_permission_id:{user_id}" + +async def get_user_permission_link(cache: UserApiKeyCache, user_id: str) -> str | None: + key: Final = user_permission_link_key(user_id) + local: Final = _PERMISSION_LINK_VALUE.validate_python(await cache.async_get_cache(key=key, local_only=True)) + if local is not None: + return local if isinstance(local, str) else None + if cache.redis_cache is None: + return None + batch: Final = active_request_redis_batch(cache.redis_cache) + if batch is not None and batch.read_as_missing(key): + return None + with suppress(Exception): + store: Final[McpUserPermissionLinkStore] = TransportMcpUserPermissionLinkStore(cache.redis_cache.transport()) + result: Final = await store.get(user_id) + if isinstance(result, RedisFailure): + return None + if result is not None: + await cache.async_set_cache(key=key, value=result, local_only=True) + return result + return None + + +async def fill_user_permission_link( + cache: UserApiKeyCache, user_id: str, permission_id: str | None, *, ttl: float +) -> None: + await cache.async_set_cache( + key=user_permission_link_key(user_id), + value=permission_id or USER_NO_MCP_PERMISSION_SENTINEL, + ttl=ttl, + local_only=True, + ) + if cache.redis_cache is not None: + with suppress(Exception): + store: Final[McpUserPermissionLinkStore] = TransportMcpUserPermissionLinkStore( + cache.redis_cache.transport() + ) + _ = await store.fill(user_id, permission_id, ttl=timedelta(seconds=int(ttl))) + + +async def invalidate_user_permission_link(cache: UserApiKeyCache, user_id: str) -> None: + from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import publish_auth_cache_invalidation + + key: Final = user_permission_link_key(user_id) + cache.in_memory_cache.delete_cache(key) + if cache.redis_cache is not None: + with suppress(Exception): + store: Final[McpUserPermissionLinkStore] = TransportMcpUserPermissionLinkStore( + cache.redis_cache.transport() + ) + _ = await store.invalidate(user_id) + await publish_auth_cache_invalidation(cache_key=key) def object_permission_cache_key(object_permission_id: str) -> str: diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 54a2d4062d2..34b6e1480de 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -50,8 +50,8 @@ from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_ from litellm.proxy.common_utils.openai_error_payload import proxy_exception_for from litellm.proxy.common_utils.user_api_key_cache import ( AUTH_OBJECTS_TARGET, + invalidate_user_permission_link, object_permission_cache_key, - user_object_permission_id_cache_key, ) from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks @@ -1481,9 +1481,11 @@ async def _invalidate_cached_user_entitlement(user_id: str | None, object_permis keys: Final = ( *(object_permission_cache_key(permission_id) for permission_id in dict.fromkeys(object_permission_ids)), - *((user_object_permission_id_cache_key(user_id), user_id) if user_id is not None else ()), ) await evict_and_broadcast(cache_keys=keys, user_api_key_cache=user_api_key_cache) + if user_id is not None: + await invalidate_user_permission_link(user_api_key_cache, user_id) + await evict_and_broadcast(cache_keys=(user_id,), user_api_key_cache=user_api_key_cache) async def _update_single_user_helper( diff --git a/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py index db640b6d689..67a3f2fb590 100644 --- a/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py @@ -61,8 +61,9 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg from litellm.proxy.common_utils.user_api_key_cache import ( AUTH_OBJECTS_TARGET, USER_NO_MCP_PERMISSION_SENTINEL, + fill_user_permission_link, get_management_object_ttl, - user_object_permission_id_cache_key, + get_user_permission_link, ) from litellm.proxy.mcp_server.catalog import global_manager from litellm.proxy.mcp_server.oauth_utils import ( @@ -4149,9 +4150,8 @@ class MCPRequestHandler: from litellm.proxy.auth.auth_checks import get_user_object from litellm.proxy.proxy_server import user_api_key_cache - cache_key: Final = user_object_permission_id_cache_key(user_id) try: - cached: Final[object] = None if check_db_only else await user_api_key_cache.async_get_cache(key=cache_key) + cached: Final = None if check_db_only else await get_user_permission_link(user_api_key_cache, user_id) if cached == USER_NO_MCP_PERMISSION_SENTINEL: return None if isinstance(cached, str) and cached: @@ -4165,9 +4165,10 @@ class MCPRequestHandler: ) linked: Final[object] = getattr(user_row, "object_permission_id", None) if user_row is not None else None object_permission_id: Final = linked if isinstance(linked, str) and linked else None - await user_api_key_cache.async_set_cache( - key=cache_key, - value=object_permission_id or USER_NO_MCP_PERMISSION_SENTINEL, + await fill_user_permission_link( + user_api_key_cache, + user_id, + object_permission_id, ttl=get_management_object_ttl(user_api_key_cache), ) return object_permission_id diff --git a/litellm/redis/stores/mcp_user_permission_links.py b/litellm/redis/stores/mcp_user_permission_links.py new file mode 100644 index 00000000000..565325b6130 --- /dev/null +++ b/litellm/redis/stores/mcp_user_permission_links.py @@ -0,0 +1,64 @@ +import ast +import json +from dataclasses import dataclass +from datetime import timedelta +from typing import Final, Protocol + +from pydantic import TypeAdapter, ValidationError + +from litellm.redis.transport.contract import RedisFailure, RedisTransport + +USER_NO_MCP_PERMISSION_SENTINEL: Final = "__user_no_mcp_permission__" +_VALUE: Final = TypeAdapter(object) + + +def user_permission_link_key(user_id: str) -> str: + return f"user_object_permission_id:{user_id}" + + +class McpUserPermissionLinkStore(Protocol): + async def get(self, user_id: str) -> str | None | RedisFailure: ... + + async def fill(self, user_id: str, permission_id: str | None, *, ttl: timedelta) -> bool | None | RedisFailure: ... + + async def invalidate(self, user_id: str) -> int | RedisFailure: ... + + +@dataclass(frozen=True, slots=True) +class TransportMcpUserPermissionLinkStore: + transport: RedisTransport + + async def get(self, user_id: str) -> str | None | RedisFailure: + reply: Final = await self.transport.get(user_permission_link_key(user_id)) + if reply is None or isinstance(reply, RedisFailure): + return reply + try: + decoded: Final = reply.decode("utf-8") + try: + value: Final = _VALUE.validate_json(decoded) + except ValidationError: + literal: Final = _VALUE.validate_python(ast.literal_eval(decoded)) + return _permission_id(literal) + return _permission_id(value) + except (UnicodeDecodeError, ValueError, SyntaxError): + return RedisFailure(kind="decode", message="Invalid user permission link") + + async def fill(self, user_id: str, permission_id: str | None, *, ttl: timedelta) -> bool | None | RedisFailure: + seconds: Final = int(ttl.total_seconds()) + key: Final = user_permission_link_key(user_id) + value: Final = json.dumps(permission_id or USER_NO_MCP_PERMISSION_SENTINEL) + expiry: Final = timedelta(seconds=seconds) + batch: Final = self.transport.request_batch() + if batch is not None: + _ = batch.set_seconds(key, value, ttl=expiry) + return None + return await self.transport.set_seconds(key, value, ttl=expiry) + + async def invalidate(self, user_id: str) -> int | RedisFailure: + return await self.transport.delete(user_permission_link_key(user_id)) + + +def _permission_id(value: object) -> str | RedisFailure: + if isinstance(value, str): + return value + return RedisFailure(kind="decode", message="User permission link is not a string") diff --git a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py index 39586a1a204..299a0bf29a1 100644 --- a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py @@ -9,6 +9,7 @@ from types import SimpleNamespace from typing import TYPE_CHECKING, Final, cast from unittest.mock import AsyncMock, MagicMock +import fakeredis import httpx import pytest from fastapi import HTTPException @@ -41,6 +42,7 @@ from litellm.proxy.management_endpoints.internal_user_endpoints import ( from litellm.proxy.proxy_server import app from litellm.types.proxy.auth.auth_checks import UserNotFoundError from litellm.types.proxy.management_endpoints.internal_user_endpoints import InsensitiveContains +from tests._support.fake_redis import fake_redis_cache from tests.unit.proxy.management_endpoints.jwt_key_mapping_doubles import ( CascadingJWTMappingTable, JWTMappingRow, @@ -870,7 +872,6 @@ async def test_new_user_license_over_limit(mocker): """ Test that /user/new endpoint raises an error when license is over the user limit """ - from fastapi import HTTPException from litellm.proxy._types import NewUserRequest, UserAPIKeyAuth from litellm.proxy.management_endpoints.internal_user_endpoints import new_user @@ -1389,9 +1390,7 @@ async def test_user_info_url_encoding_plus_character(mocker): # Mock get_data to return user when called with user_id, empty list for keys async def mock_get_data(*args, **kwargs): - if kwargs.get("table_name") == "key": - return [] - elif kwargs.get("table_name") == "team": + if kwargs.get("table_name") == "key" or kwargs.get("table_name") == "team": return [] elif kwargs.get("user_id") is not None: return mock_user @@ -1524,7 +1523,7 @@ async def test_new_user_default_teams_flow(mocker): - Teams ARE sent to _add_user_to_team """ import litellm - from litellm.proxy._types import NewUserRequest, NewUserRequestTeam, UserAPIKeyAuth + from litellm.proxy._types import NewUserRequest, UserAPIKeyAuth from litellm.proxy.management_endpoints.internal_user_endpoints import new_user # Mock the prisma client @@ -1672,9 +1671,7 @@ async def test_new_user_clears_recent_missing_user_lookup(mocker: MockerFixture) ) mocker.patch( "litellm.proxy.management_endpoints.internal_user_endpoints.generate_key_helper_fn", - new=mocker.AsyncMock( - return_value={"user_id": user_id, "token": "sk-sso-created-user", "expires": None} - ), + new=mocker.AsyncMock(return_value={"user_id": user_id, "token": "sk-sso-created-user", "expires": None}), ) try: @@ -1884,8 +1881,6 @@ async def test_check_duplicate_user_email_case_insensitive(mocker): assert user_email_clause.get("mode") == "insensitive" assert user_email_clause.get("equals") == "newuser@example.com" - return None # No existing user found - mock_prisma_client.db.litellm_usertable.find_first = mock_find_first_no_duplicate # Should not raise any exception when no duplicate is found @@ -1938,7 +1933,6 @@ async def test_check_duplicate_user_id(mocker): user_id_clause = where_clause.get("user_id", {}) assert user_id_clause.get("equals") == "new-user-id" assert "mode" not in user_id_clause - return None mock_prisma_client.db.litellm_usertable.find_first = mock_find_first_no_duplicate @@ -2301,7 +2295,11 @@ def test_update_internal_user_params_preserves_model_budget_presence_and_neighbo "user_alias": "Spruce", } - replacement: Final = {"model-spruce": {"budget_limit": 0, "time_period": "1d"}, "model-birch": 5.0, "model-cedar": 0} + replacement: Final = { + "model-spruce": {"budget_limit": 0, "time_period": "1d"}, + "model-birch": 5.0, + "model-cedar": 0, + } request: Final = UpdateUserRequest( user_id="user-spruce", model_max_budget=replacement, @@ -2343,10 +2341,14 @@ async def test_user_model_budget_update_by_email_refreshes_cached_user(mocker: M prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=saved_user) prisma_client.get_data = mocker.AsyncMock(return_value=[saved_user]) prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": saved_user.user_id, "data": saved_user}) - mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency + mocker.patch( + "litellm.proxy.proxy_server.prisma_client", prisma_client + ) # test-quality-ok: substitute the database dependency cache: Final = UserApiKeyCache() await cache.async_set_cache(key=saved_user.user_id, value=saved_user, model_type=LiteLLM_UserTable) - mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache + mocker.patch( + "litellm.proxy.proxy_server.user_api_key_cache", cache + ) # test-quality-ok: exercise a real isolated cache broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", new_callable=mocker.AsyncMock, @@ -2382,10 +2384,14 @@ async def test_user_status_update_refreshes_cached_user( prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=saved_user) prisma_client.get_data = mocker.AsyncMock(return_value=[saved_user]) prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": saved_user.user_id, "data": saved_user}) - mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency + mocker.patch( + "litellm.proxy.proxy_server.prisma_client", prisma_client + ) # test-quality-ok: substitute the database dependency cache: Final = UserApiKeyCache() await cache.async_set_cache(key=saved_user.user_id, value=saved_user, model_type=LiteLLM_UserTable) - mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache + mocker.patch( + "litellm.proxy.proxy_server.user_api_key_cache", cache + ) # test-quality-ok: exercise a real isolated cache broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", new_callable=mocker.AsyncMock, @@ -2419,10 +2425,14 @@ async def test_bulk_user_model_budget_clear_serializes_and_refreshes_cache(mocke prisma_client: Final = mocker.MagicMock() prisma_client.db.litellm_usertable.find_many = mocker.AsyncMock(return_value=[saved_user]) prisma_client.db.litellm_usertable.update_many = mocker.AsyncMock(return_value=1) - mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency + mocker.patch( + "litellm.proxy.proxy_server.prisma_client", prisma_client + ) # test-quality-ok: substitute the database dependency cache: Final = UserApiKeyCache() await cache.async_set_cache(key=saved_user.user_id, value=saved_user, model_type=LiteLLM_UserTable) - mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache + mocker.patch( + "litellm.proxy.proxy_server.user_api_key_cache", cache + ) # test-quality-ok: exercise a real isolated cache broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", new_callable=mocker.AsyncMock, @@ -2454,7 +2464,9 @@ async def test_bulk_user_model_budget_clear_serializes_and_refreshes_cache(mocke @pytest.mark.asyncio @pytest.mark.parametrize("all_users", [False, True], ids=["single-user", "bulk-all-users"]) -async def test_user_max_budget_update_evicts_cached_user_on_every_worker(mocker: MockerFixture, all_users: bool) -> None: +async def test_user_max_budget_update_evicts_cached_user_on_every_worker( + mocker: MockerFixture, all_users: bool +) -> None: from litellm.proxy._types import LiteLLM_UserTable from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.management_endpoints.internal_user_endpoints import _update_single_user_helper, bulk_user_update @@ -2467,10 +2479,14 @@ async def test_user_max_budget_update_evicts_cached_user_on_every_worker(mocker: prisma_client.db.litellm_usertable.update_many = mocker.AsyncMock(return_value=1) prisma_client.get_data = mocker.AsyncMock(return_value=[saved_user]) prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": saved_user.user_id, "data": saved_user}) - mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency + mocker.patch( + "litellm.proxy.proxy_server.prisma_client", prisma_client + ) # test-quality-ok: substitute the database dependency cache: Final = UserApiKeyCache() await cache.async_set_cache(key=saved_user.user_id, value=saved_user, model_type=LiteLLM_UserTable) - mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache + mocker.patch( + "litellm.proxy.proxy_server.user_api_key_cache", cache + ) # test-quality-ok: exercise a real isolated cache broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", new_callable=mocker.AsyncMock, @@ -2875,24 +2891,12 @@ async def test_delete_user_cleans_up_created_by_invitation_links(mocker): mock_prisma_client.db.litellm_teamtable.find_many = mocker.AsyncMock(return_value=[]) # Mock all delete_many calls - mock_prisma_client.db.litellm_verificationtoken.find_many = mocker.AsyncMock( - return_value=[] - ) - mock_prisma_client.db.litellm_verificationtoken.delete_many = mocker.AsyncMock( - return_value=0 - ) - mock_prisma_client.db.litellm_invitationlink.delete_many = mocker.AsyncMock( - return_value=1 - ) - mock_prisma_client.db.litellm_organizationmembership.delete_many = mocker.AsyncMock( - return_value=0 - ) - mock_prisma_client.db.litellm_teammembership.delete_many = mocker.AsyncMock( - return_value=0 - ) - mock_prisma_client.db.litellm_usertable.delete_many = mocker.AsyncMock( - return_value=1 - ) + mock_prisma_client.db.litellm_verificationtoken.find_many = mocker.AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_verificationtoken.delete_many = mocker.AsyncMock(return_value=0) + mock_prisma_client.db.litellm_invitationlink.delete_many = mocker.AsyncMock(return_value=1) + mock_prisma_client.db.litellm_organizationmembership.delete_many = mocker.AsyncMock(return_value=0) + mock_prisma_client.db.litellm_teammembership.delete_many = mocker.AsyncMock(return_value=0) + mock_prisma_client.db.litellm_usertable.delete_many = mocker.AsyncMock(return_value=1) _serve_locked_user_delete(mock_prisma_client, mock_user_row) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) @@ -2985,9 +2989,15 @@ async def test_delete_user_evicts_jwt_key_mapping_cache_of_its_keys(mocker): mock_prisma_client.db.litellm_usertable.delete_many = mocker.AsyncMock(return_value=1) _serve_locked_user_delete(mock_prisma_client, user_row) - mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) # test-quality-ok: substitute the database dependency - mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache - mocker.patch("litellm.proxy.proxy_server.proxy_logging_obj", None) # test-quality-ok: delete_user reads it off proxy_server at call time + mocker.patch( + "litellm.proxy.proxy_server.prisma_client", mock_prisma_client + ) # test-quality-ok: substitute the database dependency + mocker.patch( + "litellm.proxy.proxy_server.user_api_key_cache", cache + ) # test-quality-ok: exercise a real isolated cache + mocker.patch( + "litellm.proxy.proxy_server.proxy_logging_obj", None + ) # test-quality-ok: delete_user reads it off proxy_server at call time await delete_user( data=DeleteUserRequest(user_ids=["jwt-user"]), @@ -3160,9 +3170,7 @@ async def test_user_info_v2_proxy_admin_can_query_any_user(mocker: MockerFixture return None mock_prisma_client: Final = _PrismaClient( - db=_Database( - litellm_usertable=_UserTable(find_unique=mocker.AsyncMock(side_effect=mock_find_unique)) - ) + db=_Database(litellm_usertable=_UserTable(find_unique=mocker.AsyncMock(side_effect=mock_find_unique))) ) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) @@ -3556,7 +3564,7 @@ async def test_user_info_v2_team_admin_can_query_team_member(mocker): """ from fastapi import Request - from litellm.proxy._types import LiteLLM_TeamTable, UserInfoV2Response + from litellm.proxy._types import UserInfoV2Response from litellm.proxy.management_endpoints.internal_user_endpoints import user_info_v2 mock_prisma_client = mocker.MagicMock() @@ -4419,8 +4427,7 @@ async def test_user_update_persists_mcp_entitlement_and_links_it(mocker): ) mock_prisma_client = _object_permission_mocks(mocker) - cache = mocker.MagicMock() - cache.async_delete_cache = mocker.AsyncMock() + cache: Final = UserApiKeyCache(redis_cache=fake_redis_cache(fakeredis.FakeServer(), namespace="entitlements")) mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) await _update_single_user_helper( @@ -4456,8 +4463,10 @@ async def test_user_update_invalidates_the_cached_entitlement(mocker): ) _object_permission_mocks(mocker) - cache = mocker.MagicMock() - cache.async_delete_cache = mocker.AsyncMock() + cache: Final = UserApiKeyCache(redis_cache=fake_redis_cache(fakeredis.FakeServer(), namespace="entitlements")) + keys: Final = ("object_permission_id:perm-new", "user_object_permission_id:target-user", "target-user") + for key in keys: + await cache.async_set_cache(key, "stale", ttl=60) mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) await _update_single_user_helper( @@ -4468,12 +4477,8 @@ async def test_user_update_invalidates_the_cached_entitlement(mocker): user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN), ) - deleted = {call.kwargs["key"] for call in cache.async_delete_cache.call_args_list} - assert deleted == { - "object_permission_id:perm-new", - "user_object_permission_id:target-user", - "target-user", - } + for key in keys: + assert await cache.async_get_cache(key) is None @pytest.mark.asyncio @@ -4483,9 +4488,10 @@ async def test_user_update_broadcasts_the_entitlement_invalidation_to_other_work ) _object_permission_mocks(mocker) - cache: Final = mocker.MagicMock() - cache.async_delete_cache = mocker.AsyncMock() - mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: substitute the cache dependency + cache: Final = UserApiKeyCache(redis_cache=fake_redis_cache(fakeredis.FakeServer(), namespace="entitlements")) + mocker.patch( + "litellm.proxy.proxy_server.user_api_key_cache", cache + ) # test-quality-ok: substitute the cache dependency broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", new_callable=mocker.AsyncMock, @@ -4520,8 +4526,10 @@ async def test_admin_can_clear_a_users_mcp_entitlement(mocker): ) mock_prisma_client = _object_permission_mocks(mocker, "perm-existing") - cache = mocker.MagicMock() - cache.async_delete_cache = mocker.AsyncMock() + cache: Final = UserApiKeyCache(redis_cache=fake_redis_cache(fakeredis.FakeServer(), namespace="entitlements")) + keys: Final = ("object_permission_id:perm-existing", "user_object_permission_id:target-user", "target-user") + for key in keys: + await cache.async_set_cache(key, "stale", ttl=60) mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) await _update_single_user_helper( @@ -4534,12 +4542,8 @@ async def test_admin_can_clear_a_users_mcp_entitlement(mocker): assert "object_permission" not in written mock_prisma_client.db.litellm_objectpermissiontable.upsert.assert_not_called() - deleted = {call.kwargs["key"] for call in cache.async_delete_cache.call_args_list} - assert deleted == { - "object_permission_id:perm-existing", - "user_object_permission_id:target-user", - "target-user", - } + for key in keys: + assert await cache.async_get_cache(key) is None @pytest.mark.asyncio @@ -4554,8 +4558,15 @@ async def test_user_update_invalidates_both_the_old_and_new_permission_rows(mock ) _object_permission_mocks(mocker, "perm-existing") - cache = mocker.MagicMock() - cache.async_delete_cache = mocker.AsyncMock() + cache: Final = UserApiKeyCache(redis_cache=fake_redis_cache(fakeredis.FakeServer(), namespace="entitlements")) + keys: Final = ( + "object_permission_id:perm-existing", + "object_permission_id:perm-new", + "user_object_permission_id:target-user", + "target-user", + ) + for key in keys: + await cache.async_set_cache(key, "stale", ttl=60) mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) await _update_single_user_helper( @@ -4566,13 +4577,8 @@ async def test_user_update_invalidates_both_the_old_and_new_permission_rows(mock user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN), ) - deleted = {call.kwargs["key"] for call in cache.async_delete_cache.call_args_list} - assert deleted == { - "object_permission_id:perm-existing", - "object_permission_id:perm-new", - "user_object_permission_id:target-user", - "target-user", - } + for key in keys: + assert await cache.async_get_cache(key) is None @pytest.mark.asyncio @@ -4755,7 +4761,7 @@ async def test_user_new_persists_model_max_budget(monkeypatch, model_max_budget, async def get_data(self, *args, **kwargs): return None - import litellm.proxy.proxy_server as proxy_server + from litellm.proxy import proxy_server monkeypatch.setattr(proxy_server, "prisma_client", _FakePrisma(), raising=False) # model_max_budget is an enterprise feature; without this the call is rejected @@ -5032,12 +5038,18 @@ async def test_delete_user_evicts_cached_user_rows(mocker: MockerFixture) -> Non prisma_client.db.litellm_teammembership.delete_many = mocker.AsyncMock(return_value=0) prisma_client.db.litellm_usertable.delete_many = mocker.AsyncMock(return_value=1) _serve_locked_user_delete(prisma_client, deleted) - mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency + mocker.patch( + "litellm.proxy.proxy_server.prisma_client", prisma_client + ) # test-quality-ok: substitute the database dependency cache: Final = UserApiKeyCache() for row in (deleted, survivor): await cache.async_set_cache(key=row.user_id, value=row, model_type=LiteLLM_UserTable) - mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache - mocker.patch("litellm.proxy.proxy_server.proxy_logging_obj", None) # test-quality-ok: delete_user reads it off proxy_server at call time + mocker.patch( + "litellm.proxy.proxy_server.user_api_key_cache", cache + ) # test-quality-ok: exercise a real isolated cache + mocker.patch( + "litellm.proxy.proxy_server.proxy_logging_obj", None + ) # test-quality-ok: delete_user reads it off proxy_server at call time broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", new_callable=mocker.AsyncMock, @@ -5143,7 +5155,9 @@ async def test_ui_view_users_answers_a_db_outage_as_503_no_db_connection_not_as_ assert raised.value.code == "503" assert raised.value.type == ProxyErrorTypes.no_db_connection assert isinstance(raised.value.__cause__, httpx.ConnectError) - outage_logs: Final = [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING and "ConnectError" in r.getMessage()] + outage_logs: Final = [ + r.getMessage() for r in caplog.records if r.levelno == logging.WARNING and "ConnectError" in r.getMessage() + ] assert outage_logs == ["Database unavailable during user search: ConnectError"] @@ -5301,7 +5315,9 @@ async def test_user_update_without_password_revokes_nothing(_admin_prisma, mocke @pytest.mark.asyncio -async def test_user_update_reports_a_bad_budget_duration_as_a_request_error_not_an_auth_error(_admin_prisma: MagicMock) -> None: +async def test_user_update_reports_a_bad_budget_duration_as_a_request_error_not_an_auth_error( + _admin_prisma: MagicMock, +) -> None: from litellm.proxy.management_endpoints.internal_user_endpoints import user_update with pytest.raises(ProxyException) as exc_info: diff --git a/tests/unit/proxy/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/unit/proxy/mcp_server/auth/test_user_api_key_auth_mcp.py index ee3ef3f5323..230a3c84790 100644 --- a/tests/unit/proxy/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/unit/proxy/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -8,6 +8,7 @@ from types import SimpleNamespace from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch +import fakeredis import pytest from fastapi import HTTPException from fastapi.testclient import TestClient @@ -26,6 +27,7 @@ from litellm.proxy._types import ( SpecialMCPServerNames, UserAPIKeyAuth, ) +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, UnloadableEntitlementError, @@ -33,6 +35,48 @@ from litellm.proxy.mcp_server.auth.user_api_key_auth_mcp import ( is_mcp_admitted_user_subject, ) from litellm.types.agents import AgentCaller +from tests._support.fake_redis import fake_redis_cache + + +@pytest.mark.asyncio +@pytest.mark.parametrize("permission_id", ["policy", None]) +async def test_user_permission_link_fills_from_cached_user_row_and_survives_redis_outage( + monkeypatch: pytest.MonkeyPatch, permission_id: str | None +) -> None: + server: Final = fakeredis.FakeServer() + redis: Final = fake_redis_cache(server, namespace="auth") + cache: Final = UserApiKeyCache(redis_cache=redis, default_in_memory_ttl=41) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache) + await cache.async_set_cache( + "link-user", LiteLLM_UserTable(user_id="link-user", object_permission_id=permission_id), local_only=True + ) + prisma: Final = MagicMock() + assert await MCPRequestHandler._user_object_permission_id("link-user", prisma) == permission_id + assert await redis.async_get_cache("user_object_permission_id:link-user") == ( + permission_id or "__user_no_mcp_permission__" + ) + cache.in_memory_cache.delete_cache("user_object_permission_id:link-user") + server.connected = False + assert await MCPRequestHandler._user_object_permission_id("link-user", prisma) == permission_id + assert await cache.async_get_cache("user_object_permission_id:link-user", local_only=True) == ( + permission_id or "__user_no_mcp_permission__" + ) + prisma.db.litellm_usertable.find_unique.assert_not_called() + + +@pytest.mark.asyncio +async def test_fresh_user_policy_does_not_use_redis_sentinel(monkeypatch: pytest.MonkeyPatch) -> None: + redis: Final = fake_redis_cache(fakeredis.FakeServer(), namespace="auth") + cache: Final = UserApiKeyCache(redis_cache=redis) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache) + await redis.async_set_cache("user_object_permission_id:link-user", "__user_no_mcp_permission__", ttl=60) + prisma: Final = MagicMock() + prisma.db.litellm_usertable.find_unique = AsyncMock(side_effect=RuntimeError("database unavailable")) + assert await MCPRequestHandler._user_object_permission_id("link-user", prisma) is None # pyright: ignore[reportPrivateUsage] # resolver internals under test + prisma.db.litellm_usertable.find_unique.assert_not_called() + with pytest.raises(HTTPException) as error: + await MCPRequestHandler._user_object_permission_id("link-user", prisma, check_db_only=True) # pyright: ignore[reportPrivateUsage] # resolver internals under test + assert error.value.status_code == 503 @pytest.mark.parametrize( @@ -11309,11 +11353,11 @@ async def test_fresh_mcp_user_permission_link_ignores_cached_and_replica_grants( database.writer_db.litellm_usertable.find_unique = AsyncMock(return_value=current) database.db.litellm_usertable.find_unique = AsyncMock(return_value=cached) monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) - assert await MCPRequestHandler._user_object_permission_id("fresh-human", database, check_db_only=True) == "current" + assert await MCPRequestHandler._user_object_permission_id("fresh-human", database, check_db_only=True) == "current" # pyright: ignore[reportPrivateUsage] # resolver internals under test database.db.litellm_usertable.find_unique.assert_not_awaited() database.writer_db.litellm_usertable.find_unique.side_effect = failure with pytest.raises(HTTPException) as denied: - await MCPRequestHandler._user_object_permission_id("fresh-human", database, check_db_only=True) + await MCPRequestHandler._user_object_permission_id("fresh-human", database, check_db_only=True) # pyright: ignore[reportPrivateUsage] # resolver internals under test assert denied.value.status_code == 503 @@ -12451,7 +12495,6 @@ async def test_personal_key_display_name_billing_uses_original_tool_grant( async def test_personal_key_display_name_billing_uses_allowed_servers( monkeypatch: pytest.MonkeyPatch, personal_key_owner_teams: UserAPIKeyAuth ) -> None: - from litellm.proxy import proxy_server from litellm.proxy.auth import auth_checks from litellm.proxy.mcp_server.mcp_server_manager import global_mcp_server_manager from litellm.types.mcp import MCPTransport diff --git a/tests/unit/redis/stores/test_mcp_user_permission_links.py b/tests/unit/redis/stores/test_mcp_user_permission_links.py new file mode 100644 index 00000000000..c669da07863 --- /dev/null +++ b/tests/unit/redis/stores/test_mcp_user_permission_links.py @@ -0,0 +1,198 @@ +from datetime import timedelta +from typing import Final + +import fakeredis +import pytest +import pytest_asyncio +from pytest import MonkeyPatch +from redis.exceptions import ConnectionError as RedisConnectionError + +from litellm.caching.redis_batch import active_request_redis_batch, request_redis_batch_scope +from litellm.caching.redis_cache import RedisCache +from litellm.proxy.common_utils.user_api_key_cache import ( + UserApiKeyCache, + fill_user_permission_link, + get_user_permission_link, + invalidate_user_permission_link, +) +from litellm.redis.stores.mcp_user_permission_links import TransportMcpUserPermissionLinkStore +from litellm.redis.transport.contract import RedisFailure +from tests._support.fake_redis import fake_redis_cache + + +@pytest_asyncio.fixture(loop_scope="function") +async def link_cache(monkeypatch: MonkeyPatch) -> tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache]: + monkeypatch.setattr("time.time", lambda: 1000.0) + server: Final = fakeredis.FakeServer() + redis: Final = fake_redis_cache(server, namespace="links") + cache: Final = UserApiKeyCache(redis_cache=redis, default_in_memory_ttl=37.9) + return server, redis, cache + + +@pytest.mark.asyncio +@pytest.mark.parametrize("value", ['policy é "quoted"', "__user_no_mcp_permission__", ""]) +async def test_old_writer_new_reader( + link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], value: str +) -> None: + _, redis, cache = link_cache + await redis.async_set_cache("user_object_permission_id:u", value, ttl=37.9) + assert await get_user_permission_link(cache, "u") == value + assert await cache.async_get_cache("user_object_permission_id:u", local_only=True) == value + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "value,wire", + [(None, b'"__user_no_mcp_permission__"'), ('policy é "quoted"', b'"policy \\u00e9 \\"quoted\\""')], +) +@pytest.mark.parametrize("ttl", [37.9, 37.999999999]) +async def test_new_writer_old_reader_and_whole_second_expiry( + link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], value: str | None, wire: bytes, ttl: float +) -> None: + server, redis, cache = link_cache + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await fill_user_permission_link(cache, "u", value, ttl=ttl) + assert await raw.get("links:user_object_permission_id:u") == wire + assert await raw.pttl("links:user_object_permission_id:u") == 37000 + assert await redis.async_get_cache("user_object_permission_id:u") == (value or "__user_no_mcp_permission__") + assert await raw.keys("*") == [b"links:user_object_permission_id:u"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("wire", [b"broken", b"42", b"null", b"{}", b"[]", b"\xff"]) +async def test_bad_payload_is_decode_failure_and_caller_cache_miss( + link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], wire: bytes +) -> None: + server, redis, cache = link_cache + await fakeredis.FakeAsyncRedis(server=server).set("links:user_object_permission_id:u", wire) + result: Final = await TransportMcpUserPermissionLinkStore(redis.transport()).get("u") + assert isinstance(result, RedisFailure) and result.kind == "decode" + assert await get_user_permission_link(cache, "u") is None + + +@pytest.mark.asyncio +async def test_literal_legacy_string_and_local_non_string( + link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], +) -> None: + server, _, cache = link_cache + await fakeredis.FakeAsyncRedis(server=server).set("links:user_object_permission_id:u", b"'legacy'") + assert await get_user_permission_link(cache, "u") == "legacy" + await cache.async_set_cache("user_object_permission_id:u", {"invalid": True}, local_only=True) + assert await get_user_permission_link(cache, "u") is None + + +@pytest.mark.asyncio +async def test_outage_preserves_local_fill_and_invalidation( + link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], +) -> None: + server, redis, cache = link_cache + server.connected = False + assert await get_user_permission_link(cache, "u") is None + await fill_user_permission_link(cache, "u", "policy", ttl=30) + assert await get_user_permission_link(cache, "u") == "policy" + await invalidate_user_permission_link(cache, "u") + assert await cache.async_get_cache("user_object_permission_id:u", local_only=True) is None + store: Final = TransportMcpUserPermissionLinkStore(redis.transport()) + for _ in range(10): + assert isinstance(await store.get("u"), RedisFailure) + result: Final = await store.get("u") + assert isinstance(result, RedisFailure) and result.kind == "breaker_open" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ttl", [37.9, 37.999999999]) +async def test_request_fill_stays_queued_and_clears_prior_miss( + link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], ttl: float +) -> None: + server, redis, cache = link_cache + raw: Final = fakeredis.FakeAsyncRedis(server=server) + with request_redis_batch_scope(): + batch: Final = active_request_redis_batch(redis) + assert batch is not None + assert await batch.mget(("user_object_permission_id:u",)) == {"user_object_permission_id:u": None} + await redis.async_set_cache("user_object_permission_id:u", "external", ttl=60) + assert await get_user_permission_link(cache, "u") is None + await fill_user_permission_link(cache, "u", "policy", ttl=ttl) + assert await get_user_permission_link(cache, "u") == "policy" + assert await raw.get("links:user_object_permission_id:u") == b'"external"' + await batch.flush() + assert await raw.get("links:user_object_permission_id:u") == b'"policy"' + assert await raw.pttl("links:user_object_permission_id:u") == 37000 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ttl", [0.9, 0.999999999, 0, -1]) +async def test_rejected_expiry_keeps_existing_redis_value_and_new_local_value( + link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], ttl: float +) -> None: + _, redis, cache = link_cache + await redis.async_set_cache("user_object_permission_id:u", "old", ttl=60) + await fill_user_permission_link(cache, "u", "new", ttl=ttl) + assert await cache.async_get_cache("user_object_permission_id:u", local_only=True) == ("new" if ttl >= 0 else None) + assert await redis.async_get_cache("user_object_permission_id:u") == "old" + result: Final = await TransportMcpUserPermissionLinkStore(redis.transport()).fill( + "u", "new", ttl=timedelta(seconds=int(ttl)) + ) + assert isinstance(result, RedisFailure) and result.kind == "server" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ttl", [0.999999999, 0, -1]) +async def test_rejected_batched_expiry_keeps_existing_value( + link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], ttl: float +) -> None: + _, redis, cache = link_cache + await redis.async_set_cache("user_object_permission_id:u", "old", ttl=60) + with request_redis_batch_scope() as scope: + await fill_user_permission_link(cache, "u", "new", ttl=ttl) + await scope.flush_all() + assert await redis.async_get_cache("user_object_permission_id:u") == "old" + + +@pytest.mark.asyncio +async def test_unclassified_client_failure_keeps_cache_optional( + link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], monkeypatch: MonkeyPatch +) -> None: + _, redis, cache = link_cache + client: Final = redis.init_async_client() + assert isinstance(client, fakeredis.FakeAsyncRedis) + await client.ping() + connection: Final = await client.connection_pool.get_connection() + read_response: Final = connection.read_response + + async def unavailable_reply(**kwargs: object) -> object: + await read_response(**kwargs) + raise RuntimeError("client callback failed") + + monkeypatch.setattr(connection, "read_response", unavailable_reply) + await client.connection_pool.release(connection) + assert await cache.async_get_cache("user_object_permission_id:u") is None + assert await get_user_permission_link(cache, "u") is None + await fill_user_permission_link(cache, "u", "policy", ttl=30) + assert await get_user_permission_link(cache, "u") == "policy" + await invalidate_user_permission_link(cache, "u") + assert await cache.async_get_cache("user_object_permission_id:u", local_only=True) is None + + +@pytest.mark.asyncio +async def test_lost_set_reply_does_not_roll_back_applied_write( + link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], monkeypatch: MonkeyPatch +) -> None: + server, redis, cache = link_cache + client: Final = redis.init_async_client() + assert isinstance(client, fakeredis.FakeAsyncRedis) + await client.ping() + connection: Final = await client.connection_pool.get_connection() + read_response: Final = connection.read_response + + async def lost_reply(**kwargs: object) -> object: + await read_response(**kwargs) + raise RedisConnectionError("reply lost after write") + + monkeypatch.setattr(connection, "read_response", lost_reply) + await client.connection_pool.release(connection) + await fill_user_permission_link(cache, "u", "applied", ttl=37.9) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + assert await raw.get("links:user_object_permission_id:u") == b'"applied"' + assert await raw.pttl("links:user_object_permission_id:u") == 37000 + assert await get_user_permission_link(cache, "u") == "applied" From c7b1c8b2696bb31ac60386fdc3a81c0b4aa95cb2 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 20:21:18 +0000 Subject: [PATCH 06/17] refactor(redis): migrate agent MCP permission links to typed store --- .../proxy/common_utils/user_api_key_cache.py | 44 ++++ .../mcp_server/auth/user_api_key_auth_mcp.py | 17 +- .../stores/mcp_agent_permission_links.py | 59 +++++ .../auth/test_user_api_key_auth_mcp.py | 39 ++++ .../stores/test_mcp_agent_permission_links.py | 209 ++++++++++++++++++ 5 files changed, 360 insertions(+), 8 deletions(-) create mode 100644 litellm/redis/stores/mcp_agent_permission_links.py create mode 100644 tests/unit/redis/stores/test_mcp_agent_permission_links.py diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index f86920ebd92..5e3a3946604 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -16,6 +16,12 @@ from litellm.caching.redis_batch import active_request_redis_batch from litellm.caching.redis_cache import RedisCache from litellm.constants import DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec +from litellm.redis.stores.mcp_agent_permission_links import ( + AGENT_NO_MCP_PERMISSION_SENTINEL, + McpAgentPermissionLinkStore, + TransportMcpAgentPermissionLinkStore, + agent_permission_link_key, +) from litellm.redis.stores.mcp_user_permission_links import ( USER_NO_MCP_PERMISSION_SENTINEL as _USER_NO_PERMISSION, ) @@ -368,6 +374,44 @@ def object_permission_cache_key(object_permission_id: str) -> str: return f"object_permission_id:{object_permission_id}" +async def get_agent_permission_link(cache: UserApiKeyCache, agent_id: str) -> str | None: + key: Final = agent_permission_link_key(agent_id) + local: Final = _PERMISSION_LINK_VALUE.validate_python(await cache.async_get_cache(key=key, local_only=True)) + if local is not None: + return local if isinstance(local, str) else None + if cache.redis_cache is None: + return None + batch: Final = active_request_redis_batch(cache.redis_cache) + if batch is not None and batch.read_as_missing(key): + return None + with suppress(Exception): + store: Final[McpAgentPermissionLinkStore] = TransportMcpAgentPermissionLinkStore(cache.redis_cache.transport()) + result: Final = await store.get(agent_id) + if isinstance(result, RedisFailure): + return None + if result is not None: + await cache.async_set_cache(key=key, value=result, local_only=True) + return result + return None + + +async def fill_agent_permission_link( + cache: UserApiKeyCache, agent_id: str, permission_id: str | None, *, ttl: float +) -> None: + await cache.async_set_cache( + key=agent_permission_link_key(agent_id), + value=permission_id or AGENT_NO_MCP_PERMISSION_SENTINEL, + ttl=ttl, + local_only=True, + ) + if cache.redis_cache is not None: + with suppress(Exception): + store: Final[McpAgentPermissionLinkStore] = TransportMcpAgentPermissionLinkStore( + cache.redis_cache.transport() + ) + _ = await store.fill(agent_id, permission_id, ttl=timedelta(seconds=int(ttl))) + + #: Cached under ``tag_registry_cache_key`` when the table exceeds ``TAG_REGISTRY_MAX_SIZE``: #: registry unusable, fall back to the per-tag lookup. TAG_REGISTRY_OVERFLOW_SENTINEL: Final = "__tag_registry_overflow__" diff --git a/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py index 67a3f2fb590..c671ad7d8c4 100644 --- a/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py @@ -61,7 +61,9 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg from litellm.proxy.common_utils.user_api_key_cache import ( AUTH_OBJECTS_TARGET, USER_NO_MCP_PERMISSION_SENTINEL, + fill_agent_permission_link, fill_user_permission_link, + get_agent_permission_link, get_management_object_ttl, get_user_permission_link, ) @@ -91,6 +93,7 @@ from litellm.proxy.mcp_server.utils import ( asgi_scope_fields, strip_known_server_prefix, ) +from litellm.redis.stores.mcp_agent_permission_links import AGENT_NO_MCP_PERMISSION_SENTINEL from litellm.repositories.table_repositories import ( AgentsRepository, MCPServerRepository, @@ -4447,9 +4450,7 @@ class MCPRequestHandler: return list(end_user_tools) return list(set(allowed_tools) & set(end_user_tools)) - # Sentinel stored in cache when an agent has no object_permission, so we - # don't re-query the DB on every MCP request for that agent. - _AGENT_NO_PERMISSION_SENTINEL = "__agent_no_mcp_permission__" + _AGENT_NO_PERMISSION_SENTINEL: Final = AGENT_NO_MCP_PERMISSION_SENTINEL @staticmethod @with_service_target(AUTH_OBJECTS_TARGET) @@ -4462,9 +4463,8 @@ class MCPRequestHandler: a link we DID resolve can make the caller deny.""" from litellm.proxy.proxy_server import user_api_key_cache - cache_key: Final = f"agent_object_permission_id:{agent_id}" try: - cached: Final = _CACHED_VALUE.validate_python(await user_api_key_cache.async_get_cache(key=cache_key)) + cached: Final = await get_agent_permission_link(user_api_key_cache, agent_id) if cached == MCPRequestHandler._AGENT_NO_PERMISSION_SENTINEL: return None if isinstance(cached, str) and cached: @@ -4472,9 +4472,10 @@ class MCPRequestHandler: agent_row: Final = await AgentsRepository(prisma_client).table.find_unique(where={"agent_id": agent_id}) linked: Final[object] = getattr(agent_row, "object_permission_id", None) if agent_row is not None else None object_permission_id: Final = linked if isinstance(linked, str) and linked else None - await user_api_key_cache.async_set_cache( - key=cache_key, - value=object_permission_id or MCPRequestHandler._AGENT_NO_PERMISSION_SENTINEL, + await fill_agent_permission_link( + user_api_key_cache, + agent_id, + object_permission_id, ttl=get_management_object_ttl(user_api_key_cache), ) return object_permission_id diff --git a/litellm/redis/stores/mcp_agent_permission_links.py b/litellm/redis/stores/mcp_agent_permission_links.py new file mode 100644 index 00000000000..fe28319614d --- /dev/null +++ b/litellm/redis/stores/mcp_agent_permission_links.py @@ -0,0 +1,59 @@ +import ast +import json +from dataclasses import dataclass +from datetime import timedelta +from typing import Final, Protocol + +from pydantic import TypeAdapter, ValidationError + +from litellm.redis.transport.contract import RedisFailure, RedisTransport + +AGENT_NO_MCP_PERMISSION_SENTINEL: Final = "__agent_no_mcp_permission__" +_VALUE: Final = TypeAdapter(object) + + +def agent_permission_link_key(agent_id: str) -> str: + return f"agent_object_permission_id:{agent_id}" + + +class McpAgentPermissionLinkStore(Protocol): + async def get(self, agent_id: str) -> str | None | RedisFailure: ... + + async def fill(self, agent_id: str, permission_id: str | None, *, ttl: timedelta) -> bool | None | RedisFailure: ... + + +@dataclass(frozen=True, slots=True) +class TransportMcpAgentPermissionLinkStore: + transport: RedisTransport + + async def get(self, agent_id: str) -> str | None | RedisFailure: + reply: Final = await self.transport.get(agent_permission_link_key(agent_id)) + if reply is None or isinstance(reply, RedisFailure): + return reply + try: + decoded: Final = reply.decode("utf-8") + try: + value: Final = _VALUE.validate_json(decoded) + except ValidationError: + literal: Final = _VALUE.validate_python(ast.literal_eval(decoded)) + return _permission_id(literal) + return _permission_id(value) + except (UnicodeDecodeError, ValueError, SyntaxError): + return RedisFailure(kind="decode", message="Invalid agent permission link") + + async def fill(self, agent_id: str, permission_id: str | None, *, ttl: timedelta) -> bool | None | RedisFailure: + seconds: Final = int(ttl.total_seconds()) + key: Final = agent_permission_link_key(agent_id) + value: Final = json.dumps(permission_id or AGENT_NO_MCP_PERMISSION_SENTINEL) + expiry: Final = timedelta(seconds=seconds) + batch: Final = self.transport.request_batch() + if batch is not None: + _ = batch.set_seconds(key, value, ttl=expiry) + return None + return await self.transport.set_seconds(key, value, ttl=expiry) + + +def _permission_id(value: object) -> str | RedisFailure: + if isinstance(value, str): + return value + return RedisFailure(kind="decode", message="Agent permission link is not a string") diff --git a/tests/unit/proxy/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/unit/proxy/mcp_server/auth/test_user_api_key_auth_mcp.py index 230a3c84790..b912b399058 100644 --- a/tests/unit/proxy/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/unit/proxy/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -79,6 +79,45 @@ async def test_fresh_user_policy_does_not_use_redis_sentinel(monkeypatch: pytest assert error.value.status_code == 503 +@pytest.mark.asyncio +@pytest.mark.parametrize("permission_id", ["agent-policy", None]) +async def test_agent_permission_link_fills_and_hits_redis_without_reloading_agent( + monkeypatch: pytest.MonkeyPatch, permission_id: str | None +) -> None: + server: Final = fakeredis.FakeServer() + redis: Final = fake_redis_cache(server, namespace="auth") + cache: Final = UserApiKeyCache(redis_cache=redis, default_in_memory_ttl=41) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache) + prisma: Final = MagicMock() + prisma.db.litellm_agentstable.find_unique = AsyncMock( + return_value=SimpleNamespace(object_permission_id=permission_id) + ) + assert await MCPRequestHandler._agent_object_permission_id("link-agent", prisma) == permission_id # pyright: ignore[reportPrivateUsage] # resolver internals under test + assert await redis.async_get_cache("agent_object_permission_id:link-agent") == ( + permission_id or "__agent_no_mcp_permission__" + ) + cache.in_memory_cache.delete_cache("agent_object_permission_id:link-agent") + assert await MCPRequestHandler._agent_object_permission_id("link-agent", prisma) == permission_id # pyright: ignore[reportPrivateUsage] # resolver internals under test + server.connected = False + assert await MCPRequestHandler._agent_object_permission_id("link-agent", prisma) == permission_id # pyright: ignore[reportPrivateUsage] # resolver internals under test + prisma.db.litellm_agentstable.find_unique.assert_awaited_once_with(where={"agent_id": "link-agent"}) + + +@pytest.mark.asyncio +async def test_unavailable_agent_policy_keeps_no_ceiling_without_negative_cache( + monkeypatch: pytest.MonkeyPatch, +) -> None: + server: Final = fakeredis.FakeServer() + cache: Final = UserApiKeyCache(redis_cache=fake_redis_cache(server, namespace="auth")) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache) + prisma: Final = MagicMock() + prisma.db.litellm_agentstable.find_unique = AsyncMock(side_effect=RuntimeError("database unavailable")) + server.connected = False + assert await MCPRequestHandler._agent_object_permission_id("link-agent", prisma) is None # pyright: ignore[reportPrivateUsage] # resolver internals under test + assert await cache.async_get_cache("agent_object_permission_id:link-agent", local_only=True) is None + prisma.db.litellm_agentstable.find_unique.assert_awaited_once_with(where={"agent_id": "link-agent"}) + + @pytest.mark.parametrize( "headers,expected", [ diff --git a/tests/unit/redis/stores/test_mcp_agent_permission_links.py b/tests/unit/redis/stores/test_mcp_agent_permission_links.py new file mode 100644 index 00000000000..0b02b50af2c --- /dev/null +++ b/tests/unit/redis/stores/test_mcp_agent_permission_links.py @@ -0,0 +1,209 @@ +from datetime import timedelta +from typing import Final + +import fakeredis +import pytest +import pytest_asyncio +from pytest import MonkeyPatch +from redis.exceptions import ConnectionError as RedisConnectionError + +from litellm.caching.redis_batch import active_request_redis_batch, request_redis_batch_scope +from litellm.caching.redis_cache import RedisCache +from litellm.proxy.common_utils.user_api_key_cache import ( + UserApiKeyCache, + fill_agent_permission_link, + get_agent_permission_link, +) +from litellm.redis.stores.mcp_agent_permission_links import TransportMcpAgentPermissionLinkStore +from litellm.redis.transport.contract import RedisFailure +from tests._support.fake_redis import fake_redis_cache + + +@pytest_asyncio.fixture(loop_scope="function") +async def link_cache(monkeypatch: MonkeyPatch) -> tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache]: + monkeypatch.setattr("time.time", lambda: 1000.0) + server: Final = fakeredis.FakeServer() + redis: Final = fake_redis_cache(server, namespace="links") + cache: Final = UserApiKeyCache(redis_cache=redis, default_in_memory_ttl=37.9) + return server, redis, cache + + +@pytest.mark.asyncio +@pytest.mark.parametrize("value", ['policy é "quoted"', "__agent_no_mcp_permission__", ""]) +async def test_old_writer_new_reader( + link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], value: str +) -> None: + _, redis, cache = link_cache + await redis.async_set_cache("agent_object_permission_id:u", value, ttl=37.9) + assert await get_agent_permission_link(cache, "u") == value + assert await cache.async_get_cache("agent_object_permission_id:u", local_only=True) == value + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "value,wire", + [(None, b'"__agent_no_mcp_permission__"'), ('policy é "quoted"', b'"policy \\u00e9 \\"quoted\\""')], +) +@pytest.mark.parametrize("ttl", [37.9, 37.999999999]) +async def test_new_writer_old_reader_and_whole_second_expiry( + link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], value: str | None, wire: bytes, ttl: float +) -> None: + server, redis, cache = link_cache + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await fill_agent_permission_link(cache, "u", value, ttl=ttl) + assert await raw.get("links:agent_object_permission_id:u") == wire + assert await raw.pttl("links:agent_object_permission_id:u") == 37000 + assert await redis.async_get_cache("agent_object_permission_id:u") == (value or "__agent_no_mcp_permission__") + assert await raw.keys("*") == [b"links:agent_object_permission_id:u"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("wire", [b"broken", b"42", b"null", b"{}", b"[]", b"\xff"]) +async def test_bad_payload_is_decode_failure_and_caller_cache_miss( + link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], wire: bytes +) -> None: + server, redis, cache = link_cache + await fakeredis.FakeAsyncRedis(server=server).set("links:agent_object_permission_id:u", wire) + result: Final = await TransportMcpAgentPermissionLinkStore(redis.transport()).get("u") + assert isinstance(result, RedisFailure) and result.kind == "decode" + assert await get_agent_permission_link(cache, "u") is None + + +@pytest.mark.asyncio +async def test_literal_legacy_string_and_local_non_string( + link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], +) -> None: + server, _, cache = link_cache + await fakeredis.FakeAsyncRedis(server=server).set("links:agent_object_permission_id:u", b"'legacy'") + assert await get_agent_permission_link(cache, "u") == "legacy" + await cache.async_set_cache("agent_object_permission_id:u", {"invalid": True}, local_only=True) + assert await get_agent_permission_link(cache, "u") is None + + +@pytest.mark.asyncio +async def test_agent_link_does_not_overwrite_user_link_or_permission_row( + link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], +) -> None: + server, redis, cache = link_cache + await redis.async_set_cache("user_object_permission_id:u", "user-policy", ttl=60) + await redis.async_set_cache("object_permission_id:u", {"mcp_servers": ["server"]}, ttl=60) + await fill_agent_permission_link(cache, "u", "agent-policy", ttl=37.9) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + assert await raw.mget( + "links:agent_object_permission_id:u", "links:user_object_permission_id:u", "links:object_permission_id:u" + ) == [b'"agent-policy"', b'"user-policy"', b'{"mcp_servers": ["server"]}'] + + +@pytest.mark.asyncio +async def test_outage_preserves_local_fill_and_open_breaker( + link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], +) -> None: + server, redis, cache = link_cache + server.connected = False + assert await get_agent_permission_link(cache, "u") is None + await fill_agent_permission_link(cache, "u", "policy", ttl=30) + assert await get_agent_permission_link(cache, "u") == "policy" + cache.in_memory_cache.delete_cache("agent_object_permission_id:u") + assert await cache.async_get_cache("agent_object_permission_id:u", local_only=True) is None + store: Final = TransportMcpAgentPermissionLinkStore(redis.transport()) + for _ in range(10): + assert isinstance(await store.get("u"), RedisFailure) + result: Final = await store.get("u") + assert isinstance(result, RedisFailure) and result.kind == "breaker_open" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ttl", [37.9, 37.999999999]) +async def test_request_fill_stays_queued_and_clears_prior_miss( + link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], ttl: float +) -> None: + server, redis, cache = link_cache + raw: Final = fakeredis.FakeAsyncRedis(server=server) + with request_redis_batch_scope(): + batch: Final = active_request_redis_batch(redis) + assert batch is not None + assert await batch.mget(("agent_object_permission_id:u",)) == {"agent_object_permission_id:u": None} + await redis.async_set_cache("agent_object_permission_id:u", "external", ttl=60) + assert await get_agent_permission_link(cache, "u") is None + await fill_agent_permission_link(cache, "u", "policy", ttl=ttl) + assert await get_agent_permission_link(cache, "u") == "policy" + assert await raw.get("links:agent_object_permission_id:u") == b'"external"' + await batch.flush() + assert await raw.get("links:agent_object_permission_id:u") == b'"policy"' + assert await raw.pttl("links:agent_object_permission_id:u") == 37000 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ttl", [0.9, 0.999999999, 0, -1]) +async def test_rejected_expiry_keeps_existing_redis_value_and_new_local_value( + link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], ttl: float +) -> None: + _, redis, cache = link_cache + await redis.async_set_cache("agent_object_permission_id:u", "old", ttl=60) + await fill_agent_permission_link(cache, "u", "new", ttl=ttl) + assert await cache.async_get_cache("agent_object_permission_id:u", local_only=True) == ("new" if ttl >= 0 else None) + assert await redis.async_get_cache("agent_object_permission_id:u") == "old" + result: Final = await TransportMcpAgentPermissionLinkStore(redis.transport()).fill( + "u", "new", ttl=timedelta(seconds=int(ttl)) + ) + assert isinstance(result, RedisFailure) and result.kind == "server" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ttl", [0.999999999, 0, -1]) +async def test_rejected_batched_expiry_keeps_existing_value( + link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], ttl: float +) -> None: + _, redis, cache = link_cache + await redis.async_set_cache("agent_object_permission_id:u", "old", ttl=60) + with request_redis_batch_scope() as scope: + await fill_agent_permission_link(cache, "u", "new", ttl=ttl) + await scope.flush_all() + assert await redis.async_get_cache("agent_object_permission_id:u") == "old" + + +@pytest.mark.asyncio +async def test_unclassified_client_failure_keeps_cache_optional( + link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], monkeypatch: MonkeyPatch +) -> None: + _, redis, cache = link_cache + client: Final = redis.init_async_client() + assert isinstance(client, fakeredis.FakeAsyncRedis) + await client.ping() + connection: Final = await client.connection_pool.get_connection() + read_response: Final = connection.read_response + + async def unavailable_reply(**kwargs: object) -> object: + await read_response(**kwargs) + raise RuntimeError("client callback failed") + + monkeypatch.setattr(connection, "read_response", unavailable_reply) + await client.connection_pool.release(connection) + assert await cache.async_get_cache("agent_object_permission_id:a") is None + assert await get_agent_permission_link(cache, "a") is None + await fill_agent_permission_link(cache, "a", "policy", ttl=30) + assert await get_agent_permission_link(cache, "a") == "policy" + + +@pytest.mark.asyncio +async def test_lost_set_reply_does_not_roll_back_applied_write( + link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], monkeypatch: MonkeyPatch +) -> None: + server, redis, cache = link_cache + client: Final = redis.init_async_client() + assert isinstance(client, fakeredis.FakeAsyncRedis) + await client.ping() + connection: Final = await client.connection_pool.get_connection() + read_response: Final = connection.read_response + + async def lost_reply(**kwargs: object) -> object: + await read_response(**kwargs) + raise RedisConnectionError("reply lost after write") + + monkeypatch.setattr(connection, "read_response", lost_reply) + await client.connection_pool.release(connection) + await fill_agent_permission_link(cache, "u", "applied", ttl=37.9) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + assert await raw.get("links:agent_object_permission_id:u") == b'"applied"' + assert await raw.pttl("links:agent_object_permission_id:u") == 37000 + assert await get_agent_permission_link(cache, "u") == "applied" From 371444b989e63112aa63bbfd4cc3476e81cccf0d Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 20:53:37 +0000 Subject: [PATCH 07/17] refactor(redis): move scheduler queue snapshots behind a typed store --- litellm/redis/stores/scheduler_queue.py | 71 +++++++++ litellm/scheduler.py | 55 +++++-- .../unit/redis/stores/test_scheduler_queue.py | 150 ++++++++++++++++++ tests/unit/test_scheduler.py | 51 ++---- 4 files changed, 277 insertions(+), 50 deletions(-) create mode 100644 litellm/redis/stores/scheduler_queue.py create mode 100644 tests/unit/redis/stores/test_scheduler_queue.py diff --git a/litellm/redis/stores/scheduler_queue.py b/litellm/redis/stores/scheduler_queue.py new file mode 100644 index 00000000000..3df2469eb6a --- /dev/null +++ b/litellm/redis/stores/scheduler_queue.py @@ -0,0 +1,71 @@ +import ast +import json +from collections.abc import Sequence +from dataclasses import dataclass +from datetime import timedelta +from typing import Final, Protocol, TypeAlias + +from pydantic import JsonValue, TypeAdapter, ValidationError + +from litellm.redis.transport.contract import RedisFailure, RedisTransport + +QUEUE_PREFIX: Final = "scheduler:queue" +QueueEntry: TypeAlias = tuple[int, str] +_ENTRIES: Final = TypeAdapter(list[QueueEntry]) +_JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) +_OBJECT: Final = TypeAdapter(object) + + +def queue_key(model_name: str) -> str: + return f"{QUEUE_PREFIX}:{model_name}" + + +def queue_entries(value: object) -> Sequence[QueueEntry]: + return _ENTRIES.validate_python(value) if isinstance(value, list) else [] + + +def _decode(raw: bytes) -> object: + decoded: Final = raw.decode("utf-8") + try: + return _JSON.validate_json(decoded) + except ValueError: + return _OBJECT.validate_python(ast.literal_eval(decoded)) + + +@dataclass(frozen=True, slots=True) +class InvalidQueue: + error: ValidationError + + +class SchedulerQueueStore(Protocol): + async def read_queue(self, model_name: str) -> tuple[QueueEntry, ...] | None | InvalidQueue | RedisFailure: ... + + async def replace_queue( + self, model_name: str, queue: Sequence[QueueEntry], *, ttl: timedelta | None + ) -> None | RedisFailure: ... + + +@dataclass(frozen=True, slots=True) +class TransportSchedulerQueueStore: + transport: RedisTransport + + async def read_queue(self, model_name: str) -> tuple[QueueEntry, ...] | None | InvalidQueue | RedisFailure: + raw: Final = await self.transport.get(queue_key(model_name)) + if raw is None or isinstance(raw, RedisFailure): + return raw + try: + value: Final = _decode(raw) + except (ValueError, SyntaxError, UnicodeError) as error: + return RedisFailure(kind="decode", message=str(error)) + if value is None: + return None + try: + return tuple(queue_entries(value)) + except ValidationError as error: + return InvalidQueue(error) + + async def replace_queue( + self, model_name: str, queue: Sequence[QueueEntry], *, ttl: timedelta | None + ) -> None | RedisFailure: + reply: Final = await self.transport.set_seconds(queue_key(model_name), json.dumps(list(queue)), ttl=ttl) + return reply if isinstance(reply, RedisFailure) else None diff --git a/litellm/scheduler.py b/litellm/scheduler.py index 73196d24e4d..a0bc8021c44 100644 --- a/litellm/scheduler.py +++ b/litellm/scheduler.py @@ -3,7 +3,8 @@ import enum import heapq import time from collections.abc import Awaitable, Callable, Sequence -from typing import Final, TypeAlias +from datetime import timedelta +from typing import Final from pydantic import TypeAdapter @@ -12,15 +13,24 @@ from litellm._internal_context import with_service_target from litellm.caching.caching import DualCache, RedisCache from litellm.constants import DEFAULT_IN_MEMORY_TTL, DEFAULT_POLLING_INTERVAL from litellm.exceptions import Timeout +from litellm.redis.stores.scheduler_queue import ( + QUEUE_PREFIX, + InvalidQueue, + QueueEntry, + SchedulerQueueStore, + TransportSchedulerQueueStore, + queue_entries, + queue_key, +) +from litellm.redis.transport.contract import RedisFailure from litellm.types.llms.base import LiteLLMBaseModel SCHEDULER_QUEUE_TARGET: Final = "scheduler_queue" -QueueEntry: TypeAlias = tuple[int, str] -_QUEUE_ENTRIES: Final = TypeAdapter(list[QueueEntry]) +_OBJECT: Final = TypeAdapter(object) class SchedulerCacheKeys(enum.Enum): - queue = "scheduler:queue" + queue = QUEUE_PREFIX default_in_memory_ttl = DEFAULT_IN_MEMORY_TTL # cache queue in-memory for 5s when redis cache available @@ -152,18 +162,43 @@ class Scheduler: (priority, request_id) tuple the heap operations compare against. """ if self.cache is not None: - _cache_key: Final = f"{SchedulerCacheKeys.queue.value}:{model_name}" - response: Final = await self.cache.async_get_cache(key=_cache_key) - if not isinstance(response, list): + key: Final = queue_key(model_name) + local: Final = _OBJECT.validate_python(await self.cache.async_get_cache(key=key, local_only=True)) + if isinstance(local, InvalidQueue): + raise local.error + if local is not None: + return list(queue_entries(local)) + redis: Final = self.cache.redis_cache + if redis is None: return [] - return _QUEUE_ENTRIES.validate_python(response) + store: Final[SchedulerQueueStore] = TransportSchedulerQueueStore(redis.transport()) + response: Final = await store.read_queue(model_name) + if isinstance(response, InvalidQueue): + await self._save_local_queue(response, model_name) + raise response.error + if response is None or isinstance(response, RedisFailure): + return [] + await self._save_local_queue(response, model_name) + return list(response) return self.queue + async def _save_local_queue(self, queue: Sequence[QueueEntry] | InvalidQueue, model_name: str) -> None: + value: Final = queue if isinstance(queue, InvalidQueue) else list(queue) + await self.cache.async_set_cache(key=queue_key(model_name), value=value, local_only=True) + @with_service_target(SCHEDULER_QUEUE_TARGET) async def save_queue(self, queue: list[QueueEntry], model_name: str) -> None: """ Save the updated queue of the model group """ if self.cache is not None: - _cache_key: Final = f"{SchedulerCacheKeys.queue.value}:{model_name}" - await self.cache.async_set_cache(key=_cache_key, value=queue) + await self._save_local_queue(queue, model_name) + redis: Final = self.cache.redis_cache + if redis is not None: + ttl: Final = ( + int(self.cache.default_in_memory_ttl) + if self.cache.default_in_memory_ttl is not None + else redis.default_ttl + ) + store: Final[SchedulerQueueStore] = TransportSchedulerQueueStore(redis.transport()) + await store.replace_queue(model_name, queue, ttl=timedelta(seconds=ttl)) diff --git a/tests/unit/redis/stores/test_scheduler_queue.py b/tests/unit/redis/stores/test_scheduler_queue.py new file mode 100644 index 00000000000..ac5f0a4842c --- /dev/null +++ b/tests/unit/redis/stores/test_scheduler_queue.py @@ -0,0 +1,150 @@ +from datetime import timedelta +from typing import Final, Protocol, cast +from unittest.mock import patch + +import fakeredis +import pytest +from pydantic import ValidationError +from redis.exceptions import ConnectionError as RedisConnectionError + +from litellm.redis.stores.scheduler_queue import TransportSchedulerQueueStore +from litellm.redis.transport.contract import Arg, RedisFailure +from litellm.scheduler import FlowItem, Scheduler +from tests._support.fake_redis import fake_redis_cache + + +class LegacyQueueCache(Protocol): + async def async_set_cache(self, key: str, value: object, *, ttl: int) -> bool | None: ... + + async def async_get_cache(self, key: str) -> object: ... + + +class RedisQueueWriter(Protocol): + async def set( + self, name: str, value: Arg, *, ex: timedelta | None = None, px: int | None = None, nx: bool = False + ) -> object: ... + + +@pytest.mark.asyncio +@pytest.mark.parametrize("seconds", [1.9, 0.9, 0, -0.9, -1.9]) +@patch("time.time", return_value=1800000000.0) +async def test_queue_fractional_seconds_match_legacy_set_expiry( + _clock: object, server: fakeredis.FakeServer, seconds: float +) -> None: + cache: Final = fake_redis_cache(server, namespace="tenant") + raw: Final = fakeredis.FakeAsyncRedis(server=server) + key: Final = "tenant:scheduler:queue:model" + await raw.set(key, b'[[2, "older"]]') + store: Final = TransportSchedulerQueueStore(cache.transport()) + result: Final = await store.replace_queue("model", ((1, "new"),), ttl=timedelta(seconds=seconds)) + if seconds >= 1: + assert result is None + assert await raw.get(key) == b'[[1, "new"]]' + assert await raw.pttl(key) == 1000 + else: + assert isinstance(result, RedisFailure) and result.kind == "server" + assert await raw.get(key) == b'[[2, "older"]]' + assert await raw.pttl(key) == -1 + + +@pytest.mark.asyncio +@patch("time.time", return_value=1800000000.0) +async def test_queue_rolling_deploy_preserves_bytes_ttl_and_local_snapshot( + _clock: object, server: fakeredis.FakeServer +) -> None: + cache: Final = fake_redis_cache(server, namespace="tenant") + raw: Final = fakeredis.FakeAsyncRedis(server=server) + scheduler: Final = Scheduler(redis_cache=cache) + legacy: Final = cast( + LegacyQueueCache, cache + ) # cast-ok: narrow the untyped legacy API for wire-compatibility probes + key: Final = "scheduler:queue:mødel" + with patch("time.time", return_value=1800000000.0): + await legacy.async_set_cache(key, [[2, "older-雪"]], ttl=5) + assert await scheduler.get_queue("mødel") == [(2, "older-雪")] + await scheduler.add_request(FlowItem(priority=1, request_id="new-雪", model_name="mødel")) + assert await raw.get("tenant:" + key) == b'[[1, "new-\\u96ea"], [2, "older-\\u96ea"]]' + assert await legacy.async_get_cache(key) == [[1, "new-雪"], [2, "older-雪"]] + assert await raw.pttl("tenant:" + key) == 5000 + await raw.delete("tenant:" + key) + assert await scheduler.get_queue("mødel") == [(1, "new-雪"), (2, "older-雪")] + scheduler.cache.in_memory_cache.flush_cache() + assert await scheduler.get_queue("mødel") == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("raw_value", [b"garbage", b"null", b"{}", b'"text"', b"\xff"]) +async def test_queue_unreadable_or_non_list_snapshot_is_empty(server: fakeredis.FakeServer, raw_value: bytes) -> None: + cache: Final = fake_redis_cache(server) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await raw.set("scheduler:queue:model", raw_value) + assert await Scheduler(redis_cache=cache).get_queue("model") == [] + + +@pytest.mark.asyncio +async def test_queue_invalid_entries_still_raise_and_python_literals_still_read(server: fakeredis.FakeServer) -> None: + cache: Final = fake_redis_cache(server) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await raw.set("scheduler:queue:valid", b"[(2, 'old')]") + await raw.set("scheduler:queue:invalid", b'[["bad", "id"]]') + scheduler: Final = Scheduler(redis_cache=cache) + assert await scheduler.get_queue("valid") == [(2, "old")] + with pytest.raises(ValidationError, match="int_parsing"): + await scheduler.get_queue("invalid") + await raw.set("scheduler:queue:invalid", b'[[1, "fixed"]]') + with pytest.raises(ValidationError, match="int_parsing"): + await scheduler.get_queue("invalid") + scheduler.cache.in_memory_cache.flush_cache() + assert await scheduler.get_queue("invalid") == [(1, "fixed")] + + +@pytest.mark.asyncio +async def test_queue_outage_keeps_local_writes_and_uses_shared_breaker() -> None: + server: Final = fakeredis.FakeServer() + cache: Final = fake_redis_cache(server) + scheduler: Final = Scheduler(redis_cache=cache) + server.connected = False + await scheduler.save_queue([(1, "local")], "model") + assert await scheduler.get_queue("model") == [(1, "local")] + scheduler.cache.in_memory_cache.flush_cache() + assert await scheduler.get_queue("model") == [] + for _ in range(5): + await scheduler.get_queue("model") + result: Final = await TransportSchedulerQueueStore(cache.transport()).read_queue("model") + assert isinstance(result, RedisFailure) and result.kind == "breaker_open" + + +@pytest.mark.asyncio +@patch("time.time", return_value=1800000000.0) +async def test_queue_late_redis_attachment_uses_backend_default_ttl( + _clock: object, server: fakeredis.FakeServer +) -> None: + scheduler: Final = Scheduler() + await scheduler.save_queue([(1, "local")], "model") + cache: Final = fake_redis_cache(server, namespace="tenant") + scheduler.cache.attach_redis_cache(cache) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + assert await raw.get("tenant:scheduler:queue:model") is None + with patch("time.time", return_value=1800000000.0): + await scheduler.save_queue([(2, "shared")], "model") + assert await raw.pttl("tenant:scheduler:queue:model") == cache.default_ttl * 1000 + + +@pytest.mark.asyncio +async def test_queue_unknown_applied_write_is_not_retried_or_rolled_back(server: fakeredis.FakeServer) -> None: + cache: Final = fake_redis_cache(server, namespace="tenant") + original_set: Final = cast(type[RedisQueueWriter], fakeredis.FakeAsyncRedis).set # cast-ok: typed redis-py boundary + raw: Final = fakeredis.FakeAsyncRedis(server=server) + scheduler: Final = Scheduler(redis_cache=cache) + + async def applied_then_disconnected( + name: str, value: Arg, *, ex: timedelta | None = None, px: int | None = None, nx: bool = False + ) -> None: + await original_set(raw, name, value, ex=ex, px=px, nx=nx) + raise RedisConnectionError("response lost after write") + + with patch.object(fakeredis.FakeAsyncRedis, "set", side_effect=applied_then_disconnected) as write: + await scheduler.save_queue([(1, "written")], "model") + assert write.call_count == 1 + assert await raw.get("tenant:scheduler:queue:model") == b'[[1, "written"]]' + assert await scheduler.get_queue("model") == [(1, "written")] diff --git a/tests/unit/test_scheduler.py b/tests/unit/test_scheduler.py index e0968e9b2cf..0583cbc74b8 100644 --- a/tests/unit/test_scheduler.py +++ b/tests/unit/test_scheduler.py @@ -3,17 +3,18 @@ import asyncio import importlib -import json import os from collections.abc import Sequence from typing import Final +import fakeredis import pytest import litellm from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.scheduler import FlowItem, Scheduler, SchedulerCacheKeys from litellm.utils import _invalidate_model_cost_lowercase_map +from tests._support.fake_redis import fake_redis_cache from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -35,17 +36,7 @@ async def test_scheduler_diff_model_names(): @pytest.mark.asyncio async def test_scheduler_poll_persists_queue_to_cache(): - class StubRedisCache: - def __init__(self): - self.store = {} - - async def async_get_cache(self, key, **kwargs): - return self.store.get(key) - - async def async_set_cache(self, key, value, **kwargs): - self.store[key] = value - - redis_cache = StubRedisCache() + redis_cache = fake_redis_cache(fakeredis.FakeServer()) scheduler = Scheduler(redis_cache=redis_cache) item1 = FlowItem(priority=0, request_id="10", model_name="gpt-3.5-turbo") @@ -56,7 +47,8 @@ async def test_scheduler_poll_persists_queue_to_cache(): await scheduler.poll(request=item1, health_deployments=[]) queue_key = f"{SchedulerCacheKeys.queue.value}:{item1.model_name}" - updated_queue = redis_cache.store[queue_key] + updated_queue = await redis_cache.async_get_cache(queue_key) + assert updated_queue is not None assert updated_queue[0][1] == "11" @@ -196,25 +188,16 @@ async def test_poll_during_cooldown_admits_an_erased_request_that_outranks_the_q assert await scheduler.get_queue("sched-model") == [(2, "still-queued")] -class _ExpiringCache: - def __init__(self) -> None: - self.store: dict[str, object] = {} - - async def async_get_cache(self, key: str, **kwargs: object) -> object: - return self.store.get(key) - - async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None: - self.store[key] = value - - @pytest.mark.asyncio async def test_wait_for_turn_during_cooldown_survives_the_queue_key_expiring(): - cache: Final = _ExpiringCache() + server: Final = fakeredis.FakeServer() + cache: Final = fake_redis_cache(server) + raw: Final = fakeredis.FakeAsyncRedis(server=server) scheduler: Final = Scheduler(redis_cache=cache) scheduler.cache.in_memory_cache.cache_dict.clear() async def no_healthy_deployments_after_the_key_expired() -> Sequence[object]: - cache.store.clear() + await raw.flushall() scheduler.cache.in_memory_cache.cache_dict.clear() return () @@ -227,22 +210,9 @@ async def test_wait_for_turn_during_cooldown_survives_the_queue_key_expiring(): assert await scheduler.get_queue("sched-model") == [] - -class _JsonRoundTripRedisCache: - def __init__(self) -> None: - self.store: dict[str, str] = {} - - async def async_get_cache(self, key: str, **kwargs: object) -> object: - raw: Final = self.store.get(key) - return None if raw is None else json.loads(raw) - - async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None: - self.store[key] = json.dumps(value) - - @pytest.mark.asyncio async def test_second_replica_enqueues_behind_a_queue_decoded_from_redis(): - redis_cache: Final = _JsonRoundTripRedisCache() + redis_cache: Final = fake_redis_cache(fakeredis.FakeServer()) replica_a: Final = Scheduler(redis_cache=redis_cache) replica_b: Final = Scheduler(redis_cache=redis_cache) waiting_on_a: Final = FlowItem(priority=1, request_id="waiting-on-a", model_name="sched-model") @@ -257,6 +227,7 @@ async def test_second_replica_enqueues_behind_a_queue_decoded_from_redis(): await replica_b.remove_request(request_id="waiting-on-a", model_name="sched-model") assert await replica_b.get_queue("sched-model") == [] + class _PausesAfterEnqueueScheduler(Scheduler): def __init__(self) -> None: super().__init__() From b82443abc758cc6afb993bbd6904f1731d16d2e1 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 20:53:38 +0000 Subject: [PATCH 08/17] refactor(redis): move global router usage behind a typed store --- litellm/caching/dual_cache.py | 3 + litellm/redis/stores/router_global_usage.py | 111 +++++++++ litellm/router.py | 159 +++++++------ litellm/types/router.py | 6 +- .../router_code_coverage.py | 2 + .../redis/stores/test_router_global_usage.py | 215 ++++++++++++++++++ tests/unit/test_internal_context.py | 2 +- .../test_router/test_router_helper_utils.py | 103 +++------ 8 files changed, 461 insertions(+), 140 deletions(-) create mode 100644 litellm/redis/stores/router_global_usage.py create mode 100644 tests/unit/redis/stores/test_router_global_usage.py diff --git a/litellm/caching/dual_cache.py b/litellm/caching/dual_cache.py index e6214a3650c..d929953f3c1 100644 --- a/litellm/caching/dual_cache.py +++ b/litellm/caching/dual_cache.py @@ -336,6 +336,9 @@ class DualCache(BaseCache): else: self.last_redis_batch_access_time[key] = previous_time + def rollback_redis_batch_reads(self, previous_access_times: Mapping[str, float | None]) -> None: + self._rollback_redis_batch_key_reservations(dict(previous_access_times)) + async def _prepare_batch_get( self, keys: list[str], local_only: bool, throttle_redis: bool = True, **kwargs: object ) -> PendingBatchRead: diff --git a/litellm/redis/stores/router_global_usage.py b/litellm/redis/stores/router_global_usage.py new file mode 100644 index 00000000000..8471cba1a54 --- /dev/null +++ b/litellm/redis/stores/router_global_usage.py @@ -0,0 +1,111 @@ +import ast +from collections.abc import Sequence +from dataclasses import dataclass +from datetime import timedelta +from typing import Final, Literal, Protocol + +from pydantic import JsonValue, TypeAdapter + +from litellm.redis.transport.contract import RedisFailure, RedisTransport + +TPM_KEY_TEMPLATE: Final = "global_router:{id}:{model}:tpm:{current_minute}" +RPM_KEY_TEMPLATE: Final = "global_router:{id}:{model}:rpm:{current_minute}" +_JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) + + +@dataclass(frozen=True, slots=True) +class GlobalUsageCounter: + deployment_id: str + model: str | None + minute: str + metric: Literal["tpm", "rpm"] + + +def usage_key(counter: GlobalUsageCounter) -> str: + template: Final = TPM_KEY_TEMPLATE if counter.metric == "tpm" else RPM_KEY_TEMPLATE + return template.format(id=counter.deployment_id, model=counter.model, current_minute=counter.minute) + + +def _decode(raw: bytes | None) -> JsonValue: + if raw is None: + return None + text: Final = raw.decode("utf-8") + try: + return _JSON.validate_json(text) + except ValueError: + return _JSON.validate_python(ast.literal_eval(text)) + + +@dataclass(frozen=True, slots=True) +class GlobalUsageRead: + values: tuple[JsonValue | bytes, ...] + failure: RedisFailure | None = None + + +class RouterGlobalUsageStore(Protocol): + async def read_global_usage(self, counters: Sequence[GlobalUsageCounter]) -> GlobalUsageRead | RedisFailure: ... + + async def record_global_usage( + self, deployment_id: str, model: str | None, minute: str, *, tokens: float, requests: int, ttl: timedelta + ) -> tuple[float, float] | RedisFailure: ... + + async def record_global_failure( + self, deployment_id: str, model: str | None, minute: str, *, ttl: timedelta + ) -> float | RedisFailure: ... + + +@dataclass(frozen=True, slots=True) +class TransportRouterGlobalUsageStore: + transport: RedisTransport + + async def read_global_usage(self, counters: Sequence[GlobalUsageCounter]) -> GlobalUsageRead | RedisFailure: + raw: Final = await self.transport.mget(tuple(usage_key(counter) for counter in counters)) + if isinstance(raw, RedisFailure): + return raw + try: + return GlobalUsageRead(tuple(_decode(value) for value in raw)) + except (ValueError, SyntaxError, UnicodeError) as error: + return GlobalUsageRead(raw, RedisFailure(kind="decode", message=str(error))) + + async def record_global_usage( + self, deployment_id: str, model: str | None, minute: str, *, tokens: float, requests: int, ttl: timedelta + ) -> tuple[float, float] | RedisFailure: + batch: Final = self.transport.post_call_batch() or self.transport.batch() + tpm: Final = usage_key(GlobalUsageCounter(deployment_id, model, minute, "tpm")) + rpm: Final = usage_key(GlobalUsageCounter(deployment_id, model, minute, "rpm")) + token_result: Final = batch.incrbyfloat(tpm, tokens) + token_expiry: Final = batch.expire_seconds(tpm, ttl) + request_result: Final = batch.incrbyfloat(rpm, requests) + request_expiry: Final = batch.expire_seconds(rpm, ttl) + token_value: Final = await token_result + token_expired: Final = await token_expiry + request_value: Final = await request_result + request_expired: Final = await request_expiry + if isinstance(token_value, RedisFailure): + return token_value + if isinstance(token_expired, RedisFailure): + return token_expired + if isinstance(request_value, RedisFailure): + return request_value + if isinstance(request_expired, RedisFailure): + return request_expired + return token_value, request_value + + async def record_global_failure( + self, deployment_id: str, model: str | None, minute: str, *, ttl: timedelta + ) -> float | RedisFailure: + key: Final = usage_key(GlobalUsageCounter(deployment_id, model, minute, "rpm")) + batch: Final = self.transport.batch() + increment: Final = batch.incrbyfloat(key, 1) + expiry: Final = batch.ttl(key) + value: Final = await increment + remaining: Final = await expiry + if isinstance(value, RedisFailure): + return value + if isinstance(remaining, RedisFailure): + return remaining + if remaining == -1: + expired: Final = await self.transport.expire_seconds(key, ttl) + if isinstance(expired, RedisFailure): + return expired + return value diff --git a/litellm/router.py b/litellm/router.py index 233e7581a40..1809f861ad8 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -58,7 +58,7 @@ from collections.abc import ( ) from contextvars import ContextVar from dataclasses import dataclass -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from functools import lru_cache, partial from itertools import chain from types import MappingProxyType @@ -182,6 +182,13 @@ from litellm.llms.openai_like.model_info import ( MODEL_INFO_REFRESH_SECONDS, get_openai_compatible_model_info, ) +from litellm.redis.stores.router_global_usage import ( + GlobalUsageCounter, + RouterGlobalUsageStore, + TransportRouterGlobalUsageStore, + usage_key, +) +from litellm.redis.transport.contract import RedisFailure from litellm.responses.bridge_check import keep_registered_mode from litellm.responses.retrieve_stream_query import requests_streamed_retrieval, without_retrieve_stream_defaults from litellm.router_api import ( @@ -683,6 +690,7 @@ _DEPLOYMENT_PARAMS_ADAPTER: Final[TypeAdapter[Mapping[str, object] | GenericLite Mapping[str, object] | GenericLiteLLMParams ) _OPAQUE_VALUE_ADAPTER: Final = TypeAdapter(object) +_GLOBAL_USAGE_LOCAL_VALUES: Final[TypeAdapter[list[object] | None]] = TypeAdapter(list[object] | None) _DEPLOYMENT_SELECTED_EVENT: Final = "litellm.request.deployment_selected" _CONTENT_BLOCKS_ADAPTER: Final = TypeAdapter(list[dict[str, object]]) _DEPLOYMENTS_ADAPTER: Final = TypeAdapter(Sequence[Mapping[str, object]]) @@ -9861,25 +9869,34 @@ class Router(DeprecatedRouterMembers): return None current_minute: Final = get_utc_datetime().strftime("%H-%M") # use the same timezone regardless of system clock - tpm_key: Final = RouterCacheEnum.TPM.value.format( - id=deployment_id, current_minute=current_minute, model=deployment_name - ) - rpm_key: Final = RouterCacheEnum.RPM.value.format( - id=deployment_id, current_minute=current_minute, model=deployment_name - ) + tpm_key: Final = usage_key(GlobalUsageCounter(deployment_id, deployment_name, current_minute, "tpm")) + rpm_key: Final = usage_key(GlobalUsageCounter(deployment_id, deployment_name, current_minute, "rpm")) pipeline_operations: Final[list[RedisPipelineIncrementOperation]] = [ RedisPipelineIncrementOperation(key=key, increment_value=increment_value, ttl=RoutingArgs.ttl.value) for key, increment_value in ((tpm_key, total_tokens), (rpm_key, rpm_increment)) ] - post_increment_values: Final = await self._core_cache.async_increment_cache_pipeline( - increment_list=pipeline_operations, - parent_otel_span=parent_otel_span, + local_values: Final = await self._core_cache.async_increment_cache_pipeline( + increment_list=pipeline_operations, local_only=True, parent_otel_span=parent_otel_span ) - if post_increment_values is not None and self._core_cache.redis_cache is not None: - for operation, value in zip(pipeline_operations, post_increment_values): - await self._core_cache.async_set_cache( - operation["key"], int(value), local_only=True, ttl=RoutingArgs.ttl.value - ) + redis: Final = self._core_cache.redis_cache + if redis is not None: + store: Final[RouterGlobalUsageStore] = TransportRouterGlobalUsageStore(redis.transport()) + shared_values: Final = await store.record_global_usage( + deployment_id, + deployment_name, + current_minute, + tokens=total_tokens, + requests=rpm_increment, + ttl=timedelta(seconds=RoutingArgs.ttl.value), + ) + post_increment_values: Final = ( + (None if shared_values.kind == "breaker_open" else local_values) + if isinstance(shared_values, RedisFailure) + else shared_values + ) + if post_increment_values is not None: + for operation, value in zip(pipeline_operations, post_increment_values): + await self._set_global_usage_local(operation["key"], int(value), ttl=RoutingArgs.ttl.value) return tpm_key def _deployment_tracks_usage(self, deployment_id: str, model_group: str) -> bool: @@ -10098,15 +10115,26 @@ class Router(DeprecatedRouterMembers): dt: Final = get_utc_datetime() current_minute: Final = dt.strftime("%H-%M") # use the same timezone regardless of system clock - ## RPM - rpm_key: Final = RouterCacheEnum.RPM.value.format(id=id, current_minute=current_minute, model=deployment_name) + counter: Final = GlobalUsageCounter( + str(_OPAQUE_VALUE_ADAPTER.validate_python(id)), + None if deployment_name is None else str(_OPAQUE_VALUE_ADAPTER.validate_python(deployment_name)), + current_minute, + "rpm", + ) with service_target(ROUTER_USAGE_TARGET): await self._core_cache.async_increment_cache( - key=rpm_key, + key=usage_key(counter), value=1, + local_only=True, parent_otel_span=parent_otel_span, ttl=RoutingArgs.ttl.value, ) + redis: Final = self._core_cache.redis_cache + if redis is not None: + store: Final[RouterGlobalUsageStore] = TransportRouterGlobalUsageStore(redis.transport()) + await store.record_global_failure( + counter.deployment_id, counter.model, counter.minute, ttl=timedelta(seconds=RoutingArgs.ttl.value) + ) def _get_metadata_variable_name_from_kwargs(self, kwargs: dict) -> Literal["metadata", "litellm_metadata"]: """ @@ -13204,60 +13232,59 @@ class Router(DeprecatedRouterMembers): """ dt: Final = get_utc_datetime() current_minute: Final = dt.strftime("%H-%M") # use the same timezone regardless of system clock - tpm_keys: Final[list[str]] = [] - rpm_keys: Final[list[str]] = [] - model_list: Final = self._core_get_model_list(model_name=model_group) - if model_list is None: # no matching deployments + if model_list is None: return None, None - - for model in model_list: - id: str | None = model.get("model_info", {}).get("id") - litellm_model: str | None = model["litellm_params"].get( - "model" - ) # USE THE MODEL SENT TO litellm.completion() - consistent with how global_router cache is written. - if id is None or litellm_model is None: - continue - tpm_keys.append( - RouterCacheEnum.TPM.value.format( - id=id, - model=litellm_model, - current_minute=current_minute, - ) + identities: Final = tuple( + ( + _OPAQUE_VALUE_ADAPTER.validate_python(model.get("model_info", {}).get("id")), + _OPAQUE_VALUE_ADAPTER.validate_python(model["litellm_params"].get("model")), ) - rpm_keys.append( - RouterCacheEnum.RPM.value.format( - id=id, - model=litellm_model, - current_minute=current_minute, - ) - ) - combined_tpm_rpm_keys: Final = tpm_keys + rpm_keys - - combined_tpm_rpm_values: Final = await self._core_cache.async_batch_get_cache(keys=combined_tpm_rpm_keys) - if combined_tpm_rpm_values is None: + for model in model_list + ) + tpm_counters: Final = tuple( + GlobalUsageCounter(str(deployment_id), model_name, current_minute, "tpm") + for deployment_id, model_name in identities + if deployment_id is not None and isinstance(model_name, str) + ) + rpm_counters: Final = tuple( + GlobalUsageCounter(counter.deployment_id, counter.model, counter.minute, "rpm") for counter in tpm_counters + ) + counters: Final = tpm_counters + rpm_counters + keys: Final = tuple(usage_key(counter) for counter in counters) + local: Final = _GLOBAL_USAGE_LOCAL_VALUES.validate_python( + await self._core_cache.async_batch_get_cache(keys=list(keys), local_only=True) + ) + if local is None: return None, None + remote: Final = await self._read_global_usage_remote(counters) + values: Final = tuple(remote.get(key, value) for key, value in zip(keys, local)) + tpm_values: Final = tuple(value for value in values[: len(tpm_counters)] if isinstance(value, int)) + rpm_values: Final = tuple(value for value in values[len(tpm_counters) :] if isinstance(value, int)) + return (sum(tpm_values) if tpm_values else None, sum(rpm_values) if rpm_values else None) - tpm_usage_list: Final[list | None] = combined_tpm_rpm_values[: len(tpm_keys)] - rpm_usage_list: Final[list | None] = combined_tpm_rpm_values[len(tpm_keys) :] + async def _read_global_usage_remote(self, counters: Sequence[GlobalUsageCounter]) -> Mapping[str, object]: + keys: Final = tuple(usage_key(counter) for counter in counters) + due_keys, previous_access_times = self._core_cache.reserve_redis_batch_reads(keys) + due: Final = frozenset(due_keys) + redis: Final = self._core_cache.redis_cache + if redis is None or not due: + return {} + requested: Final = tuple(counter for counter in counters if usage_key(counter) in due) + store: Final[RouterGlobalUsageStore] = TransportRouterGlobalUsageStore(redis.transport()) + shared: Final = await store.read_global_usage(requested) + if isinstance(shared, RedisFailure): + if shared.kind == "breaker_open": + self._core_cache.rollback_redis_batch_reads(previous_access_times) + return {} + remote: Final = MappingProxyType(dict(zip((usage_key(counter) for counter in requested), shared.values))) + for key, value in remote.items(): + if value is not None: + await self._set_global_usage_local(key, value, ttl=self._core_cache.default_in_memory_ttl) + return remote - ## TPM - tpm_usage: int | None = None - if tpm_usage_list is not None: - for t in tpm_usage_list: - if isinstance(t, int): - if tpm_usage is None: - tpm_usage = 0 - tpm_usage += t - ## RPM - rpm_usage: int | None = None - if rpm_usage_list is not None: - for t in rpm_usage_list: - if isinstance(t, int): - if rpm_usage is None: - rpm_usage = 0 - rpm_usage += t - return tpm_usage, rpm_usage + async def _set_global_usage_local(self, key: str, value: object, *, ttl: float | None) -> None: + await self._core_cache.async_set_cache(key=key, value=value, local_only=True, ttl=ttl) async def get_model_group_io_token_usage(self, model_group: str) -> tuple[int | None, int | None]: """ diff --git a/litellm/types/router.py b/litellm/types/router.py index 726a328e913..d4aedfb45c6 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -33,6 +33,8 @@ from litellm._uuid import uuid from litellm.litellm_core_utils.completion_timeout import CompletionTimeout from litellm.litellm_core_utils.core_helpers import normalize_drop_params from litellm.litellm_core_utils.provider_affinity import validate_provider_affinity_header_name +from litellm.redis.stores.router_global_usage import RPM_KEY_TEMPLATE as _RPM_KEY_TEMPLATE +from litellm.redis.stores.router_global_usage import TPM_KEY_TEMPLATE as _TPM_KEY_TEMPLATE from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.router_weights import RouterWeights @@ -1253,8 +1255,8 @@ class RoutingStrategy(enum.Enum): class RouterCacheEnum(enum.Enum): - TPM = "global_router:{id}:{model}:tpm:{current_minute}" - RPM = "global_router:{id}:{model}:rpm:{current_minute}" + TPM = _TPM_KEY_TEMPLATE + RPM = _RPM_KEY_TEMPLATE ITPM = "global_router:{id}:{model}:itpm:{current_minute}" OTPM = "global_router:{id}:{model}:otpm:{current_minute}" diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index 7e05e98d917..791515f3b41 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -175,6 +175,8 @@ ignored_function_names = [ "_cannot_run_batches", # Tested through the Azure and OpenAI mixed-group acreate_batch tests in test_router.py "_responses_continuation", # Tested through the Responses mid-stream fallback tests in tests/unit/llms/base_llm/responses/test_transformation.py "_responses_continuation_ceilings", # Tested through the Responses mid-stream fallback tests in tests/unit/llms/base_llm/responses/test_transformation.py + "_read_global_usage_remote", # Tested through the shared Redis usage header tests in test_router_helper_utils.py + "_set_global_usage_local", # Tested through the shared Redis usage header tests in test_router_helper_utils.py "_team_deployments_named", # Tested through the team lookups in the concurrent delete tests in test_router.py "_served_openai_models", # Cached as served_openai_models, tested through the pre-call check token count tests in test_router.py "priority_header_update", # Tested through the priority header tests in test_common_request_processing.py diff --git a/tests/unit/redis/stores/test_router_global_usage.py b/tests/unit/redis/stores/test_router_global_usage.py new file mode 100644 index 00000000000..cd3c1182731 --- /dev/null +++ b/tests/unit/redis/stores/test_router_global_usage.py @@ -0,0 +1,215 @@ +from collections.abc import Iterator, Mapping +from datetime import datetime, timedelta, timezone, tzinfo +from types import SimpleNamespace +from typing import Final, Protocol, cast +from unittest.mock import patch + +import fakeredis +import pytest +from redis.asyncio.client import Pipeline +from redis.exceptions import ConnectionError as RedisConnectionError + +from litellm import Router +from litellm.caching.redis_batch import request_redis_batch_scope +from litellm.caching.redis_cache import RedisCache +from litellm.redis.stores.router_global_usage import GlobalUsageCounter, TransportRouterGlobalUsageStore +from litellm.redis.transport.contract import RedisFailure +from tests._support.fake_redis import fake_redis_cache + +TPM: Final = "global_router:deployment:openai/test-雪:tpm:12-34" +RPM: Final = "global_router:deployment:openai/test-雪:rpm:12-34" + + +class Clock(datetime): + @classmethod + def now(cls, tz: tzinfo | None = None) -> datetime: + return datetime(2026, 1, 1, 12, 34, 0, tzinfo=timezone.utc).astimezone(tz) + + +@pytest.fixture(autouse=True) +def clock() -> Iterator[None]: + with ( + patch("litellm.utils.datetime", SimpleNamespace(datetime=Clock, timezone=timezone)), + patch("time.time", return_value=1767270840.0), + ): + yield + + +class LegacyUsageCache(Protocol): + async def async_set_cache(self, key: str, value: object, *, ttl: int) -> bool | None: ... + + async def async_get_cache(self, key: str) -> object: ... + + +class ExecutablePipeline(Protocol): + async def execute(self, raise_on_error: bool = True) -> object: ... + + +class UsageCallbacks(Protocol): + async def deployment_callback_on_success( + self, kwargs: Mapping[str, object], completion_response: object, start_time: object, end_time: object + ) -> object: ... + + async def async_deployment_callback_on_failure( + self, kwargs: Mapping[str, object], completion_response: object, start_time: object, end_time: object + ) -> object: ... + + +def _router(cache: RedisCache | None) -> Router: + router: Final = Router( + model_list=[ + { + "model_name": "group", + "litellm_params": {"model": "openai/test-雪", "api_key": "fake"}, + "model_info": {"id": "deployment"}, + "tpm": 1000, + "rpm": 100, + } + ] + ) + if cache is not None: + router.health().store.attach_redis_cache(cache) + return router + + +async def _record(router: Router, tokens: float = 3, *, failure: bool = False) -> None: + callbacks: Final = cast(UsageCallbacks, router) # cast-ok: typed view of the legacy callback signatures + metadata: Final = {"deployment": "openai/test-雪", "model_group": "group"} + kwargs: Final[Mapping[str, object]] = { + "call_type": "acompletion", + "standard_logging_object": {"total_tokens": tokens}, + "litellm_params": {"metadata": metadata, "model_info": {"id": "deployment"}}, + } + if failure: + await callbacks.async_deployment_callback_on_failure(kwargs, None, None, None) + else: + await callbacks.deployment_callback_on_success(kwargs, None, None, None) + + +@pytest.mark.asyncio +async def test_usage_old_new_bytes_namespace_and_post_result_reconciliation(server: fakeredis.FakeServer) -> None: + cache: Final = fake_redis_cache(server, namespace="tenant") + legacy: Final = cast(LegacyUsageCache, cache) # cast-ok: legacy writer/reader for rolling-deploy probes + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await legacy.async_set_cache(TPM, 12, ttl=60) + await legacy.async_set_cache(RPM, 2, ttl=60) + router: Final = _router(cache) + assert await router.get_model_group_usage("group") == (12, 2) + with request_redis_batch_scope(): + await _record(router, 3.5) + assert await raw.mget("tenant:" + TPM, "tenant:" + RPM) == [b"15.5", b"3"] + assert await raw.exists(TPM, RPM) == 0 + assert await raw.pttl("tenant:" + TPM) == 60000 + assert await raw.pttl("tenant:" + RPM) == 60000 + assert await legacy.async_get_cache(TPM) == 15.5 + assert await legacy.async_get_cache(RPM) == 3 + assert await router.get_model_group_usage("group") == (15, 3) + assert await _router(cache).get_model_group_usage("group") == (None, 3) + + +@pytest.mark.asyncio +async def test_failure_preserves_window_but_success_refreshes_it(server: fakeredis.FakeServer) -> None: + cache: Final = fake_redis_cache(server) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + router: Final = _router(cache) + await _record(router, failure=True) + assert await raw.get(RPM) == b"1" + assert await raw.pttl(RPM) == 60000 + with patch("time.time", return_value=1767270850.0): + await _record(router, failure=True) + assert await raw.get(RPM) == b"2" + assert await raw.pttl(RPM) == 50000 + await _record(router) + assert await raw.get(RPM) == b"3" + assert await raw.pttl(RPM) == 60000 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("seconds", [1.9, 0.9, 0, -0.9, -1.9]) +@pytest.mark.parametrize("failure", [False, True]) +async def test_usage_expiry_keeps_legacy_second_precision( + server: fakeredis.FakeServer, seconds: float, failure: bool +) -> None: + cache: Final = fake_redis_cache(server, namespace="tenant") + raw: Final = fakeredis.FakeAsyncRedis(server=server) + store: Final = TransportRouterGlobalUsageStore(cache.transport()) + await raw.set("tenant:" + RPM, b"10") + if failure: + assert ( + await store.record_global_failure("deployment", "openai/test-雪", "12-34", ttl=timedelta(seconds=seconds)) + == 11 + ) + else: + assert await store.record_global_usage( + "deployment", "openai/test-雪", "12-34", tokens=3, requests=1, ttl=timedelta(seconds=seconds) + ) == (3, 11) + with patch("time.time", return_value=1767270840.1): + assert await raw.get("tenant:" + RPM) == (b"11" if seconds >= 1 else None) + assert await raw.pttl("tenant:" + RPM) == (900 if seconds >= 1 else -2) + + +@pytest.mark.asyncio +async def test_missing_usage_is_throttled_and_late_attachment_does_not_backfill(server: fakeredis.FakeServer) -> None: + cache: Final = fake_redis_cache(server) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + router: Final = _router(None) + await _record(router) + router.health().store.attach_redis_cache(cache) + assert await raw.exists(TPM, RPM) == 0 + reader: Final = _router(cache) + assert await reader.get_model_group_usage("group") == (None, None) + await _record(router) + assert await raw.mget(TPM, RPM) == [b"3", b"1"] + assert await reader.get_model_group_usage("group") == (None, None) + with patch("time.time", return_value=1767270860.0): + assert await reader.get_model_group_usage("group") == (3, 1) + + +@pytest.mark.asyncio +async def test_malformed_batch_caches_raw_values_until_local_expiry(server: fakeredis.FakeServer) -> None: + cache: Final = fake_redis_cache(server) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await raw.set(TPM, b"bad-json") + await raw.set(RPM, b"5") + router: Final = _router(cache) + assert await router.get_model_group_usage("group") == (None, None) + await raw.set(TPM, b"10") + assert await router.get_model_group_usage("group") == (None, 5) + assert await _router(cache).get_model_group_usage("group") == (10, 5) + + +@pytest.mark.asyncio +async def test_outage_keeps_local_counts_and_shared_breaker() -> None: + server: Final = fakeredis.FakeServer() + cache: Final = fake_redis_cache(server) + router: Final = _router(cache) + server.connected = False + await _record(router) + assert await router.get_model_group_usage("group") == (3, 1) + store: Final = TransportRouterGlobalUsageStore(cache.transport()) + for _ in range(5): + await store.read_global_usage((GlobalUsageCounter("deployment", "openai/test-雪", "12-34", "tpm"),)) + result: Final = await store.record_global_failure( + "deployment", "openai/test-雪", "12-34", ttl=timedelta(seconds=60) + ) + assert isinstance(result, RedisFailure) and result.kind == "breaker_open" + server.connected = True + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", [False, True]) +async def test_unknown_applied_usage_write_is_not_replayed(server: fakeredis.FakeServer, failure: bool) -> None: + cache: Final = fake_redis_cache(server) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + router: Final = _router(cache) + original_execute: Final = cast(type[ExecutablePipeline], Pipeline).execute # cast-ok: typed redis-py execution seam + + async def apply_then_disconnect(pipe: ExecutablePipeline, raise_on_error: bool = True) -> object: + await original_execute(pipe, raise_on_error=raise_on_error) + raise RedisConnectionError("response lost after pipeline applied") + + with patch.object(Pipeline, "execute", new=apply_then_disconnect): + await _record(router, failure=failure) + assert await raw.mget(TPM, RPM) == [None if failure else b"3", b"1"] + assert await raw.pttl(RPM) == (-1 if failure else 60000) + assert await router.get_model_group_usage("group") == (None if failure else 3, 1) diff --git a/tests/unit/test_internal_context.py b/tests/unit/test_internal_context.py index 2155a75858d..6ead20c4cd5 100644 --- a/tests/unit/test_internal_context.py +++ b/tests/unit/test_internal_context.py @@ -25,7 +25,7 @@ _REPO: Final = Path(__file__).resolve().parents[2] _REDIS_PRODUCER_ROOTS: Final = ("litellm", "enterprise") # The cache implementations and facades: they emit the service events, their callers declare the family. -_CACHE_LAYER_DIRS: Final = ("litellm/caching", "litellm/_v2/cache", "litellm/redis/transport") +_CACHE_LAYER_DIRS: Final = ("litellm/caching", "litellm/_v2/cache", "litellm/redis/transport", "litellm/redis/stores") # Helpers that act on a cache handed in by the declaring caller, or forward to the response-cache facade. _CACHE_PARAMETER_HELPERS: Final = frozenset( { diff --git a/tests/unit/test_router/test_router_helper_utils.py b/tests/unit/test_router/test_router_helper_utils.py index 5f31ded626d..9084307d0c0 100644 --- a/tests/unit/test_router/test_router_helper_utils.py +++ b/tests/unit/test_router/test_router_helper_utils.py @@ -10,6 +10,7 @@ from types import MethodType from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import fakeredis import httpx import pytest import respx @@ -33,6 +34,7 @@ from litellm.types.utils import ( StandardLoggingModelInformation, StandardLoggingPayload, ) +from tests._support.fake_redis import fake_redis_cache def create_standard_logging_payload() -> StandardLoggingPayload: @@ -458,9 +460,7 @@ async def test_deployment_callback_on_success(sync_mode): ] router = Router(model_list=model_list) # Get the actual deployment ID that was generated - gpt_deployment = router.get_deployment_by_model_group_name( - model_group_name="gpt-5-mini" - ) + gpt_deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-5-mini") deployment_id = gpt_deployment["model_info"]["id"] standard_logging_payload = create_standard_logging_payload() @@ -565,9 +565,16 @@ async def _finish_35_token_websocket_turn(routed: Mapping[str, object], call_id: litellm_call_id=call_id, function_id=call_id, ) - turn.update_environment_variables(litellm_params={"litellm_metadata": routed["litellm_metadata"]}, optional_params={}) + turn.update_environment_variables( + litellm_params={"litellm_metadata": routed["litellm_metadata"]}, optional_params={} + ) await turn.dispatch_success_handlers( - [{"type": "response.completed", "response": {"usage": {"input_tokens": 20, "output_tokens": 15, "total_tokens": 35}}}], + [ + { + "type": "response.completed", + "response": {"usage": {"input_tokens": 20, "output_tokens": 15, "total_tokens": 35}}, + } + ], prefer_async_handlers=True, ) @@ -1327,9 +1334,7 @@ async def test_routing_strategy_pre_call_checks(model_list, sync_mode): router = Router(model_list=model_list) - deployment = router.get_deployment_by_model_group_name( - model_group_name="gpt-5-mini" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-5-mini") litellm_logging_obj = Logging( model="gpt-5-mini", @@ -1344,9 +1349,7 @@ async def test_routing_strategy_pre_call_checks(model_list, sync_mode): router.routing_strategy_pre_call_checks(deployment) else: ## NO EXCEPTION - await router.async_routing_strategy_pre_call_checks( - deployment, litellm_logging_obj - ) + await router.async_routing_strategy_pre_call_checks(deployment, litellm_logging_obj) ## WITH EXCEPTION - rate limit error with patch.object( @@ -1361,27 +1364,19 @@ async def test_routing_strategy_pre_call_checks(model_list, sync_mode): ), ): with pytest.raises(litellm.RateLimitError): - await router.async_routing_strategy_pre_call_checks( - deployment, litellm_logging_obj - ) + await router.async_routing_strategy_pre_call_checks(deployment, litellm_logging_obj) ## WITH EXCEPTION - generic error - with patch.object( - callback, "async_pre_call_check", AsyncMock(side_effect=Exception("Error")) - ): + with patch.object(callback, "async_pre_call_check", AsyncMock(side_effect=Exception("Error"))): with pytest.raises(Exception, match="Error"): - await router.async_routing_strategy_pre_call_checks( - deployment, litellm_logging_obj - ) + await router.async_routing_strategy_pre_call_checks(deployment, litellm_logging_obj) @pytest.mark.parametrize( "set_supported_environments, supported_environments, is_supported", [(True, ["staging"], True), (False, None, True), (True, ["development"], False)], ) -def test_create_deployment( - model_list, set_supported_environments, supported_environments, is_supported -): +def test_create_deployment(model_list, set_supported_environments, supported_environments, is_supported): """Test if the '_create_deployment' function is working correctly""" router = Router(model_list=model_list) @@ -1415,20 +1410,14 @@ def test_deployment_is_active_for_environment( ): """Test if the '_deployment_is_active_for_environment' function is working correctly""" router = Router(model_list=model_list) - deployment = router.get_deployment_by_model_group_name( - model_group_name="gpt-5-mini" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-5-mini") if set_supported_environments: os.environ["LITELLM_ENVIRONMENT"] = "staging" deployment["model_info"]["supported_environments"] = supported_environments if is_supported: - assert ( - router.deployment_is_active_for_environment(deployment=deployment) is True - ) + assert router.deployment_is_active_for_environment(deployment=deployment) is True else: - assert ( - router.deployment_is_active_for_environment(deployment=deployment) is False - ) + assert router.deployment_is_active_for_environment(deployment=deployment) is False def test_set_model_list(model_list): @@ -1755,31 +1744,14 @@ def test_deployment_tracks_usage_for_every_limit_source( assert router._deployment_tracks_usage(deployment_id, "my-unmapped-model") is tracks -def _shared_redis_stub(store: dict) -> MagicMock: - from litellm.caching.redis_cache import RedisCache - - async def increment_pipeline(increment_list, **kwargs): - for op in increment_list: - store[op["key"]] = store.get(op["key"], 0.0) + op["increment_value"] - return [store[op["key"]] for op in increment_list] - - async def batch_get(keys, **kwargs): - return {key: store.get(key) for key in keys} - - redis_stub = MagicMock(spec=RedisCache) - redis_stub.async_increment_pipeline = increment_pipeline - redis_stub.async_batch_get_cache = batch_get - return redis_stub - - @pytest.mark.asyncio @pytest.mark.usefixtures("router_minute_pinned") async def test_headers_on_fresh_worker_reflect_shared_redis_usage(): - store: dict = {} + server: Final = fakeredis.FakeServer() worker_a = _rpm_tpm_router("lit-3058-workers") worker_b = _rpm_tpm_router("lit-3058-workers") - worker_a.cache = DualCache(redis_cache=_shared_redis_stub(store), in_memory_cache=InMemoryCache()) - worker_b.cache = DualCache(redis_cache=_shared_redis_stub(store), in_memory_cache=InMemoryCache()) + worker_a.cache = DualCache(redis_cache=fake_redis_cache(server), in_memory_cache=InMemoryCache()) + worker_b.cache = DualCache(redis_cache=fake_redis_cache(server), in_memory_cache=InMemoryCache()) messages = [{"role": "user", "content": "hi"}] tokens_on_a = 0 @@ -1840,27 +1812,19 @@ async def test_get_model_group_io_token_usage_sums_across_deployments(): minute = get_utc_datetime().strftime("%H-%M") keys_and_values = [ ( - RouterCacheEnum.ITPM.value.format( - id="io-usage-dep-1", model="openai/gpt-4o-mini", current_minute=minute - ), + RouterCacheEnum.ITPM.value.format(id="io-usage-dep-1", model="openai/gpt-4o-mini", current_minute=minute), 30, ), ( - RouterCacheEnum.OTPM.value.format( - id="io-usage-dep-1", model="openai/gpt-4o-mini", current_minute=minute - ), + RouterCacheEnum.OTPM.value.format(id="io-usage-dep-1", model="openai/gpt-4o-mini", current_minute=minute), 10, ), ( - RouterCacheEnum.ITPM.value.format( - id="io-usage-dep-2", model="openai/gpt-4o", current_minute=minute - ), + RouterCacheEnum.ITPM.value.format(id="io-usage-dep-2", model="openai/gpt-4o", current_minute=minute), 70, ), ( - RouterCacheEnum.OTPM.value.format( - id="io-usage-dep-2", model="openai/gpt-4o", current_minute=minute - ), + RouterCacheEnum.OTPM.value.format(id="io-usage-dep-2", model="openai/gpt-4o", current_minute=minute), 20, ), ] @@ -2182,9 +2146,7 @@ def test_initialize_alerting(model_list): from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.types.router import AlertingConfig - router = Router( - model_list=model_list, alerting_config=AlertingConfig(webhook_url="test") - ) + router = Router(model_list=model_list, alerting_config=AlertingConfig(webhook_url="test")) router._initialize_alerting() callback_added = False @@ -2459,7 +2421,6 @@ async def test_discard_releases_per_request_override_selector(): assert selector_ref() is None - def test_discard_registers_the_selector_of_a_same_strategy_survivor_built_without_a_model_list(): retired: Final = Router(model_list=[], routing_strategy="latency-based-routing") survivor: Final = Router(routing_strategy="latency-based-routing") @@ -2537,6 +2498,7 @@ async def test_discard_releases_its_router_budget_limiter(): assert budget_limiter_ref() is None + def test_initialize_assistants_endpoint(model_list): """Test if the 'initialize_assistants_endpoint' function is working correctly""" router = Router(model_list=model_list) @@ -2612,7 +2574,6 @@ def test_deployments_by_pattern(model_list): ], ) def test_pattern_match_deployment_set_model_name(user_request_model, model_name, litellm_model, expected_model): - from re import Match from litellm.router_utils.pattern_match_deployments import PatternMatchRouter @@ -2667,7 +2628,6 @@ def test_add_optional_pre_call_checks(model_list): @pytest.mark.asyncio async def test_async_callback_filter_deployments(model_list): - from litellm.router_strategy.budget_limiter import RouterBudgetLimiting router = Router(model_list=model_list) @@ -2699,7 +2659,7 @@ def test_cached_get_model_group_info(model_list): def test_init_responses_api_endpoints(model_list): """Test if the '_init_responses_api_endpoints' function is working correctly""" - from typing import Callable + from collections.abc import Callable router = Router(model_list=model_list) @@ -3554,6 +3514,7 @@ def test_resolve_model_name_from_model_id_keeps_a_wildcard_name_only_for_the_ids router = Router(model_list=[{"model_name": "*", "litellm_params": {"model": "openai/*", "api_key": "sk-fake"}}]) assert router.resolve_model_name_from_model_id(model_id, custom_llm_provider) == expected + def test_get_valid_args(): """Test get_valid_args static method returns valid Router.__init__ arguments""" From 0cba26b52099b6ec4d58a1d8e7e77ecd830d4e69 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 20:53:39 +0000 Subject: [PATCH 09/17] refactor(redis): move Claude session router bindings behind a typed store --- .../stores/claude_code_router_binding.py | 69 + litellm/router.py | 93 +- .../router_code_coverage.py | 2 + .../stores/test_claude_code_router_binding.py | 223 ++ tests/unit/test_router/test_router.py | 2503 +++++++++-------- 5 files changed, 1633 insertions(+), 1257 deletions(-) create mode 100644 litellm/redis/stores/claude_code_router_binding.py create mode 100644 tests/unit/redis/stores/test_claude_code_router_binding.py diff --git a/litellm/redis/stores/claude_code_router_binding.py b/litellm/redis/stores/claude_code_router_binding.py new file mode 100644 index 00000000000..048c01f18ae --- /dev/null +++ b/litellm/redis/stores/claude_code_router_binding.py @@ -0,0 +1,69 @@ +import ast +import json +from dataclasses import dataclass +from datetime import timedelta +from typing import Final, Protocol + +from pydantic import JsonValue, TypeAdapter + +from litellm.redis.transport.contract import RedisFailure, RedisTransport + +_PREFIX: Final = "claude_code_session_router:v1:" +_JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) +_OBJECT: Final = TypeAdapter(object) + + +@dataclass(frozen=True, slots=True) +class SessionRouterScope: + caller: str + session_id: str + + +def binding_key(scope: SessionRouterScope) -> str: + return f"{_PREFIX}{scope.caller}:{scope.session_id}" + + +def binding_scope(key: str) -> SessionRouterScope | None: + if not key.startswith(_PREFIX): + return None + caller, separator, session_id = key[len(_PREFIX) :].rpartition(":") + return SessionRouterScope(caller, session_id) if separator and caller and session_id else None + + +def _decode(raw: bytes) -> object: + text: Final = raw.decode("utf-8") + try: + return _JSON.validate_json(text) + except ValueError: + return _OBJECT.validate_python(ast.literal_eval(text)) + + +class ClaudeCodeRouterBindingStore(Protocol): + async def read_binding(self, scope: SessionRouterScope) -> str | None | RedisFailure: ... + + async def bind_router(self, scope: SessionRouterScope, model: str, *, ttl: timedelta) -> None | RedisFailure: ... + + async def remove_binding(self, scope: SessionRouterScope) -> None | RedisFailure: ... + + +@dataclass(frozen=True, slots=True) +class TransportClaudeCodeRouterBindingStore: + transport: RedisTransport + + async def read_binding(self, scope: SessionRouterScope) -> str | None | RedisFailure: + raw: Final = await self.transport.get(binding_key(scope)) + if raw is None or isinstance(raw, RedisFailure): + return raw + try: + value: Final = _decode(raw) + except (ValueError, SyntaxError, UnicodeError) as error: + return RedisFailure(kind="decode", message=str(error)) + return value if isinstance(value, str) else None + + async def bind_router(self, scope: SessionRouterScope, model: str, *, ttl: timedelta) -> None | RedisFailure: + result: Final = await self.transport.set_seconds(binding_key(scope), json.dumps(model), ttl=ttl) + return result if isinstance(result, RedisFailure) else None + + async def remove_binding(self, scope: SessionRouterScope) -> None | RedisFailure: + result: Final = await self.transport.delete(binding_key(scope)) + return result if isinstance(result, RedisFailure) else None diff --git a/litellm/router.py b/litellm/router.py index 1809f861ad8..76275bb9618 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -83,7 +83,6 @@ from litellm.caching.caching import ( RedisCache, RedisClusterCache, ) -from litellm.caching.redis_cache import log_redis_failure from litellm.constants import ( BAGGAGE_MAPPED_TAGS_METADATA_KEY, CLIENT_OUTPUT_CEILING_METADATA_KEY, @@ -182,6 +181,13 @@ from litellm.llms.openai_like.model_info import ( MODEL_INFO_REFRESH_SECONDS, get_openai_compatible_model_info, ) +from litellm.redis.stores.claude_code_router_binding import ( + ClaudeCodeRouterBindingStore, + SessionRouterScope, + TransportClaudeCodeRouterBindingStore, + binding_key, + binding_scope, +) from litellm.redis.stores.router_global_usage import ( GlobalUsageCounter, RouterGlobalUsageStore, @@ -16238,6 +16244,10 @@ class Router(DeprecatedRouterMembers): ) def _claude_code_session_router_cache_key(self, request_kwargs: Mapping[str, object]) -> str | None: + scope: Final = self._claude_code_session_router_scope(request_kwargs) + return None if scope is None else binding_key(scope) + + def _claude_code_session_router_scope(self, request_kwargs: Mapping[str, object]) -> SessionRouterScope | None: session_id: Final = self._request_header(request_kwargs, "x-claude-code-session-id") if session_id is None or _CLAUDE_CODE_SESSION_ID_RE.fullmatch(session_id) is None: return None @@ -16248,33 +16258,52 @@ class Router(DeprecatedRouterMembers): caller_scope: Final = metadata.get("user_api_key_hash") if not isinstance(caller_scope, str) or not caller_scope: return None - return f"claude_code_session_router:v1:{caller_scope}:{session_id}" + return SessionRouterScope(caller_scope, session_id) @with_service_target(CLAUDE_CODE_SESSION_ROUTER_BINDING_TARGET) - async def _delete_claude_code_session_router_binding(self, cache_key: str) -> None: + async def _delete_claude_code_session_router_binding(self, cache_key: SessionRouterScope | str) -> None: + scope: Final = binding_scope(cache_key) if isinstance(cache_key, str) else cache_key + if scope is None: + return try: - await self._claude_code_session_router_cache.async_delete_cache(key=cache_key) - except Exception as e: # noqa: BLE001 # cache cleanup must not fail an otherwise routable request - verbose_router_logger.warning( - "Failed to delete Claude Code session router binding; the binding may remain until its TTL expires: %s", - e, - ) + session_cache: Final = self._claude_code_session_router_cache + session_cache.in_memory_cache.delete_cache(binding_key(scope)) + redis: Final = session_cache.redis_cache + if redis is not None: + store: Final[ClaudeCodeRouterBindingStore] = TransportClaudeCodeRouterBindingStore(redis.transport()) + await store.remove_binding(scope) + except Exception as e: # noqa: BLE001 # local cache cleanup must not fail an otherwise routable request + verbose_router_logger.warning("Failed to delete Claude Code session router binding: %s", e) @with_service_target(CLAUDE_CODE_SESSION_ROUTER_BINDING_TARGET) - async def _get_claude_code_session_router_binding(self, cache_key: str) -> object: - session_cache: Final = self._claude_code_session_router_cache - try: - if session_cache.redis_cache is None: - return await session_cache.async_get_cache(key=cache_key) - return await session_cache.redis_cache.async_get_cache(key=cache_key) - except Exception as e: # noqa: BLE001 # an optional binding must not make routing depend on Redis - log_redis_failure( - verbose_router_logger, - logging.WARNING, - "Failed to read Claude Code session router binding; using the requested model", - e, - ) + async def _get_claude_code_session_router_binding(self, cache_key: SessionRouterScope | str) -> str | None: + scope: Final = binding_scope(cache_key) if isinstance(cache_key, str) else cache_key + if scope is None: return None + session_cache: Final = self._claude_code_session_router_cache + redis: Final = session_cache.redis_cache + if redis is None: + local: Final = _OPAQUE_VALUE_ADAPTER.validate_python( + await session_cache.async_get_cache(key=binding_key(scope), local_only=True) + ) + return local if isinstance(local, str) else None + store: Final[ClaudeCodeRouterBindingStore] = TransportClaudeCodeRouterBindingStore(redis.transport()) + shared: Final = await store.read_binding(scope) + return None if isinstance(shared, RedisFailure) else shared + + @with_service_target(CLAUDE_CODE_SESSION_ROUTER_BINDING_TARGET) + async def _bind_claude_code_session_router(self, scope: SessionRouterScope, model: str) -> None: + session_cache: Final = self._claude_code_session_router_cache + await session_cache.async_set_cache( + key=binding_key(scope), + value=model, + ttl=_CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS, + local_only=True, + ) + redis: Final = session_cache.redis_cache + if redis is not None: + store: Final[ClaudeCodeRouterBindingStore] = TransportClaudeCodeRouterBindingStore(redis.transport()) + await store.bind_router(scope, model, ttl=timedelta(seconds=_CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS)) @with_service_target(CLAUDE_CODE_SESSION_ROUTER_BINDING_TARGET) async def _resolve_claude_code_session_router( @@ -16294,36 +16323,28 @@ class Router(DeprecatedRouterMembers): ) ): return registered_model_name - cache_key: Final = self._claude_code_session_router_cache_key(request_kwargs) - if cache_key is None or not isinstance(request_kwargs, dict): + scope: Final = self._claude_code_session_router_scope(request_kwargs) + if scope is None or not isinstance(request_kwargs, dict): return registered_model_name if request_kwargs.get("fallback_depth") not in (None, 0): return registered_model_name agent_id: Final = self._request_header(request_kwargs, "x-claude-code-agent-id") if agent_id is not None: - bound_model: Final = await self._get_claude_code_session_router_binding(cache_key) + bound_model: Final = await self._get_claude_code_session_router_binding(scope) if not isinstance(bound_model, str): return registered_model_name bound_registered_model: Final = self._core_get_model_from_alias(model=bound_model) or bound_model if self._select_pre_routing_strategy(bound_registered_model, request_kwargs) is None: - await self._delete_claude_code_session_router_binding(cache_key) + await self._delete_claude_code_session_router_binding(scope) return registered_model_name - await self._claude_code_session_router_cache.async_set_cache( - key=cache_key, - value=bound_model, - ttl=_CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS, - ) + await self._bind_claude_code_session_router(scope, bound_model) self._stamp_or_clear_metadata_key(request_kwargs, "model_group", bound_model) return bound_registered_model if self._select_pre_routing_strategy(registered_model_name, request_kwargs) is None: return registered_model_name - await self._claude_code_session_router_cache.async_set_cache( - key=cache_key, - value=model, - ttl=_CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS, - ) + await self._bind_claude_code_session_router(scope, model) return registered_model_name async def _core_async_pre_routing_hook( diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index 791515f3b41..9a5a37d53ef 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -177,6 +177,8 @@ ignored_function_names = [ "_responses_continuation_ceilings", # Tested through the Responses mid-stream fallback tests in tests/unit/llms/base_llm/responses/test_transformation.py "_read_global_usage_remote", # Tested through the shared Redis usage header tests in test_router_helper_utils.py "_set_global_usage_local", # Tested through the shared Redis usage header tests in test_router_helper_utils.py + "_claude_code_session_router_scope", # Tested through the Claude Code session router binding tests in test_router.py + "_bind_claude_code_session_router", # Tested through the Claude Code session router binding tests in test_router.py "_team_deployments_named", # Tested through the team lookups in the concurrent delete tests in test_router.py "_served_openai_models", # Cached as served_openai_models, tested through the pre-call check token count tests in test_router.py "priority_header_update", # Tested through the priority header tests in test_common_request_processing.py diff --git a/tests/unit/redis/stores/test_claude_code_router_binding.py b/tests/unit/redis/stores/test_claude_code_router_binding.py new file mode 100644 index 00000000000..35c45c74b5f --- /dev/null +++ b/tests/unit/redis/stores/test_claude_code_router_binding.py @@ -0,0 +1,223 @@ +import json +from collections.abc import Iterator, Mapping +from datetime import timedelta +from typing import Final, Protocol, cast +from unittest.mock import patch + +import fakeredis +import pytest +from pydantic import TypeAdapter +from redis.exceptions import ConnectionError as RedisConnectionError + +from litellm import Router +from litellm.caching.redis_cache import RedisCache +from litellm.redis.stores.claude_code_router_binding import SessionRouterScope, TransportClaudeCodeRouterBindingStore +from litellm.redis.transport.contract import Arg, RedisFailure +from litellm.types.router import PreRoutingHookResponse +from tests._support.fake_redis import fake_redis_cache + +KEY: Final = "claude_code_session_router:v1:key:scope:session-1234" +MODEL: Final = 'smart-雪 "router"' +SCOPE: Final = SessionRouterScope("key:scope", "session-1234") +_METADATA: Final = TypeAdapter(dict[str, object]) + + +class BindingRouter(Protocol): + def update_redis_cache(self, cache: RedisCache) -> None: ... + + async def async_pre_routing_hook( + self, model: str, request_kwargs: Mapping[str, object] + ) -> PreRoutingHookResponse | None: ... + + +class LegacyBindingCache(Protocol): + async def async_get_cache(self, key: str) -> object: ... + + async def async_set_cache(self, key: str, value: str, *, ttl: int) -> None: ... + + +@pytest.fixture(autouse=True) +def clock() -> Iterator[None]: + with patch("time.time", return_value=1800000000.0): + yield + + +def _router(cache: RedisCache | None) -> BindingRouter: + router: Final = Router( + model_list=[ + {"model_name": "plain", "litellm_params": {"model": "openai/gpt-4.1-mini", "api_key": "fake"}}, + *( + { + "model_name": model, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "tiers": dict.fromkeys(("SIMPLE", "MEDIUM", "COMPLEX", "REASONING"), "plain") + }, + }, + } + for model in (MODEL, "premium-router") + ), + ] + ) + typed: Final = cast(BindingRouter, router) # cast-ok: typed view of the actual binding caller, without replacing it + if cache is not None: + typed.update_redis_cache(cache) + return typed + + +def _kwargs( + *, agent: bool = False, caller: str = "key:scope", session: str = "session-1234", depth: int = 0 +) -> Mapping[str, object]: + return { + "metadata": {"user_api_key_hash": caller}, + "proxy_server_request": { + "headers": { + "X-Claude-Code-Session-Id": session, + **({"x-claude-code-agent-id": "agent-1234"} if agent else {}), + } + }, + "fallback_depth": depth, + } + + +async def _route(router: BindingRouter, model: str, kwargs: Mapping[str, object]) -> str: + await router.async_pre_routing_hook(model, dict(kwargs)) + bound: Final = _METADATA.validate_python(kwargs["metadata"]).get("model_group") + return bound if isinstance(bound, str) else model + + +@pytest.mark.asyncio +async def test_binding_old_new_wire_namespace_and_read_refresh(server: fakeredis.FakeServer) -> None: + cache: Final = fake_redis_cache(server, namespace="tenant") + legacy: Final = cast(LegacyBindingCache, cache) # cast-ok: rolling-deploy writer and reader + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await legacy.async_set_cache(KEY, MODEL, ttl=3600) + router: Final = _router(cache) + assert await _route(router, "plain", _kwargs(agent=True)) == MODEL + assert await raw.get("tenant:" + KEY) == json.dumps(MODEL).encode() + assert await raw.exists(KEY) == 0 + with patch("time.time", return_value=1800000100.0): + assert await raw.pttl("tenant:" + KEY) == 3500000 + assert await _route(router, "plain", _kwargs(agent=True)) == MODEL + assert await raw.pttl("tenant:" + KEY) == 3600000 + assert await _route(router, "premium-router", _kwargs()) == "premium-router" + assert await legacy.async_get_cache(KEY) == "premium-router" + assert await raw.get("tenant:" + KEY) == b'"premium-router"' + + +@pytest.mark.asyncio +async def test_shared_miss_and_remote_update_override_a_fresh_local_binding(server: fakeredis.FakeServer) -> None: + cache: Final = fake_redis_cache(server) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + main: Final = _router(cache) + agent: Final = _router(cache) + await _route(main, MODEL, _kwargs()) + assert await _route(agent, "plain", _kwargs(agent=True)) == MODEL + await _route(main, "premium-router", _kwargs()) + assert await _route(agent, "plain", _kwargs(agent=True)) == "premium-router" + await raw.delete(KEY) + assert await _route(agent, "plain", _kwargs(agent=True)) == "plain" + + +@pytest.mark.asyncio +async def test_local_bindings_survive_without_redis_but_are_not_backfilled_on_attach( + server: fakeredis.FakeServer, +) -> None: + router: Final = _router(None) + await _route(router, MODEL, _kwargs()) + assert await _route(router, "plain", _kwargs(agent=True)) == MODEL + router.update_redis_cache(fake_redis_cache(server)) + assert await _route(router, "plain", _kwargs(agent=True)) == "plain" + assert await fakeredis.FakeAsyncRedis(server=server).exists(KEY) == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("raw_value", [b"null", b"4", b"[]", b'{"router": "other"}', b"not-json", b"\xff"]) +async def test_non_string_and_malformed_bindings_are_ignored_not_rewritten( + server: fakeredis.FakeServer, raw_value: bytes +) -> None: + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await raw.set(KEY, raw_value, ex=600) + assert await _route(_router(fake_redis_cache(server)), "plain", _kwargs(agent=True)) == "plain" + assert await raw.get(KEY) == raw_value + assert await raw.pttl(KEY) == 600000 + + +@pytest.mark.asyncio +async def test_legacy_python_literal_binding_refreshes_as_json(server: fakeredis.FakeServer) -> None: + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await raw.set(KEY, b"'premium-router'", ex=600) + assert await _route(_router(fake_redis_cache(server)), "plain", _kwargs(agent=True)) == "premium-router" + assert await raw.get(KEY) == b'"premium-router"' + assert await raw.pttl(KEY) == 3600000 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("value", [b'"removed-router"', b'""']) +async def test_unknown_strategy_deletes_the_binding_instead_of_tombstoning( + server: fakeredis.FakeServer, value: bytes +) -> None: + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await raw.set(KEY, value, ex=3600) + assert await _route(_router(fake_redis_cache(server)), "plain", _kwargs(agent=True)) == "plain" + assert await raw.exists(KEY) == 0 + + +@pytest.mark.asyncio +async def test_invalid_scope_other_caller_and_fallback_cannot_reuse_the_binding(server: fakeredis.FakeServer) -> None: + router: Final = _router(fake_redis_cache(server)) + await _route(router, MODEL, _kwargs()) + assert await _route(router, "plain", _kwargs(agent=True, caller="other")) == "plain" + assert await _route(router, "plain", _kwargs(agent=True, caller="")) == "plain" + assert await _route(router, "plain", _kwargs(agent=True, session="short")) == "plain" + assert await _route(router, "plain", _kwargs(agent=True, depth=1)) == "plain" + assert await fakeredis.FakeAsyncRedis(server=server).get(KEY) == json.dumps(MODEL).encode() + + +@pytest.mark.asyncio +async def test_binding_outage_never_uses_stale_local_and_shares_the_breaker() -> None: + server: Final = fakeredis.FakeServer() + cache: Final = fake_redis_cache(server) + router: Final = _router(cache) + await _route(router, MODEL, _kwargs()) + server.connected = False + assert await _route(router, "plain", _kwargs(agent=True)) == "plain" + assert await _route(router, "premium-router", _kwargs()) == "premium-router" + store: Final = TransportClaudeCodeRouterBindingStore(cache.transport()) + for _ in range(5): + await store.read_binding(SCOPE) + result: Final = await store.bind_router(SCOPE, MODEL, ttl=timedelta(seconds=3600)) + assert isinstance(result, RedisFailure) and result.kind == "breaker_open" + + +class RedisWriter(Protocol): + async def set( + self, name: str, value: Arg, *, ex: timedelta | None = None, px: int | None = None, nx: bool = False + ) -> object: ... + + +@pytest.mark.asyncio +async def test_unknown_applied_binding_write_remains_routable_without_replay(server: fakeredis.FakeServer) -> None: + cache: Final = fake_redis_cache(server) + router: Final = _router(cache) + original_set: Final = cast(type[RedisWriter], fakeredis.FakeAsyncRedis).set # cast-ok: typed redis-py boundary + + async def apply_then_disconnect( + client: RedisWriter, + name: str, + value: Arg, + *, + ex: timedelta | None = None, + px: int | None = None, + nx: bool = False, + ) -> object: + await original_set(client, name, value, ex=ex, px=px, nx=nx) + raise RedisConnectionError("response lost after SET applied") + + with patch.object(fakeredis.FakeAsyncRedis, "set", new=apply_then_disconnect): + assert await _route(router, MODEL, _kwargs()) == MODEL + raw: Final = fakeredis.FakeAsyncRedis(server=server) + assert await raw.get(KEY) == json.dumps(MODEL).encode() + assert await raw.pttl(KEY) == 3600000 + assert await _route(_router(cache), "plain", _kwargs(agent=True)) == MODEL diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index a6436e47507..3228feff84f 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -1,16 +1,21 @@ -import ast, asyncio, importlib, time +import ast +import asyncio import contextlib import copy import functools import gc -import itertools +import importlib +import inspect import io +import itertools import json import linecache import logging import os import sys import threading +import time +import traceback import warnings from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Iterator, Mapping from datetime import datetime, timedelta, timezone @@ -25,45 +30,48 @@ import pytest import respx from fastapi import HTTPException from opentelemetry import trace -from pydantic import BaseModel, TypeAdapter +from pydantic import TypeAdapter import litellm from litellm import APIConnectionError, Router from litellm._logging import verbose_logger, verbose_router_logger from litellm.caching import RedisCache, RedisClusterCache -from litellm.caching.caching import DualCache from litellm.caching.in_memory_cache import DEFAULT_MAX_SIZE_IN_MEMORY, InMemoryCache -from litellm.caching.redis_cache import _redis_circuit_breaker_guard from litellm.exceptions import GuardrailRaisedException, MidStreamFallbackError, ModifyResponseException -from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.otel.plumbing.context import set_request_root_span from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging -from litellm.llms.anthropic.messages.native.streaming_iterator import closing_superseded_attempt +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.llms.anthropic.messages.native.agentic_streaming_iterator import ( SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES, ) +from litellm.llms.anthropic.messages.native.streaming_iterator import closing_superseded_attempt +from litellm.llms.base_llm.vector_store.transformation import ( + LiteLLMVectorStoreEmbeddingExecutor, + RouterVectorStoreEmbeddingExecutor, +) from litellm.llms.bedrock.common_utils import BedrockError from litellm.models.access_group import LiteLLM_AccessGroupTable from litellm.proxy._types import LiteLLM_TeamTable, LitellmUserRoles, Member, ProxyException, UserAPIKeyAuth +from litellm.proxy.response_polling.polling_handler import should_use_polling_for_request from litellm.router import ( MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS, MAX_HELD_PRE_OUTPUT_RESPONSES_EVENTS, FallbackAwareAnthropicMessagesStream, - _stream_commits_now, + Span, _anthropic_stream_error_is_gateway_verdict, - _stream_fallback_error_for_raised, _anthropic_stream_forwards_ping_live, - _stream_raised_error_status, _anthropic_stream_should_decline_fallback, _is_retriable_anthropic_status, _responses_stream_holds_event, - _without_line_breaks, _search_tool_support, - Span, + _stream_commits_now, + _stream_fallback_error_for_raised, + _stream_raised_error_status, + _without_line_breaks, ) -from litellm.proxy.response_polling.polling_handler import should_use_polling_for_request from litellm.router_strategy import simple_shuffle from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter from litellm.router_strategy.savings_baseline import Baseline, _models_in @@ -94,19 +102,11 @@ from litellm.types.router import ( RouterRateLimitError, RoutingContext, ) -from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.llms.base_llm.vector_store.transformation import( - LiteLLMVectorStoreEmbeddingExecutor, - RouterVectorStoreEmbeddingExecutor, -) from litellm.types.utils import CallTypes, CredentialItem, Usage from litellm.utils import _invalidate_model_cost_lowercase_map from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome from tests.fake_openai_endpoint import FAKE_OPENAI_API_BASE from tests.large_text import text -import traceback -import inspect -from typing import List, Optional if TYPE_CHECKING: from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter @@ -351,31 +351,18 @@ def test_router_model_group_encrypted_content_affinity_callback_registration(): num_retries=0, ) callbacks = router.optional_callbacks or [] - encrypted_content_callbacks = [ - cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) - ] - deployment_callback = next( - cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck) - ) + encrypted_content_callbacks = [cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)] + deployment_callback = next(cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck)) assert len(encrypted_content_callbacks) == 1 assert encrypted_content_callbacks[0].enable_global_affinity is False - assert ( - encrypted_content_callbacks[0].model_group_affinity_config - == model_group_affinity_config - ) - assert callbacks.index(encrypted_content_callbacks[0]) < callbacks.index( - deployment_callback - ) - assert litellm.callbacks.index(encrypted_content_callbacks[0]) < ( - litellm.callbacks.index(deployment_callback) - ) + assert encrypted_content_callbacks[0].model_group_affinity_config == model_group_affinity_config + assert callbacks.index(encrypted_content_callbacks[0]) < callbacks.index(deployment_callback) + assert litellm.callbacks.index(encrypted_content_callbacks[0]) < (litellm.callbacks.index(deployment_callback)) router._add_encrypted_content_affinity_check(enable_global_affinity=True) callbacks = router.optional_callbacks or [] - encrypted_content_callbacks = [ - cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) - ] + encrypted_content_callbacks = [cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)] assert len(encrypted_content_callbacks) == 1 assert encrypted_content_callbacks[0].enable_global_affinity is True assert encrypted_content_callbacks[0].router is router @@ -406,13 +393,9 @@ async def test_encrypted_content_affinity_model_group_config_is_additive(): }, target_deployment, ] - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( - "deployment-b", "rs_test" - ) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-b", "rs_test") - assert EncryptedContentAffinityCheck.has_model_group_affinity_enabled( - {model_group: ["encrypted_content_affinity"]} - ) + assert EncryptedContentAffinityCheck.has_model_group_affinity_enabled({model_group: ["encrypted_content_affinity"]}) assert not EncryptedContentAffinityCheck.has_model_group_affinity_enabled(None) per_group_check = EncryptedContentAffinityCheck( @@ -453,10 +436,7 @@ async def test_encrypted_content_affinity_model_group_config_is_additive(): ) assert unfiltered == healthy_deployments - assert ( - "encrypted_content_affinity_enabled" - not in disabled_request_kwargs["litellm_metadata"] - ) + assert "encrypted_content_affinity_enabled" not in disabled_request_kwargs["litellm_metadata"] global_check = EncryptedContentAffinityCheck( enable_global_affinity=True, @@ -476,9 +456,7 @@ async def test_encrypted_content_affinity_model_group_config_is_additive(): ) assert globally_filtered == [target_deployment] - assert global_request_kwargs["litellm_metadata"][ - "encrypted_content_affinity_enabled" - ] + assert global_request_kwargs["litellm_metadata"]["encrypted_content_affinity_enabled"] @pytest.mark.asyncio @@ -525,18 +503,10 @@ async def test_encrypted_content_affinity_takes_priority_over_user_key_affinity( num_retries=0, ) callbacks = router.optional_callbacks or [] - deployment_callback = next( - cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck) - ) - encrypted_content_callback = next( - cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) - ) - assert callbacks.index(encrypted_content_callback) < callbacks.index( - deployment_callback - ) - assert litellm.callbacks.index(encrypted_content_callback) < ( - litellm.callbacks.index(deployment_callback) - ) + deployment_callback = next(cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck)) + encrypted_content_callback = next(cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)) + assert callbacks.index(encrypted_content_callback) < callbacks.index(deployment_callback) + assert litellm.callbacks.index(encrypted_content_callback) < (litellm.callbacks.index(deployment_callback)) cache_key = DeploymentAffinityCheck.get_affinity_cache_key( model_group=model_group, @@ -547,9 +517,7 @@ async def test_encrypted_content_affinity_takes_priority_over_user_key_affinity( value={"model_id": "deployment-a"}, ttl=60, ) - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( - "deployment-b", "rs_test" - ) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-b", "rs_test") request_kwargs = { "input": [{"type": "reasoning", "id": encoded_id}], "litellm_metadata": {"user_api_key_hash": user_api_key_hash}, @@ -759,7 +727,9 @@ async def test_async_router_acreate_file_passthrough_keeps_the_file_and_forwards from io import BytesIO from unittest.mock import MagicMock, patch - jsonl_content = b'{"custom_id": "r1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "vertex-batch"}}\n' + jsonl_content = ( + b'{"custom_id": "r1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "vertex-batch"}}\n' + ) router = litellm.Router( model_list=[ { @@ -770,9 +740,7 @@ async def test_async_router_acreate_file_passthrough_keeps_the_file_and_forwards ) with patch("litellm.acreate_file", return_value=MagicMock()) as mock_acreate_file: - await router.acreate_file( - model="vertex-batch", purpose="batch", file=BytesIO(jsonl_content), passthrough=True - ) + await router.acreate_file(model="vertex-batch", purpose="batch", file=BytesIO(jsonl_content), passthrough=True) forwarded = mock_acreate_file.call_args.kwargs assert forwarded["passthrough"] is True forwarded["file"].seek(0) @@ -1044,7 +1012,8 @@ async def test_async_router_afile_content_uses_deployment_custom_llm_provider(): This prevents "None is not a valid LlmProviders" errors when calling file content operations. """ - from unittest.mock import AsyncMock, MagicMock, patch + from unittest.mock import MagicMock, patch + from litellm.types.llms.openai import HttpxBinaryResponseContent router = litellm.Router( @@ -1248,9 +1217,7 @@ async def test_arouter_aretrieve_batch(): ], ) - with patch.object( - litellm, "aretrieve_batch", return_value=AsyncMock() - ) as mock_aretrieve_batch: + with patch.object(litellm, "aretrieve_batch", return_value=AsyncMock()) as mock_aretrieve_batch: try: response = await router.aretrieve_batch( model="gpt-3.5-turbo", @@ -1510,6 +1477,7 @@ def test_sync_deployment_callback_on_success_skips_batch_retrieves( == expected_successes ) + _ROUTING_STRATEGY_CACHE_MARKERS = ("_map", "_request_count", ":tpm:", ":rpm:") @@ -1521,8 +1489,7 @@ async def _moved_routing_counters(router, timeout: float = 2.0) -> list[str]: moved = sorted( f"{key}={cache_dict[key]}" for key in cache_dict - if any(marker in key for marker in _ROUTING_STRATEGY_CACHE_MARKERS) - and cache_dict[key] + if any(marker in key for marker in _ROUTING_STRATEGY_CACHE_MARKERS) and cache_dict[key] ) if moved: return moved @@ -1603,9 +1570,7 @@ async def test_arouter_aretrieve_file_content(): Test that router.acreate_file with JSONL file returns the correct response """ - with patch.object( - litellm, "afile_content", return_value=AsyncMock() - ) as mock_afile_content: + with patch.object(litellm, "afile_content", return_value=AsyncMock()) as mock_afile_content: router = litellm.Router( model_list=[ { @@ -1666,7 +1631,7 @@ async def test_arouter_filter_team_based_models(): assert result is not None # FAILS - with pytest.raises(Exception, match='No deployments available for selected model, Try again in') as e: + with pytest.raises(Exception, match="No deployments available for selected model, Try again in") as e: result = await router.acompletion( model="gpt-3.5-turbo", messages=[{"role": "user", "content": "Hello, world!"}], @@ -1750,9 +1715,7 @@ def test_arouter_should_include_deployment(): model=deployment_with_team_and_public_name, team_id="test-team", ) - assert ( - result is True - ), "Should return True when team_id and team_public_model_name match" + assert result is True, "Should return True when team_id and team_public_model_name match" # Test Case 2: Team-specific deployment - team_id matches but model_name doesn't match team_public_model_name result = router.should_include_deployment( @@ -1760,9 +1723,9 @@ def test_arouter_should_include_deployment(): model=deployment_with_team_and_public_name, team_id="test-team", ) - assert ( - result is False - ), "Should return False when team_id matches but model_name doesn't match team_public_model_name" + assert result is False, ( + "Should return False when team_id matches but model_name doesn't match team_public_model_name" + ) # Test Case 3: Team-specific deployment - team_id doesn't match result = router.should_include_deployment( @@ -1778,30 +1741,18 @@ def test_arouter_should_include_deployment(): model=deployment_with_team_no_public_name, team_id="test-team", ) - assert ( - result is True - ), "Should return True when team deployment has no team_public_model_name to match" + assert result is True, "Should return True when team deployment has no team_public_model_name to match" # Test Case 5: Non-team deployment - model_name matches and no team_id - result = router.should_include_deployment( - model_name="gpt-4", model=deployment_without_team, team_id=None - ) - assert ( - result is True - ), "Should return True when model_name matches and deployment has no team_id" + result = router.should_include_deployment(model_name="gpt-4", model=deployment_without_team, team_id=None) + assert result is True, "Should return True when model_name matches and deployment has no team_id" # Test Case 6: Non-team deployment - model_name matches but team_id provided (should still work) - result = router.should_include_deployment( - model_name="gpt-4", model=deployment_without_team, team_id="any-team" - ) - assert ( - result is True - ), "Should return True when model_name matches non-team deployment, regardless of team_id param" + result = router.should_include_deployment(model_name="gpt-4", model=deployment_without_team, team_id="any-team") + assert result is True, "Should return True when model_name matches non-team deployment, regardless of team_id param" # Test Case 7: Non-team deployment - model_name doesn't match - result = router.should_include_deployment( - model_name="different-model", model=deployment_without_team, team_id=None - ) + result = router.should_include_deployment(model_name="different-model", model=deployment_without_team, team_id=None) assert result is False, "Should return False when model_name doesn't match" # Test Case 8: Team deployment accessed without matching team_id @@ -1810,9 +1761,7 @@ def test_arouter_should_include_deployment(): model=deployment_with_team_and_public_name, team_id=None, ) - assert ( - result is True - ), "Should return True when matching model with exact model_name" + assert result is True, "Should return True when matching model with exact model_name" def test_arouter_responses_api_bridge(): @@ -1863,9 +1812,7 @@ def test_arouter_responses_api_bridge(): "status": "completed", "output": [], } - mock_response.text = ( - '{"id": "resp_test", "object": "response", "status": "completed", "output": []}' - ) + mock_response.text = '{"id": "resp_test", "object": "response", "status": "completed", "output": []}' with patch.object(client, "post", return_value=mock_response) as mock_post: try: @@ -1939,7 +1886,7 @@ def test_add_invalid_provider_to_router(): ], ) - with pytest.raises(Exception, match='Unsupported provider - vertex_ai_eu') as e: + with pytest.raises(Exception, match="Unsupported provider - vertex_ai_eu") as e: router.add_deployment( Deployment( model_name="vertex_ai/*", @@ -1964,7 +1911,9 @@ def registered_custom_provider(monkeypatch: pytest.MonkeyPatch) -> str: model="gpt-5.6", messages=[{"role": "user", "content": "hi"}], mock_response="served by onprem handler" ) - monkeypatch.setattr(litellm, "custom_provider_map", [{"provider": "test-onprem-llm", "custom_handler": OnPremLLM()}]) + monkeypatch.setattr( + litellm, "custom_provider_map", [{"provider": "test-onprem-llm", "custom_handler": OnPremLLM()}] + ) monkeypatch.setattr(litellm, "provider_list", list(litellm.provider_list)) monkeypatch.setattr(litellm, "_custom_providers", list(litellm._custom_providers)) return "test-onprem-llm" @@ -2041,15 +1990,9 @@ async def test_router_ageneric_api_call_with_fallbacks_helper(): }, } - with patch.object( - router, "_update_kwargs_with_deployment" - ) as mock_update_kwargs: - with patch.object( - router, "async_routing_strategy_pre_call_checks" - ) as mock_pre_call_checks: - with patch.object( - router, "_get_client", return_value=None - ) as mock_get_client: + with patch.object(router, "_update_kwargs_with_deployment") as mock_update_kwargs: + with patch.object(router, "async_routing_strategy_pre_call_checks") as mock_pre_call_checks: + with patch.object(router, "_get_client", return_value=None) as mock_get_client: result = await router._ageneric_api_call_with_fallbacks_helper( model="gpt-3.5-turbo", original_generic_function=mock_generic_function, @@ -2084,7 +2027,7 @@ async def test_router_ageneric_api_call_with_fallbacks_helper(): with patch.object(router, "async_get_available_deployment") as mock_get_deployment: mock_get_deployment.side_effect = Exception("No deployment available") - with pytest.raises(Exception, match='No deployment available') as exc_info: + with pytest.raises(Exception, match="No deployment available") as exc_info: await router._ageneric_api_call_with_fallbacks_helper( model="gpt-3.5-turbo", original_generic_function=mock_generic_function, @@ -2095,7 +2038,6 @@ async def test_router_ageneric_api_call_with_fallbacks_helper(): assert "No deployment available" in str(exc_info.value) # Test 4: Test with semaphore (rate limiting) - import asyncio async def mock_semaphore_function(**kwargs): return {"result": "semaphore_success", "model": kwargs.get("model")} @@ -2114,15 +2056,9 @@ async def test_router_ageneric_api_call_with_fallbacks_helper(): max_parallel_requests=1, model_id="deployment-1", model_group="gpt-3.5-turbo" ) - with patch.object( - router, "_update_kwargs_with_deployment" - ) as mock_update_kwargs: - with patch.object( - router, "_get_client", return_value=mock_semaphore - ) as mock_get_client: - with patch.object( - router, "async_routing_strategy_pre_call_checks" - ) as mock_pre_call_checks: + with patch.object(router, "_update_kwargs_with_deployment") as mock_update_kwargs: + with patch.object(router, "_get_client", return_value=mock_semaphore) as mock_get_client: + with patch.object(router, "async_routing_strategy_pre_call_checks") as mock_pre_call_checks: result = await router._ageneric_api_call_with_fallbacks_helper( model="gpt-3.5-turbo", original_generic_function=mock_semaphore_function, @@ -2151,16 +2087,10 @@ async def test_router_ageneric_api_call_with_fallbacks_helper(): }, } - with patch.object( - router, "_update_kwargs_with_deployment" - ) as mock_update_kwargs: - with patch.object( - router, "_get_client", return_value=None - ) as mock_get_client: - with patch.object( - router, "async_routing_strategy_pre_call_checks" - ) as mock_pre_call_checks: - with pytest.raises(Exception, match='Mock failure') as exc_info: + with patch.object(router, "_update_kwargs_with_deployment") as mock_update_kwargs: + with patch.object(router, "_get_client", return_value=None) as mock_get_client: + with patch.object(router, "async_routing_strategy_pre_call_checks") as mock_pre_call_checks: + with pytest.raises(Exception, match="Mock failure") as exc_info: await router._ageneric_api_call_with_fallbacks_helper( model="gpt-3.5-turbo", original_generic_function=mock_failing_function, @@ -2234,9 +2164,9 @@ async def test_ageneric_api_call_deployment_model_overrides_alias(): original_generic_function=capture_model, ) - assert ( - captured["model"] == "vertex_ai/gemini-2.5-flash" - ), f"Expected deployment model 'vertex_ai/gemini-2.5-flash', got '{captured['model']}'" + assert captured["model"] == "vertex_ai/gemini-2.5-flash", ( + f"Expected deployment model 'vertex_ai/gemini-2.5-flash', got '{captured['model']}'" + ) @pytest.mark.asyncio @@ -2342,14 +2272,10 @@ def test_router_get_model_access_groups_team_only_models(): ] ) - access_groups = router.get_model_access_groups( - model_name="gpt-3.5-turbo", team_id=None - ) + access_groups = router.get_model_access_groups(model_name="gpt-3.5-turbo", team_id=None) assert len(access_groups) == 0 - access_groups = router.get_model_access_groups( - model_name="gpt-3.5-turbo", team_id="team_1" - ) + access_groups = router.get_model_access_groups(model_name="gpt-3.5-turbo", team_id="team_1") assert list(access_groups.keys()) == ["default-models"] @@ -2444,9 +2370,7 @@ def test_model_group_info_cost_from_db_model_info(): ] ) - with patch.object( - router, "_core_get_deployment_model_info", side_effect=Exception("not found") - ): + with patch.object(router, "_core_get_deployment_model_info", side_effect=Exception("not found")): result = router._cached_get_model_group_info("my-custom-model") assert result is not None assert result.input_cost_per_token == 0.0001 @@ -2474,9 +2398,7 @@ def test_model_group_info_cost_none_when_db_model_info_has_no_cost(): ] ) - with patch.object( - router, "_core_get_deployment_model_info", side_effect=Exception("not found") - ): + with patch.object(router, "_core_get_deployment_model_info", side_effect=Exception("not found")): result = router._cached_get_model_group_info("my-custom-model-no-cost") assert result is not None assert result.input_cost_per_token is None @@ -2962,9 +2884,7 @@ def test_model_group_info_with_stringified_cost_values(): } return None - with patch.object( - router, "_core_get_deployment_model_info", side_effect=_model_info_with_str_costs - ): + with patch.object(router, "_core_get_deployment_model_info", side_effect=_model_info_with_str_costs): result = router._set_model_group_info( model_group="my-custom-model", user_facing_model_group_name="my-custom-model", @@ -3010,9 +2930,7 @@ def test_model_group_info_db_fallback_with_stringified_cost_values(): ] ) - with patch.object( - router, "_core_get_deployment_model_info", side_effect=Exception("not found") - ): + with patch.object(router, "_core_get_deployment_model_info", side_effect=Exception("not found")): result = router._set_model_group_info( model_group="my-custom-model", user_facing_model_group_name="my-custom-model", @@ -3272,6 +3190,7 @@ async def test_acompletion_streaming_iterator(): # Collect streamed chunks — the first chunk succeeds, then the error re-raises collected_chunks = [] + async def _drain(): async for chunk in result: collected_chunks.append(chunk) @@ -3335,7 +3254,7 @@ async def test_acompletion_streaming_iterator_continues_after_content_when_eligi """Flag on + plain-text break: the router re-enters the fallback chain with an assistant-prefill continuation and the mid-stream marker, then streams the fallback's output instead of re-raising.""" - from unittest.mock import AsyncMock, MagicMock, patch + from unittest.mock import AsyncMock, patch from litellm.exceptions import MidStreamFallbackError from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( @@ -3404,8 +3323,12 @@ async def test_acompletion_streaming_iterator_continuation_runs_post_call_rules_ enable_mid_stream_fallback_continuation=True, ) error = MidStreamFallbackError( - message="Connection lost", model="gpt-4", llm_provider="openai", generated_content="123-45-", - is_pre_first_chunk=False, emitted_disqualifying_content=False, + message="Connection lost", + model="gpt-4", + llm_provider="openai", + generated_content="123-45-", + is_pre_first_chunk=False, + emitted_disqualifying_content=False, ) fallback_chunk = litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": continuation_text}}]) @@ -3450,8 +3373,12 @@ async def test_a_second_cut_runs_post_call_rules_only_on_the_text_the_caller_rec def cut(generated_content): return MidStreamFallbackError( - message="Connection lost", model="gpt-4", llm_provider="openai", generated_content=generated_content, - is_pre_first_chunk=False, emitted_disqualifying_content=False, + message="Connection lost", + model="gpt-4", + llm_provider="openai", + generated_content=generated_content, + is_pre_first_chunk=False, + emitted_disqualifying_content=False, ) messages = [{"role": "user", "content": "Hi"}] @@ -3461,7 +3388,11 @@ async def test_a_second_cut_runs_post_call_rules_only_on_the_text_the_caller_rec first_chunk=litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": " there"}}]), ), messages=[*messages, {"role": "assistant", "content": "Hello", "prefix": True}], - initial_kwargs={"model": "gpt-4", "stream": True, MID_STREAM_CONTINUATION_KWARG: MID_STREAM_CONTINUATION_MARKER}, + initial_kwargs={ + "model": "gpt-4", + "stream": True, + MID_STREAM_CONTINUATION_KWARG: MID_STREAM_CONTINUATION_MARKER, + }, ) backup_leg = _FakeFallbackStream( litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": " friend"}}]) @@ -3500,8 +3431,12 @@ async def test_acompletion_streaming_iterator_continuation_never_repeats_the_par enable_mid_stream_fallback_continuation=True, ) error = MidStreamFallbackError( - message="Connection lost", model="gpt-4", llm_provider="openai", generated_content="Hello", - is_pre_first_chunk=False, emitted_disqualifying_content=False, + message="Connection lost", + model="gpt-4", + llm_provider="openai", + generated_content="Hello", + is_pre_first_chunk=False, + emitted_disqualifying_content=False, ) async def fallback_stream(): @@ -3537,8 +3472,12 @@ async def test_acompletion_streaming_iterator_continuation_stays_in_the_primary_ enable_mid_stream_fallback_continuation=True, ) error = MidStreamFallbackError( - message="Connection lost", model="gpt-4", llm_provider="openai", generated_content="Hello", - is_pre_first_chunk=False, emitted_disqualifying_content=False, + message="Connection lost", + model="gpt-4", + llm_provider="openai", + generated_content="Hello", + is_pre_first_chunk=False, + emitted_disqualifying_content=False, ) primary_chunk = litellm.ModelResponseStream( id="chatcmpl-primary", created=111, choices=[{"index": 0, "delta": {"role": "assistant", "content": "Hello"}}] @@ -3623,8 +3562,12 @@ async def test_acompletion_streaming_iterator_continuation_usage_sums_both_legs( messages = [{"role": "user", "content": "Explain transformers in one paragraph."}] partial = "Transformers use " error = MidStreamFallbackError( - message="Connection lost", model="gpt-4", llm_provider="openai", generated_content=partial, - is_pre_first_chunk=False, emitted_disqualifying_content=False, + message="Connection lost", + model="gpt-4", + llm_provider="openai", + generated_content=partial, + is_pre_first_chunk=False, + emitted_disqualifying_content=False, ) content_chunk = litellm.ModelResponseStream( model="gpt-4", choices=[{"index": 0, "delta": {"role": "assistant", "content": partial}}] @@ -3646,20 +3589,35 @@ class _CutSSE(httpx.AsyncByteStream): self._content: Final = content async def __aiter__(self) -> AsyncIterator[bytes]: - yield f'data: {json.dumps({"id": "cut", "object": "chat.completion.chunk", "created": 1700000000, "model": "m", "choices": [{"index": 0, "delta": {"role": "assistant", "content": self._content}, "finish_reason": None}]})}\n\n'.encode() + yield f"data: {json.dumps({'id': 'cut', 'object': 'chat.completion.chunk', 'created': 1700000000, 'model': 'm', 'choices': [{'index': 0, 'delta': {'role': 'assistant', 'content': self._content}, 'finish_reason': None}]})}\n\n".encode() raise httpx.RemoteProtocolError("peer closed connection without sending complete message body") -_BACKUP_SSE: Final = "".join( - f"data: {json.dumps(chunk)}\n\n" - for chunk in ( - {"id": "b", "object": "chat.completion.chunk", "created": 1700000000, "model": "mistral-large-latest", - "choices": [{"index": 0, "delta": {"role": "assistant", "content": " self-attention."}, "finish_reason": None}]}, - {"id": "b", "object": "chat.completion.chunk", "created": 1700000000, "model": "mistral-large-latest", - "choices": [{"index": 0, "delta": {"content": ""}, "finish_reason": "stop"}], - "usage": {"prompt_tokens": 91, "completion_tokens": 5, "total_tokens": 96}}, +_BACKUP_SSE: Final = ( + "".join( + f"data: {json.dumps(chunk)}\n\n" + for chunk in ( + { + "id": "b", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "mistral-large-latest", + "choices": [ + {"index": 0, "delta": {"role": "assistant", "content": " self-attention."}, "finish_reason": None} + ], + }, + { + "id": "b", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "mistral-large-latest", + "choices": [{"index": 0, "delta": {"content": ""}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 91, "completion_tokens": 5, "total_tokens": 96}, + }, + ) ) -) + "data: [DONE]\n\n" + + "data: [DONE]\n\n" +) class _SpendRecorder(CustomLogger): @@ -3701,7 +3659,9 @@ async def test_continued_stream_success_log_bills_every_cut_leg( recorder: Final = _SpendRecorder() monkeypatch.setattr(litellm, "callbacks", [recorder]) respx_mock.post("https://api.openai.com/v1/chat/completions").mock( - return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, stream=_CutSSE("Transformers use")) + return_value=httpx.Response( + 200, headers={"content-type": "text/event-stream"}, stream=_CutSSE("Transformers use") + ) ) respx_mock.post(host="api.deepseek.com").mock( return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, stream=_CutSSE(" learned")) @@ -3784,7 +3744,9 @@ async def test_continued_stream_sharing_the_proxy_logging_object_bills_the_cut_t recorder: Final = _SlowFailureRecorder() monkeypatch.setattr(litellm, "callbacks", [recorder]) respx_mock.post("https://api.openai.com/v1/chat/completions").mock( - return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, stream=_CutSSE("Transformers use")) + return_value=httpx.Response( + 200, headers={"content-type": "text/event-stream"}, stream=_CutSSE("Transformers use") + ) ) respx_mock.post(host="api.mistral.ai").mock( return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, text=_BACKUP_SSE) @@ -3899,7 +3861,10 @@ async def test_responses_fallback_sharing_the_proxy_logging_object_logs_the_fall 200, headers={"content-type": "text/event-stream"}, text=_responses_sse( - {"type": "response.created", "response": _responses_envelope("resp_backup", "gpt-4o-mini", "in_progress")}, + { + "type": "response.created", + "response": _responses_envelope("resp_backup", "gpt-4o-mini", "in_progress"), + }, { "type": "response.completed", "response": _responses_envelope( @@ -3973,8 +3938,13 @@ def test_carry_partial_spend_to_fallback_adds_the_partial_the_cut_stream_inherit from litellm.types.utils import Usage cut_logging: Final = LiteLLMLogging( - model="deepseek/deepseek-chat", messages=[], stream=True, call_type="acompletion", - start_time=datetime.now(), litellm_call_id="cut", function_id="cut", + model="deepseek/deepseek-chat", + messages=[], + stream=True, + call_type="acompletion", + start_time=datetime.now(), + litellm_call_id="cut", + function_id="cut", ) cut_logging.model_call_details["response_cost"] = 0.5 cut_logging.update_environment_variables( @@ -3986,8 +3956,13 @@ def test_carry_partial_spend_to_fallback_adds_the_partial_the_cut_stream_inherit response_id="chatcmpl-primary", ) fallback_logging: Final = LiteLLMLogging( - model="mistral-large-latest", messages=[], stream=True, call_type="acompletion", - start_time=datetime.now(), litellm_call_id="fallback", function_id="fallback", + model="mistral-large-latest", + messages=[], + stream=True, + call_type="acompletion", + start_time=datetime.now(), + litellm_call_id="fallback", + function_id="fallback", ) Router._carry_partial_spend_to_fallback( @@ -4016,7 +3991,10 @@ async def test_acompletion_streaming_iterator_plain_fallback_keeps_the_fallback_ from litellm.exceptions import MidStreamFallbackError error = MidStreamFallbackError( - message="Connection lost", model="gpt-4", llm_provider="anthropic", generated_content="", + message="Connection lost", + model="gpt-4", + llm_provider="anthropic", + generated_content="", is_pre_first_chunk=True, ) role_chunk = litellm.ModelResponseStream( @@ -4050,8 +4028,12 @@ async def test_acompletion_streaming_iterator_continuation_does_not_double_trail from litellm.exceptions import MidStreamFallbackError error = MidStreamFallbackError( - message="Connection lost", model="gpt-4", llm_provider="openai", generated_content=partial, - is_pre_first_chunk=False, emitted_disqualifying_content=False, + message="Connection lost", + model="gpt-4", + llm_provider="openai", + generated_content=partial, + is_pre_first_chunk=False, + emitted_disqualifying_content=False, ) fallback_items = [ litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": text}}]) for text in fallback_texts @@ -4088,8 +4070,12 @@ async def test_update_settings_turns_on_mid_stream_continuation(monkeypatch, ena router.update_settings(enable_mid_stream_fallback_continuation=enabled) error = MidStreamFallbackError( - message="Connection lost", model="gpt-4", llm_provider="openai", generated_content="Hello", - is_pre_first_chunk=False, emitted_disqualifying_content=False, + message="Connection lost", + model="gpt-4", + llm_provider="openai", + generated_content="Hello", + is_pre_first_chunk=False, + emitted_disqualifying_content=False, ) fallback_chunk = litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": " world"}}]) with patch.object( @@ -4137,8 +4123,12 @@ async def test_acompletion_streaming_iterator_declines_ineligible_after_content( enable_mid_stream_fallback_continuation=True, ) error = MidStreamFallbackError( - message="boom", model="gpt-4", llm_provider="openai", generated_content="Hello", - is_pre_first_chunk=False, **error_kwargs, + message="boom", + model="gpt-4", + llm_provider="openai", + generated_content="Hello", + is_pre_first_chunk=False, + **error_kwargs, ) source = _make_midstream_source(error) initial_kwargs = {"model": "gpt-4", "stream": True, **request_kwargs} @@ -4172,8 +4162,13 @@ async def test_acompletion_streaming_iterator_raises_content_filter_instead_of_c ) content_filter = litellm.ContentPolicyViolationError(message="filtered", model="gpt-4", llm_provider="azure") error = MidStreamFallbackError( - message="filtered", model="gpt-4", llm_provider="azure", generated_content="Step one is to obtain the", - is_pre_first_chunk=False, emitted_disqualifying_content=False, original_exception=content_filter, + message="filtered", + model="gpt-4", + llm_provider="azure", + generated_content="Step one is to obtain the", + is_pre_first_chunk=False, + emitted_disqualifying_content=False, + original_exception=content_filter, ) with patch.object(router, "async_function_with_fallbacks_common_utils", new=AsyncMock()) as mock_fallback: @@ -4232,8 +4227,12 @@ async def _continues_after_content( enable_mid_stream_fallback_continuation=True, ) error: Final = MidStreamFallbackError( - message="boom", model="gpt-4", llm_provider="openai", generated_content="Hello", - is_pre_first_chunk=False, emitted_disqualifying_content=False, + message="boom", + model="gpt-4", + llm_provider="openai", + generated_content="Hello", + is_pre_first_chunk=False, + emitted_disqualifying_content=False, ) fallback_chunk: Final = litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": " world"}}]) cut_stream: Final = _make_midstream_source(error) @@ -4448,14 +4447,19 @@ async def test_acompletion_streaming_iterator_flag_off_declines_after_content(fl **flag_kwargs, ) error = MidStreamFallbackError( - message="boom", model="gpt-4", llm_provider="openai", generated_content="Hello", - is_pre_first_chunk=False, emitted_disqualifying_content=False, + message="boom", + model="gpt-4", + llm_provider="openai", + generated_content="Hello", + is_pre_first_chunk=False, + emitted_disqualifying_content=False, ) source = _make_midstream_source(error) with patch.object(router, "async_function_with_fallbacks_common_utils", new=AsyncMock()) as mock_fallback: result = await router._acompletion_streaming_iterator( - model_response=source, messages=[{"role": "user", "content": "Hi"}], + model_response=source, + messages=[{"role": "user", "content": "Hi"}], initial_kwargs={"model": "gpt-4", "stream": True}, ) with pytest.raises(MidStreamFallbackError): @@ -4545,8 +4549,12 @@ async def test_acompletion_streaming_iterator_declines_when_request_ends_in_plai enable_mid_stream_fallback_continuation=True, ) error = MidStreamFallbackError( - message="boom", model="gpt-4", llm_provider="openai", generated_content="Hello", - is_pre_first_chunk=False, emitted_disqualifying_content=False, + message="boom", + model="gpt-4", + llm_provider="openai", + generated_content="Hello", + is_pre_first_chunk=False, + emitted_disqualifying_content=False, ) with patch.object(router, "async_function_with_fallbacks_common_utils", new=AsyncMock()) as mock_fallback: @@ -4577,7 +4585,10 @@ def test_build_completion_continuation_input_keeps_structured_assistant_content( assert [m["role"] for m in built] == ["user", "assistant"] assert built[-1] == {"role": "assistant", "content": [{"type": "text", "text": "Sure, here it is"}], "prefix": True} image_block = {"type": "image_url", "image_url": {"url": "https://example.test/a.png"}} - assert _extend_trailing_text_block([image_block], "here it is") == [image_block, {"type": "text", "text": "here it is"}] + assert _extend_trailing_text_block([image_block], "here it is") == [ + image_block, + {"type": "text", "text": "here it is"}, + ] assert _extend_trailing_text_block([], "here it is") == [{"type": "text", "text": "here it is"}] assert ( litellm.Router._build_completion_continuation_input([{"role": "assistant", "content": 42, "prefix": True}], "x") @@ -4595,11 +4606,18 @@ def test_mid_stream_continuation_eligible_allows_text_response_format(): enable_mid_stream_fallback_continuation=True, ) e = MidStreamFallbackError( - message="boom", model="gpt-4", llm_provider="openai", generated_content="Hello", - is_pre_first_chunk=False, emitted_disqualifying_content=False, + message="boom", + model="gpt-4", + llm_provider="openai", + generated_content="Hello", + is_pre_first_chunk=False, + emitted_disqualifying_content=False, ) assert router._mid_stream_continuation_eligible(e=e, request_kwargs={"response_format": {"type": "text"}}) is True - assert router._mid_stream_continuation_eligible(e=e, request_kwargs={"response_format": {"type": "json_object"}}) is False + assert ( + router._mid_stream_continuation_eligible(e=e, request_kwargs={"response_format": {"type": "json_object"}}) + is False + ) def test_raise_original_mid_stream_error_surfaces_the_provider_exception(): @@ -4607,8 +4625,11 @@ def test_raise_original_mid_stream_error_surfaces_the_provider_exception(): provider_error = RateLimitError(message="rate limited", llm_provider="openai", model="gpt-4") wrapped = MidStreamFallbackError( - message="rate limited", model="gpt-4", llm_provider="openai", - original_exception=provider_error, generated_content="Hello", + message="rate limited", + model="gpt-4", + llm_provider="openai", + original_exception=provider_error, + generated_content="Hello", ) with pytest.raises(RateLimitError) as raised: litellm.Router._raise_original_mid_stream_error(wrapped) @@ -4955,9 +4976,7 @@ def test_adopt_fallback_response_headers_replaces_rather_than_merges(): "additional_headers": {"llm_provider-x-request-id": "req-FALLBACK"}, } - wrapper.adopt_fallback_response_headers( - fallback, Router._prepare_fallback_hidden_params(fallback) - ) + wrapper.adopt_fallback_response_headers(fallback, Router._prepare_fallback_hidden_params(fallback)) assert wrapper._response_headers == {"x-request-id": "req-FALLBACK"} assert wrapper._hidden_params["model_id"] == "fallback-deployment" @@ -5029,9 +5048,7 @@ def test_adopt_fallback_response_headers_drops_headers_the_fallback_cannot_repla fallback._response_headers = None fallback._hidden_params = {"model_id": "fallback-deployment"} - wrapper.adopt_fallback_response_headers( - fallback, Router._prepare_fallback_hidden_params(fallback) - ) + wrapper.adopt_fallback_response_headers(fallback, Router._prepare_fallback_hidden_params(fallback)) assert wrapper._response_headers is None assert wrapper._hidden_params["model_id"] == "fallback-deployment" @@ -5058,9 +5075,7 @@ def test_adopt_fallback_response_headers_keeps_identity_when_fallback_has_none() hidden_params_before = wrapper._hidden_params fallback = object() - wrapper.adopt_fallback_response_headers( - fallback, Router._prepare_fallback_hidden_params(fallback) - ) + wrapper.adopt_fallback_response_headers(fallback, Router._prepare_fallback_hidden_params(fallback)) assert wrapper._response_headers is None assert wrapper._hidden_params is hidden_params_before @@ -5134,9 +5149,7 @@ async def test_set_response_headers_is_the_only_complexity_header_source_for_pro **additional_headers, ) - assert not { - key for key in proxy_headers if key.startswith("x-litellm-complexity-router-") - } + assert not {key for key in proxy_headers if key.startswith("x-litellm-complexity-router-")} @pytest.mark.asyncio @@ -6312,11 +6325,7 @@ def _make_responses_iterator( BaseResponsesAPIStreamingIterator, ) - base = ( - LiteLLMCompletionStreamingIterator - if bridge - else BaseResponsesAPIStreamingIterator - ) + base = LiteLLMCompletionStreamingIterator if bridge else BaseResponsesAPIStreamingIterator class _Iter(base): def __init__(self): @@ -6416,9 +6425,7 @@ async def test_aresponses_streaming_iterator_fallback(): BaseResponsesAPIStreamingIterator, ) - router = _make_router_with_fallback( - "anthropic/claude-sonnet-4-6", "vertex_ai/claude-sonnet-4-6" - ) + router = _make_router_with_fallback("anthropic/claude-sonnet-4-6", "vertex_ai/claude-sonnet-4-6") src = _make_responses_iterator( chunks=[MagicMock(type="response.created")], error=MidStreamFallbackError( @@ -6893,9 +6900,9 @@ async def test_aresponses_streaming_iterator_writes_litellm_metadata_on_fallback fbk = mock_fallback_utils.call_args.kwargs["kwargs"] assert "litellm_metadata" in fbk, "wrong metadata_variable_name" assert fbk["litellm_metadata"]["model_group"] == "gpt-4" - assert "model_group" not in fbk.get( - "metadata", {} - ), "model_group leaked into 'metadata' instead of 'litellm_metadata'" + assert "model_group" not in fbk.get("metadata", {}), ( + "model_group leaked into 'metadata' instead of 'litellm_metadata'" + ) @pytest.mark.asyncio @@ -7462,7 +7469,6 @@ async def test_aresponses_background_stream_surfaces_a_later_error_instead_of_a_ assert backup_route.call_count == 0 - _BACKGROUND_TEXT_DELTA: Final = { "type": "response.output_text.delta", "sequence_number": 2, @@ -7533,6 +7539,7 @@ async def test_aresponses_background_stream_content_policy_error_reads_the_same_ assert primary_route.call_count == 2 assert backup_route.call_count == 0 + @pytest.mark.asyncio async def test_aresponses_streaming_iterator_partial_content_injects_continuation(): """Mid-stream error: preserve partial assistant output and request continuation in a user turn.""" @@ -7610,9 +7617,7 @@ async def test_aresponses_streaming_iterator_combines_partial_usage(): fallback_response_object = ResponsesAPIResponse( id="resp_test", created_at=0, model="gpt-4", object="response", output=[] ) - fallback_response_object.usage = ResponseAPIUsage( - input_tokens=20, output_tokens=15, total_tokens=35 - ) + fallback_response_object.usage = ResponseAPIUsage(input_tokens=20, output_tokens=15, total_tokens=35) fallback_event = ResponseCompletedEvent( type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, response=fallback_response_object, @@ -7621,9 +7626,7 @@ async def test_aresponses_streaming_iterator_combines_partial_usage(): with ( patch( "litellm.main.stream_chunk_builder", - return_value=SimpleNamespace( - usage=SimpleNamespace(prompt_tokens=10, completion_tokens=4) - ), + return_value=SimpleNamespace(usage=SimpleNamespace(prompt_tokens=10, completion_tokens=4)), ), patch.object( router, @@ -7924,9 +7927,7 @@ def test_pre_call_checks_skips_token_count_without_max_input_tokens(monkeypatch) monkeypatch.setattr(router, "_core_get_router_model_info", lambda **kwargs: {}) calls = [] - monkeypatch.setattr( - litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000 - ) + monkeypatch.setattr(litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -7954,14 +7955,10 @@ def test_pre_call_checks_counts_once_and_filters_on_max_input_tokens(monkeypatch ], enable_pre_call_checks=True, ) - monkeypatch.setattr( - router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} - ) + monkeypatch.setattr(router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 5}) calls = [] - monkeypatch.setattr( - litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000 - ) + monkeypatch.setattr(litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -8068,14 +8065,10 @@ def test_pre_call_checks_uses_precounted_tokens(monkeypatch): ], enable_pre_call_checks=True, ) - monkeypatch.setattr( - router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} - ) + monkeypatch.setattr(router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 5}) calls = [] - monkeypatch.setattr( - litellm, "token_counter", lambda *a, **k: calls.append(1) or 1 - ) + monkeypatch.setattr(litellm, "token_counter", lambda *a, **k: calls.append(1) or 1) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -8102,9 +8095,7 @@ async def test_async_get_healthy_deployments_counts_tokens_off_the_event_loop(mo ], enable_pre_call_checks=True, ) - monkeypatch.setattr( - router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 1_000_000} - ) + monkeypatch.setattr(router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 1_000_000}) counting_threads = [] monkeypatch.setattr( @@ -8200,14 +8191,10 @@ def test_pre_call_checks_does_not_recount_inline_after_an_off_loop_failure(monke ], enable_pre_call_checks=True, ) - monkeypatch.setattr( - router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} - ) + monkeypatch.setattr(router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 5}) calls = [] - monkeypatch.setattr( - litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000 - ) + monkeypatch.setattr(litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -8235,9 +8222,7 @@ async def test_async_get_healthy_deployments_never_recounts_on_the_loop(monkeypa ], enable_pre_call_checks=True, ) - monkeypatch.setattr( - router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} - ) + monkeypatch.setattr(router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 5}) counting_threads = [] @@ -8272,9 +8257,7 @@ async def test_acount_pre_call_check_tokens_leaves_the_event_loop_free(monkeypat ], enable_pre_call_checks=True, ) - monkeypatch.setattr( - router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} - ) + monkeypatch.setattr(router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 5}) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -8310,9 +8293,7 @@ async def test_acount_pre_call_check_tokens_skips_without_max_input_tokens(monke monkeypatch.setattr(router, "_core_get_router_model_info", lambda **kwargs: {}) calls = [] - monkeypatch.setattr( - litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000 - ) + monkeypatch.setattr(litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000) count = await router._acount_pre_call_check_tokens( model="m", @@ -8340,9 +8321,7 @@ def test_pre_call_checks_counts_tokens_from_responses_input_string(monkeypatch): ], enable_pre_call_checks=True, ) - monkeypatch.setattr( - router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 1} - ) + monkeypatch.setattr(router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 1}) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -8367,9 +8346,7 @@ def test_pre_call_checks_counts_tokens_from_responses_input_list(monkeypatch): ], enable_pre_call_checks=True, ) - monkeypatch.setattr( - router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 1} - ) + monkeypatch.setattr(router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 1}) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -8411,9 +8388,7 @@ def test_pre_call_checks_counts_responses_instructions_tokens(monkeypatch): ) assert with_instructions_tokens > input_only_tokens - monkeypatch.setattr( - router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": input_only_tokens} - ) + monkeypatch.setattr(router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": input_only_tokens}) with pytest.raises(litellm.ContextWindowExceededError): router._pre_call_checks( model="m", @@ -8524,7 +8499,9 @@ def test_pre_call_checks_counts_anthropic_system_tokens(monkeypatch, system): messages = [{"role": "user", "content": "hi"}] messages_only_tokens = router._count_pre_call_check_tokens(messages=messages, input=None) - monkeypatch.setattr(router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": messages_only_tokens}) + monkeypatch.setattr( + router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": messages_only_tokens} + ) assert len(router._pre_call_checks(model="m", healthy_deployments=deployments, messages=messages)) == 1 with pytest.raises(litellm.ContextWindowExceededError): @@ -8604,7 +8581,7 @@ def test_count_pre_call_check_tokens_across_api_surfaces(): assert string_input_tokens > 0 assert list_input_tokens > 0 - with pytest.raises(ValueError, match='Either messages or input must be provided to count tokens'): + with pytest.raises(ValueError, match="Either messages or input must be provided to count tokens"): router._count_pre_call_check_tokens(messages=None, input=None) @@ -8642,9 +8619,7 @@ def test_pre_call_checks_no_messages_or_input_does_not_crash(monkeypatch): ], enable_pre_call_checks=True, ) - monkeypatch.setattr( - router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} - ) + monkeypatch.setattr(router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 5}) counted: list[dict] = [] original = router._count_pre_call_check_tokens @@ -8735,9 +8710,7 @@ def test_get_deployment_model_info_base_model_flow(): } # Test Case 1: Base model flow with custom model info that has base_model - with patch.object( - litellm, "model_cost", {"test-custom-model": mock_custom_model_info} - ): + with patch.object(litellm, "model_cost", {"test-custom-model": mock_custom_model_info}): with patch.object(litellm, "get_model_info") as mock_get_model_info: # Configure mock returns mock_get_model_info.side_effect = lambda model: { @@ -8745,15 +8718,11 @@ def test_get_deployment_model_info_base_model_flow(): "test-model": mock_litellm_model_name_info, }.get(model) - result = router.get_deployment_model_info( - model_id="test-custom-model", model_name="test-model" - ) + result = router.get_deployment_model_info(model_id="test-custom-model", model_name="test-model") # Verify that get_model_info was called for both base model and model name assert mock_get_model_info.call_count == 2 - mock_get_model_info.assert_any_call( - model="gpt-3.5-turbo" - ) # base model call + mock_get_model_info.assert_any_call(model="gpt-3.5-turbo") # base model call mock_get_model_info.assert_any_call(model="test-model") # model name call # Verify the result contains merged information @@ -8764,26 +8733,18 @@ def test_get_deployment_model_info_base_model_flow(): # 2. The result of step 1 gets merged into litellm_model_name_info (custom+base override litellm) # Fields from custom model (should override base model values) - assert ( - result["input_cost_per_token"] == 0.001 - ) # From custom model (overrides base 0.0015) - assert ( - result["output_cost_per_token"] == 0.002 - ) # From custom model (same as base) + assert result["input_cost_per_token"] == 0.001 # From custom model (overrides base 0.0015) + assert result["output_cost_per_token"] == 0.002 # From custom model (same as base) assert result["custom_field"] == "custom_value" # From custom model # Fields from base model that weren't overridden by custom assert result["max_tokens"] == 4096 # From base model assert result["litellm_provider"] == "openai" # From base model - assert ( - result["mode"] == "chat" - ) # From base model (overrides litellm "completion") + assert result["mode"] == "chat" # From base model (overrides litellm "completion") # The key field comes from base model since both base and litellm have it # and base model info overrides litellm model name info in final merge - assert ( - result["key"] == "gpt-3.5-turbo" - ) # From base model (overrides litellm key) + assert result["key"] == "gpt-3.5-turbo" # From base model (overrides litellm key) # Test Case 2: Custom model info without base_model mock_custom_model_info_no_base = { @@ -8802,9 +8763,7 @@ def test_get_deployment_model_info_base_model_flow(): "test-model": mock_litellm_model_name_info, }.get(model) - result = router.get_deployment_model_info( - model_id="test-custom-model-no-base", model_name="test-model" - ) + result = router.get_deployment_model_info(model_id="test-custom-model-no-base", model_name="test-model") # Should only call get_model_info once for model name (no base model) assert mock_get_model_info.call_count == 1 @@ -8824,9 +8783,7 @@ def test_get_deployment_model_info_base_model_flow(): "test-model": mock_litellm_model_name_info, }.get(model) - result = router.get_deployment_model_info( - model_id="non-existent-model", model_name="test-model" - ) + result = router.get_deployment_model_info(model_id="non-existent-model", model_name="test-model") # Should only call get_model_info once for model name assert mock_get_model_info.call_count == 1 @@ -8859,9 +8816,7 @@ def test_get_deployment_model_info_base_model_flow(): mock_get_model_info.side_effect = mock_get_model_info_side_effect - result = router.get_deployment_model_info( - model_id="test-custom-model-invalid", model_name="test-model" - ) + result = router.get_deployment_model_info(model_id="test-custom-model-invalid", model_name="test-model") # Should handle exception gracefully and still return merged result assert result is not None @@ -8870,12 +8825,8 @@ def test_get_deployment_model_info_base_model_flow(): # Test Case 5: Both model_cost.get() and get_model_info() return None with patch.object(litellm, "model_cost", {}): - with patch.object( - litellm, "get_model_info", side_effect=Exception("Not found") - ): - result = router.get_deployment_model_info( - model_id="non-existent", model_name="non-existent" - ) + with patch.object(litellm, "get_model_info", side_effect=Exception("Not found")): + result = router.get_deployment_model_info(model_id="non-existent", model_name="non-existent") # Should return None when no model info is found assert result is None @@ -8898,9 +8849,7 @@ def test_get_deployment_model_info_base_model_flow(): # Model NOT in built-in cost map — raise exception mock_get_model_info.side_effect = Exception("Model not in cost map") - result = router.get_deployment_model_info( - model_id="custom-model-id", model_name="unknown-model" - ) + result = router.get_deployment_model_info(model_id="custom-model-id", model_name="unknown-model") # Should return custom_model_info even when litellm_model_name_model_info is None assert result is not None @@ -8936,15 +8885,11 @@ def test_get_deployment_model_info_base_model_flow(): mock_get_model_info.side_effect = get_info_side_effect - result = router.get_deployment_model_info( - model_id="custom-with-base", model_name="unknown-model" - ) + result = router.get_deployment_model_info(model_id="custom-with-base", model_name="unknown-model") # Should return custom_model_info merged with base model info assert result is not None - assert ( - result["input_cost_per_token"] == 0.01 - ) # From custom (overrides base) + assert result["input_cost_per_token"] == 0.01 # From custom (overrides base) assert result["max_tokens"] == 8192 # From base model assert result["litellm_provider"] == "openai" # From base model @@ -8991,18 +8936,14 @@ def test_get_deployment_model_info_base_model_merge_priority(): "litellm_only_field": "litellm_value", } - with patch.object( - litellm, "model_cost", {"custom-model-id": mock_custom_model_info} - ): + with patch.object(litellm, "model_cost", {"custom-model-id": mock_custom_model_info}): with patch.object(litellm, "get_model_info") as mock_get_model_info: mock_get_model_info.side_effect = lambda model: { "gpt-4": mock_base_model_info, "test-model": mock_litellm_model_name_info, }.get(model) - result = router.get_deployment_model_info( - model_id="custom-model-id", model_name="test-model" - ) + result = router.get_deployment_model_info(model_id="custom-model-id", model_name="test-model") assert result is not None @@ -9012,29 +8953,17 @@ def test_get_deployment_model_info_base_model_merge_priority(): # 3. Result from steps 1-2 overrides litellm_model_name_info # Fields that should come from custom model info (highest priority) - assert ( - result["input_cost_per_token"] == 0.01 - ) # From custom model (overrides base 0.03) - assert ( - result["max_tokens"] == 8000 - ) # From custom model (overrides base 4096) + assert result["input_cost_per_token"] == 0.01 # From custom model (overrides base 0.03) + assert result["max_tokens"] == 8000 # From custom model (overrides base 4096) assert result["custom_only_field"] == "custom_value" # From custom model # Fields that should come from base model (not overridden by custom) - assert ( - result["output_cost_per_token"] == 0.06 - ) # From base model (not in custom) - assert ( - result["litellm_provider"] == "openai" - ) # From base model (not in custom) - assert ( - result["base_only_field"] == "base_value" - ) # From base model (not in custom) + assert result["output_cost_per_token"] == 0.06 # From base model (not in custom) + assert result["litellm_provider"] == "openai" # From base model (not in custom) + assert result["base_only_field"] == "base_value" # From base model (not in custom) # Fields that should come from litellm model name info (not overridden by custom+base) - assert ( - result["mode"] == "completion" - ) # From litellm model name info (not in custom or base) + assert result["mode"] == "completion" # From litellm model name info (not in custom or base) assert ( result["litellm_only_field"] == "litellm_value" ) # From litellm model name info (not in custom or base) @@ -9051,7 +8980,11 @@ def test_get_deployment_model_info_base_model_merge_priority(): [ ( "gpt", - {"model": "azure_ai/gpt-5.4-mini", "api_base": "https://my-resource.services.ai.azure.com", "api_key": "key"}, + { + "model": "azure_ai/gpt-5.4-mini", + "api_base": "https://my-resource.services.ai.azure.com", + "api_key": "key", + }, "gpt/openai/deployments/gpt-5.4-mini/chat/completions", "gpt-5.4-mini/openai/deployments/gpt-5.4-mini/chat/completions", ), @@ -9106,10 +9039,9 @@ def test_add_deployment_model_to_endpoint_for_llm_passthrough_route(): model="special-bedrock-model", model_name="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", ) - assert ( - result["endpoint"] - == "/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke" - ), f"Expected '/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke', got '{result['endpoint']}'" + assert result["endpoint"] == "/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke", ( + f"Expected '/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke', got '{result['endpoint']}'" + ) # Test Case 2: Bedrock invoke-with-response-stream endpoint kwargs = { @@ -9121,10 +9053,9 @@ def test_add_deployment_model_to_endpoint_for_llm_passthrough_route(): model="special-bedrock-model", model_name="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", ) - assert ( - result["endpoint"] - == "/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke-with-response-stream" - ), f"Expected streaming endpoint with stripped prefix, got '{result['endpoint']}'" + assert result["endpoint"] == "/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke-with-response-stream", ( + f"Expected streaming endpoint with stripped prefix, got '{result['endpoint']}'" + ) # Test Case 3: Bedrock converse endpoint kwargs = { @@ -9136,9 +9067,9 @@ def test_add_deployment_model_to_endpoint_for_llm_passthrough_route(): model="bedrock-model", model_name="bedrock/us.meta.llama3-8b-instruct-v1:0", ) - assert ( - result["endpoint"] == "/model/us.meta.llama3-8b-instruct-v1:0/converse" - ), f"Expected '/model/us.meta.llama3-8b-instruct-v1:0/converse', got '{result['endpoint']}'" + assert result["endpoint"] == "/model/us.meta.llama3-8b-instruct-v1:0/converse", ( + f"Expected '/model/us.meta.llama3-8b-instruct-v1:0/converse', got '{result['endpoint']}'" + ) # Test Case 4: Bedrock provider prefix auto-detected from model_name kwargs = { @@ -9149,9 +9080,9 @@ def test_add_deployment_model_to_endpoint_for_llm_passthrough_route(): model="router-model", model_name="bedrock/us.meta.llama3-8b-instruct-v1:0", ) - assert ( - result["endpoint"] == "/model/us.meta.llama3-8b-instruct-v1:0/invoke" - ), f"Expected '/model/us.meta.llama3-8b-instruct-v1:0/invoke', got '{result['endpoint']}'" + assert result["endpoint"] == "/model/us.meta.llama3-8b-instruct-v1:0/invoke", ( + f"Expected '/model/us.meta.llama3-8b-instruct-v1:0/invoke', got '{result['endpoint']}'" + ) def test_update_kwargs_with_deployment_uses_pass_through_request_timeout(): @@ -9288,14 +9219,10 @@ async def test_router_acompletion_with_unknown_model_and_default_fallback(): # Initialize the router with a default fallback router = litellm.Router(model_list=model_list, default_fallbacks=["gpt-4o"]) - messages = [ - {"role": "user", "content": "This call should succeed by falling back."} - ] + messages = [{"role": "user", "content": "This call should succeed by falling back."}] # Call completion with a model name that is NOT in the model_list - response = await router.acompletion( - model="completely-unknown-model", messages=messages - ) + response = await router.acompletion(model="completely-unknown-model", messages=messages) # Check that the call did not fail and we received a valid response object. assert response is not None @@ -9460,7 +9387,9 @@ def test_get_wildcard_deployment_usable_by_team_prefers_the_team_pattern(): pytest.param(False, "rwild-chat", "team-r", ("team-r-wildcard",), id="team-pattern-when-no-proxy-wide-match"), pytest.param(False, "rwild-chat", "team-b", (), id="another-teams-pattern-never-serves"), pytest.param(False, "rwild-chat", None, (), id="no-team-no-team-pattern"), - pytest.param(True, "rwild-chat", "team-r", ("team-r-wildcard",), id="team-pattern-beats-the-proxy-wide-pattern"), + pytest.param( + True, "rwild-chat", "team-r", ("team-r-wildcard",), id="team-pattern-beats-the-proxy-wide-pattern" + ), pytest.param(True, "rwild-chat", "team-b", ("catch-all",), id="another-team-falls-through-to-proxy-wide"), pytest.param(True, "ptu-chat", "team-r", ("ptu-wildcard",), id="only-the-most-specific-proxy-wide-pattern"), ], @@ -9469,7 +9398,13 @@ def test_wildcard_route_deployments_is_the_one_pattern_routing_serves_from( with_catch_all: bool, requested: str, team_id: str | None, expected_ids: tuple[str, ...] ): catch_all: Final = ( - ({"model_name": "*", "litellm_params": {"model": "openai/*", "api_key": "sk-open"}, "model_info": {"id": "catch-all"}},) + ( + { + "model_name": "*", + "litellm_params": {"model": "openai/*", "api_key": "sk-open"}, + "model_info": {"id": "catch-all"}, + }, + ) if with_catch_all else () ) @@ -9490,7 +9425,9 @@ def test_wildcard_route_deployments_is_the_one_pattern_routing_serves_from( ) served: Final = router.wildcard_route_deployments(requested, team_id) - early: Final = router._try_early_resolve_deployments_for_model_not_in_names(model=requested, request_team_id=team_id) + early: Final = router._try_early_resolve_deployments_for_model_not_in_names( + model=requested, request_team_id=team_id + ) early_deployments: Final = early[1] if early is not None and isinstance(early[1], list) else [] assert tuple(Deployment.model_validate(d).model_info.id for d in served) == expected_ids @@ -9553,15 +9490,10 @@ def test_get_deployment_credentials_with_provider_aws_bedrock_runtime_endpoint() ], ) - credentials = router.get_deployment_credentials_with_provider( - model_id="bedrock-claude-model" - ) + credentials = router.get_deployment_credentials_with_provider(model_id="bedrock-claude-model") assert credentials is not None - assert ( - credentials["aws_bedrock_runtime_endpoint"] - == "https://bedrock-runtime.us-east-1.amazonaws.com" - ) + assert credentials["aws_bedrock_runtime_endpoint"] == "https://bedrock-runtime.us-east-1.amazonaws.com" assert credentials["aws_access_key_id"] == "test-access-key" assert credentials["aws_secret_access_key"] == "test-secret-key" assert credentials["aws_region_name"] == "us-east-1" @@ -9588,9 +9520,7 @@ def test_get_deployment_credentials_with_provider_includes_bucket_name(): ], ) - credentials = router.get_deployment_credentials_with_provider( - model_id="vertex-gemini" - ) + credentials = router.get_deployment_credentials_with_provider(model_id="vertex-gemini") assert credentials is not None assert credentials["gcs_bucket_name"] == "my-batch-bucket" @@ -9675,9 +9605,7 @@ def test_get_deployment_credentials_with_provider_resolves_credential_name(): ], ) - credentials = router.get_deployment_credentials_with_provider( - model_id="azure-gpt-4" - ) + credentials = router.get_deployment_credentials_with_provider(model_id="azure-gpt-4") assert credentials is not None assert credentials["api_key"] == "resolved-api-key" @@ -9714,9 +9642,7 @@ def test_get_deployment_credentials_with_provider_bedrock_batch_fields(): ], ) - credentials = router.get_deployment_credentials_with_provider( - model_id="bedrock-batch-model" - ) + credentials = router.get_deployment_credentials_with_provider(model_id="bedrock-batch-model") assert credentials is not None assert credentials["custom_llm_provider"] == "bedrock" @@ -9760,9 +9686,7 @@ def test_get_deployment_credentials_with_provider_preserves_aws_auth_params(): ], ) - credentials = router.get_deployment_credentials_with_provider( - model_id="bedrock-batch-model" - ) + credentials = router.get_deployment_credentials_with_provider(model_id="bedrock-batch-model") assert credentials is not None for key, value in aws_auth_params.items(): @@ -9832,15 +9756,11 @@ def test_get_deployment_credentials_with_provider_team_wildcard_priority(): ], ) - team_credentials = router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-1" - ) + team_credentials = router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") assert team_credentials is not None assert team_credentials["api_key"] == "team-key" - global_credentials = router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2" - ) + global_credentials = router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2") assert global_credentials is not None assert global_credentials["api_key"] == "global-key" @@ -9881,15 +9801,11 @@ def test_get_deployment_credentials_with_provider_skips_other_team_deployment(): assert other_team_credentials is not None assert other_team_credentials["vertex_project"] == "shared-project" - unscoped_credentials = router.get_deployment_credentials_with_provider( - model_id="gemini-2.5-pro" - ) + unscoped_credentials = router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro") assert unscoped_credentials is not None assert unscoped_credentials["vertex_project"] == "shared-project" - owner_credentials = router.get_deployment_credentials_with_provider( - model_id="gemini-2.5-pro", team_id="team-b" - ) + owner_credentials = router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro", team_id="team-b") assert owner_credentials is not None assert owner_credentials["vertex_project"] == "team-b-project" @@ -9916,16 +9832,8 @@ def test_get_deployment_credentials_with_provider_no_fallback_to_other_team_only ], ) - assert ( - router.get_deployment_credentials_with_provider( - model_id="gemini-2.5-pro", team_id="team-a" - ) - is None - ) - assert ( - router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro") - is None - ) + assert router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro", team_id="team-a") is None + assert router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro") is None def test_deployment_usable_by_team_helpers(): @@ -9965,9 +9873,7 @@ def test_deployment_usable_by_team_helpers(): assert router._deployment_usable_by_team(shared, "team-a") is True assert router._deployment_usable_by_team(shared, None) is True - picked = router._get_model_group_deployment_usable_by_team( - model_group_name="gemini-2.5-pro", team_id="team-a" - ) + picked = router._get_model_group_deployment_usable_by_team(model_group_name="gemini-2.5-pro", team_id="team-a") assert picked is not None assert picked.litellm_params.vertex_project == "shared-project" @@ -9977,12 +9883,7 @@ def test_deployment_usable_by_team_helpers(): assert owner_picked is not None assert owner_picked.litellm_params.vertex_project == "team-b-project" - assert ( - router._get_model_group_deployment_usable_by_team( - model_group_name="unknown-model", team_id="team-a" - ) - is None - ) + assert router._get_model_group_deployment_usable_by_team(model_group_name="unknown-model", team_id="team-a") is None def test_deployment_usable_by_team_uses_dynamic_model_info_get(): @@ -10022,9 +9923,7 @@ def test_get_deployment_credentials_with_provider_skips_other_team_wildcard(): assert other_team_credentials is not None assert other_team_credentials["api_key"] == "global-key" - owner_credentials = router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-b" - ) + owner_credentials = router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-b") assert owner_credentials is not None assert owner_credentials["api_key"] == "team-b-key" @@ -10036,21 +9935,11 @@ def test_team_wildcard_credentials_not_usable_after_delete_deployment(): """ router = litellm.Router(model_list=[_team_wildcard_model(api_key="old-key")]) - assert ( - router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-1" - ) - is not None - ) + assert router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") is not None router.delete_deployment(id="team-wildcard-id") - assert ( - router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-1" - ) - is None - ) + assert router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") is None def test_global_wildcard_pattern_router_evicts_stale_entry_on_upsert_and_delete(): @@ -10123,22 +10012,13 @@ def test_team_wildcard_credentials_refreshed_on_upsert_and_set_model_list(): router = litellm.Router(model_list=[_team_wildcard_model(api_key="old-key")]) - router.upsert_deployment( - deployment=Deployment(**_team_wildcard_model(api_key="new-key")) - ) - credentials = router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-1" - ) + router.upsert_deployment(deployment=Deployment(**_team_wildcard_model(api_key="new-key"))) + credentials = router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") assert credentials is not None assert credentials["api_key"] == "new-key" router.set_model_list(model_list=[]) - assert ( - router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-1" - ) - is None - ) + assert router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") is None def test_get_available_guardrail_single_deployment(): @@ -10326,9 +10206,7 @@ async def test_anthropic_messages_call_type_is_cached(): startTime=1234567890.0, endTime=1234567891.0, completionStartTime=1234567890.5, - model_map_information=StandardLoggingModelInformation( - model_map_key="gpt-3.5-turbo", model_map_value=None - ), + model_map_information=StandardLoggingModelInformation(model_map_key="gpt-3.5-turbo", model_map_value=None), model="gpt-3.5-turbo", custom_llm_provider="anthropic", model_id="model-123", @@ -10408,12 +10286,8 @@ async def test_anthropic_messages_call_type_is_cached(): ) # This assertion will FAIL if anthropic_messages is filtered out - assert ( - cached_result is not None - ), "Model ID should be cached for anthropic_messages call type" - assert ( - cached_result["model_id"] == test_model_id - ), f"Expected {test_model_id}, got {cached_result['model_id']}" + assert cached_result is not None, "Model ID should be cached for anthropic_messages call type" + assert cached_result["model_id"] == test_model_id, f"Expected {test_model_id}, got {cached_result['model_id']}" def test_update_kwargs_with_deployment_propagates_model_tags(): @@ -10438,9 +10312,7 @@ def test_update_kwargs_with_deployment_propagates_model_tags(): ) kwargs: dict = {"metadata": {}} - deployment = router.get_deployment_by_model_group_name( - model_group_name="gpt-4o-mini" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-4o-mini") router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) # Deployment tags should be propagated to kwargs metadata @@ -10489,9 +10361,7 @@ def test_update_kwargs_with_deployment_merges_tags_without_duplicates(): # Simulate request that already has tags (from request body or key/team level) kwargs: dict = {"metadata": {"tags": ["user-tag", "shared-tag"]}} - deployment = router.get_deployment_by_model_group_name( - model_group_name="gpt-4o-mini" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-4o-mini") router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) # Both sources should be merged, no duplicates @@ -10518,9 +10388,7 @@ def test_update_kwargs_with_deployment_no_tags(): ) kwargs: dict = {"metadata": {}} - deployment = router.get_deployment_by_model_group_name( - model_group_name="gpt-4o-mini" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-4o-mini") router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) # No tags key should be added if deployment has no tags @@ -10602,9 +10470,7 @@ def test_update_kwargs_with_deployment_merges_tools(): }, ], } - deployment = router.get_deployment_by_model_group_name( - model_group_name="o3-deep-research" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="o3-deep-research") router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) # Tools should be merged: deployment first, then request @@ -10635,9 +10501,7 @@ def test_update_kwargs_with_deployment_merge_tools_deployment_only(): ) kwargs: dict = {"metadata": {}} - deployment = router.get_deployment_by_model_group_name( - model_group_name="o3-deep-research" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="o3-deep-research") router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) assert kwargs["tools"] == [{"type": "web_search"}] @@ -10666,9 +10530,7 @@ def test_update_kwargs_with_deployment_merge_tools_request_overrides_tool_choice "metadata": {}, "tool_choice": "none", } - deployment = router.get_deployment_by_model_group_name( - model_group_name="o3-deep-research" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="o3-deep-research") router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) # Request tool_choice should be preserved (merged tools still applied) @@ -10860,12 +10722,8 @@ def test_update_kwargs_with_deployment_model_info_in_litellm_metadata(): ) kwargs: dict = {} - deployment = router.get_deployment_by_model_group_name( - model_group_name="claude-sonnet-4" - ) - router._update_kwargs_with_deployment( - deployment=deployment, kwargs=kwargs, function_name="generic_api_call" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="claude-sonnet-4") + router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, function_name="generic_api_call") assert "litellm_metadata" in kwargs model_info = kwargs["litellm_metadata"]["model_info"] @@ -10897,12 +10755,8 @@ def test_update_kwargs_with_deployment_model_info_in_metadata(): ) kwargs: dict = {} - deployment = router.get_deployment_by_model_group_name( - model_group_name="claude-sonnet-4" - ) - router._update_kwargs_with_deployment( - deployment=deployment, kwargs=kwargs, function_name=None - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="claude-sonnet-4") + router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, function_name=None) assert "metadata" in kwargs model_info = kwargs["metadata"]["model_info"] @@ -11108,6 +10962,7 @@ async def test_acompletion_streaming_iterator_does_not_log_success_on_terminal_f initial_kwargs=dict(initial_kwargs), ) collected = [] + async def _drain(): async for chunk in result: collected.append(chunk) @@ -11134,6 +10989,7 @@ async def test_acompletion_streaming_iterator_does_not_log_success_on_terminal_f initial_kwargs=dict(initial_kwargs), ) collected = [] + async def _drain(): async for chunk in result: collected.append(chunk) @@ -11278,7 +11134,12 @@ def _shared_ptu_deployment(model_name: str, deployment_id: str, **litellm_params """A one-PTU Azure deployment whose whole reservation is team-a's share.""" return { "model_name": model_name, - "litellm_params": {"model": "azure/gpt-4.1", "api_key": "fake", "api_base": "https://ptu.example", **litellm_params}, + "litellm_params": { + "model": "azure/gpt-4.1", + "api_key": "fake", + "api_base": "https://ptu.example", + **litellm_params, + }, "model_info": { "id": deployment_id, "base_model": "azure/gpt-4.1", @@ -11388,23 +11249,17 @@ def test_multiregion_team_deployments_unique_model_names(): assert len(deployments) == 0 # With team_id: O(n) scan finds BOTH regional deployments - deployments = router._get_all_deployments( - model_name="claude-sonnet", team_id="metis-team" - ) + deployments = router._get_all_deployments(model_name="claude-sonnet", team_id="metis-team") assert len(deployments) == 2 deployment_names = {d["model_name"] for d in deployments} assert deployment_names == {"metis-claude-us-east-1", "metis-claude-us-west-2"} # Each deployment has a unique ID (critical for cooldown/retry to work) deployment_ids = {d["model_info"]["id"] for d in deployments} - assert ( - len(deployment_ids) == 2 - ), "Each deployment must have a unique ID for cooldown tracking" + assert len(deployment_ids) == 2, "Each deployment must have a unique ID for cooldown tracking" # Wrong team: returns nothing - deployments = router._get_all_deployments( - model_name="claude-sonnet", team_id="other-team" - ) + deployments = router._get_all_deployments(model_name="claude-sonnet", team_id="other-team") assert len(deployments) == 0 @@ -11449,12 +11304,8 @@ async def test_multiregion_team_failover_between_regions(): ) # Verify the router finds both deployments for the team - deployments = router._get_all_deployments( - model_name="claude-sonnet", team_id="metis-team" - ) - assert ( - len(deployments) == 2 - ), "Router must find both regional deployments by team_public_model_name" + deployments = router._get_all_deployments(model_name="claude-sonnet", team_id="metis-team") + assert len(deployments) == 2, "Router must find both regional deployments by team_public_model_name" # Make a normal request — should succeed from one of the regions response = await router.acompletion( @@ -11579,9 +11430,7 @@ def test_explicit_model_access_does_not_force_access_group_filtering(): }, ) - deployment_groups = [ - d.get("model_info", {}).get("access_groups") for d in deployments - ] + deployment_groups = [d.get("model_info", {}).get("access_groups") for d in deployments] assert ["AG1"] in deployment_groups assert ["AG2"] in deployment_groups @@ -11626,9 +11475,7 @@ def test_access_group_filter_empty_does_not_bypass_via_litellm_model_fallback( orig_groups = router.get_model_access_groups - def fake_get_model_access_groups( - model_name=None, model_access_group=None, team_id=None - ): + def fake_get_model_access_groups(model_name=None, model_access_group=None, team_id=None): if model_name == "gpt-5" and model_access_group is None: return {"AG1": ["gpt-5"], "AG2": ["gpt-5"]} return orig_groups( @@ -11703,9 +11550,7 @@ def test_access_group_block_does_not_silently_use_default_fallback_model( orig_groups = router.get_model_access_groups - def fake_get_model_access_groups( - model_name=None, model_access_group=None, team_id=None - ): + def fake_get_model_access_groups(model_name=None, model_access_group=None, team_id=None): if model_name == "gpt-5" and model_access_group is None: return {"AG1": ["gpt-5"], "AG2": ["gpt-5"]} return orig_groups( @@ -11772,9 +11617,7 @@ def test_access_group_block_via_litellm_model_branch_does_not_use_default_fallba orig_groups = router.get_model_access_groups - def fake_get_model_access_groups( - model_name=None, model_access_group=None, team_id=None - ): + def fake_get_model_access_groups(model_name=None, model_access_group=None, team_id=None): if model_name == "gpt-5" and model_access_group is None: return {"AG1": ["gpt-5"], "AG2": ["gpt-5"]} return orig_groups( @@ -11829,9 +11672,7 @@ def test_try_early_resolve_deployments_for_model_not_in_names(): ) assert ( - router_in_names._try_early_resolve_deployments_for_model_not_in_names( - model="gpt-5", request_team_id=None - ) + router_in_names._try_early_resolve_deployments_for_model_not_in_names(model="gpt-5", request_team_id=None) is None ) assert ( @@ -11853,10 +11694,8 @@ def test_try_early_resolve_deployments_for_model_not_in_names(): ] ) - pattern_result = ( - pattern_router._try_early_resolve_deployments_for_model_not_in_names( - model="openai/gpt-4o-mini", request_team_id=None - ) + pattern_result = pattern_router._try_early_resolve_deployments_for_model_not_in_names( + model="openai/gpt-4o-mini", request_team_id=None ) assert pattern_result is not None resolved_model, pattern_deployments = pattern_result @@ -11882,10 +11721,8 @@ def test_try_early_resolve_deployments_for_model_not_in_names(): }, } - default_result = ( - default_router._try_early_resolve_deployments_for_model_not_in_names( - model="brand-new-model", request_team_id=None - ) + default_result = default_router._try_early_resolve_deployments_for_model_not_in_names( + model="brand-new-model", request_team_id=None ) assert default_result is not None resolved_model, default_deployment = default_result @@ -11893,10 +11730,7 @@ def test_try_early_resolve_deployments_for_model_not_in_names(): assert isinstance(default_deployment, dict) assert default_deployment["litellm_params"]["model"] == "brand-new-model" # The original default_deployment must not be mutated. - assert ( - default_router.default_deployment["litellm_params"]["model"] - == "openai/will-be-overridden" - ) + assert default_router.default_deployment["litellm_params"]["model"] == "openai/will-be-overridden" def _router_with_two_deployments(blocked_flags): @@ -11944,10 +11778,7 @@ def _seed_unhealthy_states(router, unhealthy_ids, timestamp=None): ts = timestamp if timestamp is not None else time.time() router.health_state_cache.set_deployment_health_states( - { - uid: {"is_healthy": False, "timestamp": ts, "reason": "test_unhealthy"} - for uid in unhealthy_ids - } + {uid: {"is_healthy": False, "timestamp": ts, "reason": "test_unhealthy"} for uid in unhealthy_ids} ) @@ -11981,6 +11812,7 @@ async def test_health_probe_preserves_normal_caller_policy( strict_ids: tuple[str, ...], ) -> None: import time + from litellm.types.router import AllowedFailsPolicy, RouterRateLimitError router: Final = Router( @@ -12074,9 +11906,7 @@ async def test_async_get_fully_unhealthy_model_names_noop_with_allowed_fails_pol @pytest.mark.asyncio async def test_async_get_healthy_deployments_skips_blocked_deployment(): router = _router_with_two_deployments([True, False]) - healthy, all_dep = await router._async_get_healthy_deployments( - model="gpt-4o", parent_otel_span=None - ) + healthy, all_dep = await router._async_get_healthy_deployments(model="gpt-4o", parent_otel_span=None) healthy_ids = [d["model_info"]["id"] for d in healthy] assert "dep-0" not in healthy_ids assert "dep-1" in healthy_ids @@ -12085,9 +11915,7 @@ async def test_async_get_healthy_deployments_skips_blocked_deployment(): def test_get_healthy_deployments_sync_skips_blocked_deployment(): router = _router_with_two_deployments([False, True]) - healthy, all_dep = router._get_healthy_deployments( - model="gpt-4o", parent_otel_span=None - ) + healthy, all_dep = router._get_healthy_deployments(model="gpt-4o", parent_otel_span=None) healthy_ids = [d["model_info"]["id"] for d in healthy] assert "dep-0" in healthy_ids assert "dep-1" not in healthy_ids @@ -12104,9 +11932,7 @@ def test_filter_blocked_deployments_drops_blocked_keeps_unblocked(): @pytest.mark.asyncio async def test_public_async_get_healthy_deployments_skips_blocked_on_primary_path(): router = _router_with_two_deployments([True, False]) - deployments = await router.async_get_healthy_deployments( - model="gpt-4o", request_kwargs={} - ) + deployments = await router.async_get_healthy_deployments(model="gpt-4o", request_kwargs={}) assert isinstance(deployments, list) ids = [d["model_info"]["id"] for d in deployments] assert "dep-0" not in ids @@ -12194,9 +12020,7 @@ def _router_with_two_pass_through_deployments(blocked_flags): def test_get_available_deployment_for_pass_through_skips_blocked(): router = _router_with_two_pass_through_deployments([True, False]) - deployment = router.get_available_deployment_for_pass_through( - model="gpt-4o", request_kwargs={} - ) + deployment = router.get_available_deployment_for_pass_through(model="gpt-4o", request_kwargs={}) assert deployment["model_info"]["id"] == "pt-1" @@ -12205,9 +12029,7 @@ def test_get_available_deployment_for_pass_through_raises_when_dict_blocked(): router = _router_with_two_pass_through_deployments([True, True]) with pytest.raises(litellm.ServiceUnavailableError): - router.get_available_deployment_for_pass_through( - model="pt-0", request_kwargs={} - ) + router.get_available_deployment_for_pass_through(model="pt-0", request_kwargs={}) def test_get_available_deployment_for_pass_through_names_cooldown_despite_healthy_non_pass_through(): @@ -12249,9 +12071,7 @@ def test_initialize_deployment_for_pass_through_keeps_bedrock_iam_deployment(): } ] ) - assert [m["model_info"]["id"] for m in router.get_model_list()] == [ - "bedrock-iam-pt" - ] + assert [m["model_info"]["id"] for m in router.get_model_list()] == ["bedrock-iam-pt"] def test_pass_through_deployment_api_key_resolves_via_get_credentials(): @@ -12262,12 +12082,7 @@ def test_pass_through_deployment_api_key_resolves_via_get_credentials(): router = _router_with_two_pass_through_deployments([False, False]) passthrough_router = PassthroughEndpointRouter(llm_router_getter=lambda: router) assert len(router.get_model_list()) == 2 - assert ( - passthrough_router.get_credentials( - custom_llm_provider="openai", region_name=None - ) - == "sk-fake-for-tests" - ) + assert passthrough_router.get_credentials(custom_llm_provider="openai", region_name=None) == "sk-fake-for-tests" def test_get_deployment_credentials_returns_none_for_blocked_deployment(): @@ -12303,7 +12118,9 @@ def test_is_deployment_blocked_static_helper_reflects_blocked_flag(): missing_blocked = types.SimpleNamespace() assert litellm.Router._core_is_deployment_blocked(types.SimpleNamespace(model_info=missing_blocked)) is False assert ( - litellm.Router._core_is_deployment_blocked(types.SimpleNamespace(model_info=types.SimpleNamespace(blocked=True))) + litellm.Router._core_is_deployment_blocked( + types.SimpleNamespace(model_info=types.SimpleNamespace(blocked=True)) + ) is True ) @@ -12415,9 +12232,7 @@ class TestRouterRequestTimeoutPropagation: litellm.request_timeout = original_value litellm.request_timeout_explicitly_set = original_flag - def test_request_timeout_stored_independently_when_both_set( - self, explicit_request_timeout - ): + def test_request_timeout_stored_independently_when_both_set(self, explicit_request_timeout): router = self._make_router(timeout=330) assert router.timeout == 330 assert router.request_timeout == 300 @@ -12435,22 +12250,16 @@ class TestRouterRequestTimeoutPropagation: litellm.request_timeout = original_value litellm.request_timeout_explicitly_set = original_flag - def test_non_stream_prefers_request_timeout_over_router_timeout( - self, explicit_request_timeout - ): + def test_non_stream_prefers_request_timeout_over_router_timeout(self, explicit_request_timeout): router = self._make_router(timeout=330) assert router._get_non_stream_timeout(kwargs={}, data={}) == 300 - def test_stream_prefers_request_timeout_over_router_timeout( - self, explicit_request_timeout - ): + def test_stream_prefers_request_timeout_over_router_timeout(self, explicit_request_timeout): router = self._make_router(timeout=330) # stream=True resolves through _get_stream_timeout; request_timeout must win. assert router._get_timeout(kwargs={"stream": True}, data={}) == 300 - def test_explicit_stream_timeout_still_wins_over_request_timeout( - self, explicit_request_timeout - ): + def test_explicit_stream_timeout_still_wins_over_request_timeout(self, explicit_request_timeout): router = self._make_router(timeout=330, stream_timeout=45) assert router._get_stream_timeout(kwargs={}, data={}) == 45 @@ -12466,22 +12275,13 @@ class TestRouterRequestTimeoutPropagation: litellm.request_timeout = original_value litellm.request_timeout_explicitly_set = original_flag - def test_per_deployment_timeout_overrides_request_timeout( - self, explicit_request_timeout - ): + def test_per_deployment_timeout_overrides_request_timeout(self, explicit_request_timeout): router = self._make_router(timeout=330) assert router._get_non_stream_timeout(kwargs={}, data={"timeout": 120}) == 120 - def test_per_request_timeout_overrides_request_timeout( - self, explicit_request_timeout - ): + def test_per_request_timeout_overrides_request_timeout(self, explicit_request_timeout): router = self._make_router(timeout=330) - assert ( - router._get_non_stream_timeout( - kwargs={"timeout": 60}, data={"timeout": 120} - ) - == 60 - ) + assert router._get_non_stream_timeout(kwargs={"timeout": 60}, data={"timeout": 120}) == 60 def test_passthrough_prefers_request_timeout_over_router_timeout(self, explicit_request_timeout): router = self._make_router(timeout=330) @@ -12800,9 +12600,7 @@ class TestAdvisorSubCallCooldown: ) def _cooled_down_ids(self, router): - active = router.cooldown_cache.get_active_cooldowns( - model_ids=["dep-1"], parent_otel_span=None - ) + active = router.cooldown_cache.get_active_cooldowns(model_ids=["dep-1"], parent_otel_span=None) return [entry[0] for entry in active] @pytest.mark.asyncio @@ -12811,12 +12609,7 @@ class TestAdvisorSubCallCooldown: router = self._router() now = datetime.now() - assert ( - router.deployment_callback_on_failure( - self._kwargs(self._auth_error()), None, now, now - ) - is True - ) + assert router.deployment_callback_on_failure(self._kwargs(self._auth_error()), None, now, now) is True assert "dep-1" in self._cooled_down_ids(router) def test_advisor_orchestration_failure_does_not_cool_down_deployment(self): @@ -12831,12 +12624,7 @@ class TestAdvisorSubCallCooldown: mark_advisor_orchestration_failure(exception) now = datetime.now() - assert ( - router.deployment_callback_on_failure( - self._kwargs(exception), None, now, now - ) - is False - ) + assert router.deployment_callback_on_failure(self._kwargs(exception), None, now, now) is False assert "dep-1" not in self._cooled_down_ids(router) @@ -13125,13 +12913,13 @@ def test_stream_chunks_have_generated_content_detects_text_and_non_text(): audio_chunk = _chunk(audio_delta) assert _stream_chunks_have_generated_content([audio_chunk]) is True - images_delta = Delta(images=[{"image_url": {"url": "https://example.com/img.png"}, "index": 0, "type": "image_url"}]) + images_delta = Delta( + images=[{"image_url": {"url": "https://example.com/img.png"}, "index": 0, "type": "image_url"}] + ) images_chunk = _chunk(images_delta) assert _stream_chunks_have_generated_content([images_chunk]) is True - annotations_delta = Delta( - annotations=[{"type": "url_citation", "url_citation": {"url": "https://example.com"}}] - ) + annotations_delta = Delta(annotations=[{"type": "url_citation", "url_citation": {"url": "https://example.com"}}]) annotations_chunk = _chunk(annotations_delta) assert _stream_chunks_have_generated_content([annotations_chunk]) is True @@ -13175,12 +12963,8 @@ def test_get_configured_token_limits_skips_wildcard_pattern_matching(): ] ) - with patch.object( - router.pattern_router, "route", side_effect=AssertionError("pattern route called") - ): - assert router.get_configured_token_limits( - "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0" - ) == (None, None) + with patch.object(router.pattern_router, "route", side_effect=AssertionError("pattern route called")): + assert router.get_configured_token_limits("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") == (None, None) def test_get_configured_token_limits_treats_malformed_values_as_absent(): @@ -13403,13 +13187,8 @@ def test_get_model_listing_info_skips_wildcard_pattern_matching(): ] ) - with patch.object( - router.pattern_router, "route", side_effect=AssertionError("pattern route called") - ): - assert ( - router.get_model_listing_info("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") - is None - ) + with patch.object(router.pattern_router, "route", side_effect=AssertionError("pattern route called")): + assert router.get_model_listing_info("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") is None def test_get_configured_mode_reads_deployment_model_info(): @@ -13451,13 +13230,8 @@ def test_get_configured_mode_skips_wildcard_pattern_matching(): ] ) - with patch.object( - router.pattern_router, "route", side_effect=AssertionError("pattern route called") - ): - assert ( - router.get_configured_mode("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") - is None - ) + with patch.object(router.pattern_router, "route", side_effect=AssertionError("pattern route called")): + assert router.get_configured_mode("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") is None def test_get_configured_mode_treats_malformed_values_as_absent(): @@ -13516,13 +13290,8 @@ def test_get_configured_display_name_skips_wildcard_pattern_matching(): ] ) - with patch.object( - router.pattern_router, "route", side_effect=AssertionError("pattern route called") - ): - assert ( - router.get_configured_display_name("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") - is None - ) + with patch.object(router.pattern_router, "route", side_effect=AssertionError("pattern route called")): + assert router.get_configured_display_name("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") is None @pytest.mark.parametrize( @@ -13551,7 +13320,10 @@ def test_get_configured_service_tiers_returns_one_value_per_deployment_in_model_ "litellm_params": {"model": "openai/gpt-6-astra"}, "model_info": {"service_tiers": ["ultrafast"]}, }, - {"model_name": "gpt-6-astra", "litellm_params": {"model": "openai/gpt-6-astra", "api_base": "https://a.example"}}, + { + "model_name": "gpt-6-astra", + "litellm_params": {"model": "openai/gpt-6-astra", "api_base": "https://a.example"}, + }, { "model_name": "gpt-6-astra", "litellm_params": {"model": "openai/gpt-6-astra", "api_base": "https://b.example"}, @@ -13669,7 +13441,11 @@ def test_get_search_tool_support_reports_each_routable_deployment_of_the_group() @pytest.mark.parametrize("provider", ["chatgpt", "github_copilot"]) @pytest.mark.parametrize("explicit_provider", [False, True]) async def test_search_tool_catalog_lookup_needs_no_credentials( - provider: str, explicit_provider: bool, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter + provider: str, + explicit_provider: bool, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, ) -> None: monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path / "chatgpt")) monkeypatch.setenv("GITHUB_COPILOT_TOKEN_DIR", str(tmp_path / "copilot")) @@ -14269,7 +14045,12 @@ def test_credentials_for_a_mixed_group_name_keep_resolving_to_its_first_deployme _AZURE_CHAT_SKU_REJECTS_BATCH: Final = httpx.Response( 400, - json={"error": {"code": "invalid_deployment_type", "message": "The current deployment is not supported for batch jobs"}}, + json={ + "error": { + "code": "invalid_deployment_type", + "message": "The current deployment is not supported for batch jobs", + } + }, ) @@ -14290,7 +14071,9 @@ async def test_acreate_batch_on_a_mixed_group_reaches_chat_as_on_main_while_its_ _cool_down(router, *(d["model_info"]["id"] for d in router.model_list if d["model_info"]["mode"] == "batch")) with respx.mock(assert_all_called=False) as respx_mock: - chat_route: Final = respx_mock.post(host="chat-sku.openai.azure.com").mock(return_value=_AZURE_CHAT_SKU_REJECTS_BATCH) + chat_route: Final = respx_mock.post(host="chat-sku.openai.azure.com").mock( + return_value=_AZURE_CHAT_SKU_REJECTS_BATCH + ) with pytest.raises(openai.BadRequestError): await _create_batches(router, "gpt-mixed", attempts=1) @@ -14316,7 +14099,9 @@ async def test_acreate_batch_with_tags_that_match_only_the_chat_deployment_reach ) with respx.mock(assert_all_called=False) as respx_mock: - chat_route: Final = respx_mock.post(host="chat-sku.openai.azure.com").mock(return_value=_AZURE_CHAT_SKU_REJECTS_BATCH) + chat_route: Final = respx_mock.post(host="chat-sku.openai.azure.com").mock( + return_value=_AZURE_CHAT_SKU_REJECTS_BATCH + ) batch_route: Final = respx_mock.post(host="batch-sku.openai.azure.com").mock( return_value=httpx.Response(200, json=_AZURE_BATCH_CREATED) ) @@ -14406,13 +14191,16 @@ async def test_acreate_batch_request_bedrock_tags_override_deployment_tags(): mock_client = MagicMock() mock_client.post = AsyncMock(side_effect=lambda *args, **kwargs: fake_response()) - with patch.object( - CommonBatchFilesUtils, - "sign_aws_request", - return_value=({"Authorization": "signed"}, b"{}"), - ) as mock_sign, patch( - "litellm.llms.custom_httpx.llm_http_handler.get_async_httpx_client", - return_value=mock_client, + with ( + patch.object( + CommonBatchFilesUtils, + "sign_aws_request", + return_value=({"Authorization": "signed"}, b"{}"), + ) as mock_sign, + patch( + "litellm.llms.custom_httpx.llm_http_handler.get_async_httpx_client", + return_value=mock_client, + ), ): await router.acreate_batch( model="bedrock-batch-model", @@ -14441,13 +14229,13 @@ async def test_avector_store_search_injects_router(): """ from litellm.types.vector_stores import VectorStoreSearchResponse - expected_response = VectorStoreSearchResponse( - object="vector_store.search_results.page", search_query="q", data=[] - ) + expected_response = VectorStoreSearchResponse(object="vector_store.search_results.page", search_query="q", data=[]) mock_asearch = AsyncMock(return_value=expected_response) # Router.__init__ binds asearch via a local import, so patch the module # attribute before constructing the Router. - with patch("litellm.vector_stores.main.asearch", new=mock_asearch): # test-quality-ok: the SDK call is the only place the injected router kwarg is observable + with patch( + "litellm.vector_stores.main.asearch", new=mock_asearch + ): # test-quality-ok: the SDK call is the only place the injected router kwarg is observable router = litellm.Router( model_list=[ { @@ -14481,7 +14269,9 @@ async def test_avector_store_create_does_not_inject_router(): } ] ) - with patch("litellm.vector_stores.main.acreate", new=mock_acreate): # test-quality-ok: the SDK call is the only place a leaked router kwarg would surface + with patch( + "litellm.vector_stores.main.acreate", new=mock_acreate + ): # test-quality-ok: the SDK call is the only place a leaked router kwarg would surface create_response = await router.avector_store_create(model=None, custom_llm_provider="openai") assert create_response is expected_response @@ -14497,13 +14287,13 @@ def test_vector_store_search_injects_router(): """ from litellm.types.vector_stores import VectorStoreSearchResponse - expected_response = VectorStoreSearchResponse( - object="vector_store.search_results.page", search_query="q", data=[] - ) + expected_response = VectorStoreSearchResponse(object="vector_store.search_results.page", search_query="q", data=[]) mock_search = MagicMock(return_value=expected_response) # Router.__init__ binds search via a local import, so patch the module # attribute before constructing the Router. - with patch("litellm.vector_stores.main.search", new=mock_search): # test-quality-ok: the SDK call is the only place the injected router kwarg is observable + with patch( + "litellm.vector_stores.main.search", new=mock_search + ): # test-quality-ok: the SDK call is the only place the injected router kwarg is observable router = litellm.Router( model_list=[ { @@ -14512,9 +14302,7 @@ def test_vector_store_search_injects_router(): } ] ) - search_response = router.vector_store_search( - vector_store_id="v", query="q", custom_llm_provider="s3_vectors" - ) + search_response = router.vector_store_search(vector_store_id="v", query="q", custom_llm_provider="s3_vectors") assert search_response is expected_response mock_search.assert_called_once() @@ -14528,7 +14316,9 @@ def test_vector_store_create_does_not_inject_router(): mock_create = MagicMock(return_value=expected_response) # Router.__init__ binds create via a local import, so patch the module # attribute before constructing the Router. - with patch("litellm.vector_stores.main.create", new=mock_create): # test-quality-ok: the SDK call is the only place a leaked router kwarg would surface + with patch( + "litellm.vector_stores.main.create", new=mock_create + ): # test-quality-ok: the SDK call is the only place a leaked router kwarg would surface router = litellm.Router( model_list=[ { @@ -14563,9 +14353,7 @@ class TestPreRoutingStrategyRegistryLifecycle: def _complexity_router_params(default_model: str, tags=None) -> dict: return { "model": "auto_router/complexity_router", - "complexity_router_config": { - "tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"} - }, + "complexity_router_config": {"tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"}}, "complexity_router_default_model": default_model, **({"tags": tags} if tags else {}), } @@ -14870,9 +14658,7 @@ class TestPreRoutingStrategyRegistryLifecycle: deployment=Deployment( model_name="hybrid-router", litellm_params=LiteLLM_Params( - **self._hybrid_router_params( - {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"} - ) + **self._hybrid_router_params({"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"}) ), model_info=ModelInfo(id="router-1", db_model=True), ) @@ -14981,9 +14767,7 @@ class TestPreRoutingStrategyRegistryLifecycle: ({"model": "openai/gpt-4o"}, False), ] for params, expected in cases: - actual = router._deployment_participates_in_adaptive_routing( - litellm_params=LiteLLM_Params(**params) - ) + actual = router._deployment_participates_in_adaptive_routing(litellm_params=LiteLLM_Params(**params)) assert actual is expected, params["model"] @@ -15170,22 +14954,16 @@ class TestUpsertDeploymentRollback: router.delete_deployment(id="prod-1") assert router.has_model_id("prod-1") is False - router._restore_deployment_after_failed_upsert( - previous_deployment=previous, model_id="prod-1" - ) + router._restore_deployment_after_failed_upsert(previous_deployment=previous, model_id="prod-1") restored = router.get_deployment(model_id="prod-1") assert restored is not None assert restored.litellm_params.model == "gpt-4o" - router._restore_deployment_after_failed_upsert( - previous_deployment=previous, model_id="prod-1" - ) + router._restore_deployment_after_failed_upsert(previous_deployment=previous, model_id="prod-1") assert len(router.model_list) == 1 - router._restore_deployment_after_failed_upsert( - previous_deployment=None, model_id="prod-1" - ) + router._restore_deployment_after_failed_upsert(previous_deployment=None, model_id="prod-1") assert len(router.model_list) == 1 @@ -15457,26 +15235,34 @@ class TestClaudeCodeSubagentSessionRouterBinding: @pytest.mark.asyncio async def test_redis_cleanup_failure_does_not_reject_a_subagent_request(self): - from litellm.caching.caching import RedisCache + import fakeredis + from redis.exceptions import ConnectionError as RedisConnectionError + + from tests._support.fake_redis import fake_redis_cache router = self._router() del router.complexity_routers["smart-router"] - redis_cache = MagicMock(spec=RedisCache) - redis_cache.async_get_cache = AsyncMock(return_value="smart-router") - redis_cache.async_delete_cache = AsyncMock(side_effect=ConnectionError("redis unavailable")) + server: Final = fakeredis.FakeServer() + raw: Final = fakeredis.FakeAsyncRedis(server=server) + key: Final = "claude_code_session_router:v1:key-hash-a:session-1234" + await raw.set(key, b'"smart-router"', ex=3600) + redis_cache = fake_redis_cache(server) router.update_redis_cache(cache=redis_cache) - response = await router.async_pre_routing_hook( - model="expensive-model", - request_kwargs=self._request_kwargs(agent_id="agent-1234"), - ) + with patch.object(fakeredis.FakeAsyncRedis, "delete", side_effect=RedisConnectionError("unavailable")): + response = await router.async_pre_routing_hook( + model="expensive-model", + request_kwargs=self._request_kwargs(agent_id="agent-1234"), + ) assert response is None - redis_cache.async_delete_cache.assert_awaited_once() + assert await raw.get(key) == b'"smart-router"' @pytest.mark.asyncio async def test_redis_read_failure_does_not_reject_a_subagent_request(self): - from litellm.caching.caching import RedisCache + import fakeredis + + from tests._support.fake_redis import fake_redis_cache router = self._router() request_kwargs = self._request_kwargs(agent_id="agent-1234") @@ -15486,9 +15272,10 @@ class TestClaudeCodeSubagentSessionRouterBinding: cache_key, "smart-router", ) - redis_cache = MagicMock(spec=RedisCache) - redis_cache.async_get_cache = AsyncMock(side_effect=Exception("Redis circuit breaker is open")) + server: Final = fakeredis.FakeServer() + redis_cache = fake_redis_cache(server) router.update_redis_cache(cache=redis_cache) + server.connected = False response = await router.async_pre_routing_hook( model="expensive-model", @@ -15497,45 +15284,46 @@ class TestClaudeCodeSubagentSessionRouterBinding: assert response is None assert "model_group" not in request_kwargs["metadata"] - redis_cache.async_get_cache.assert_awaited_once() @pytest.mark.asyncio async def test_redis_write_failures_do_not_reject_main_or_subagent_requests(self): - from litellm.caching.caching import RedisCache + import fakeredis + from redis.exceptions import ConnectionError as RedisConnectionError + + from tests._support.fake_redis import fake_redis_cache router = self._router() - redis_cache = MagicMock(spec=RedisCache) - redis_cache.async_get_cache = AsyncMock(return_value="smart-router") - redis_cache.async_set_cache = AsyncMock(side_effect=Exception("redis unavailable")) + server: Final = fakeredis.FakeServer() + raw: Final = fakeredis.FakeAsyncRedis(server=server) + key: Final = "claude_code_session_router:v1:key-hash-a:session-1234" + await raw.set(key, b'"smart-router"', ex=3600) + redis_cache = fake_redis_cache(server) router.update_redis_cache(cache=redis_cache) - main_response = await router.async_pre_routing_hook( - model="smart-router", - request_kwargs=self._request_kwargs(), - ) - subagent_response = await router.async_pre_routing_hook( - model="expensive-model", - request_kwargs=self._request_kwargs(agent_id="agent-1234"), - ) + with patch.object(fakeredis.FakeAsyncRedis, "set", side_effect=RedisConnectionError("unavailable")): + main_response = await router.async_pre_routing_hook( + model="smart-router", + request_kwargs=self._request_kwargs(), + ) + subagent_response = await router.async_pre_routing_hook( + model="expensive-model", + request_kwargs=self._request_kwargs(agent_id="agent-1234"), + ) assert main_response is not None assert main_response.model == "cheap-model" assert subagent_response is not None assert subagent_response.model == "cheap-model" - assert redis_cache.async_set_cache.await_count == 2 + assert await raw.get(key) == b'"smart-router"' @pytest.mark.asyncio async def test_subagents_follow_the_main_threads_latest_router_across_workers(self): - from types import SimpleNamespace + import fakeredis - from litellm.caching.caching import RedisCache + from tests._support.fake_redis import fake_redis_cache - shared_binding = SimpleNamespace(value=None) - shared_redis = MagicMock(spec=RedisCache) - shared_redis.async_get_cache = AsyncMock(side_effect=lambda key, **_: shared_binding.value) - shared_redis.async_set_cache = AsyncMock( - side_effect=lambda key, value, **_: setattr(shared_binding, "value", value) - ) + server: Final = fakeredis.FakeServer() + shared_redis: Final = fake_redis_cache(server) main_worker, subagent_worker = self._router(), self._router() main_worker.update_redis_cache(cache=shared_redis) subagent_worker.update_redis_cache(cache=shared_redis) @@ -15555,27 +15343,33 @@ class TestClaudeCodeSubagentSessionRouterBinding: assert first.model == "cheap-model" assert second is not None assert second.model == "expensive-model" - assert shared_binding.value == "premium-router" + assert ( + await fakeredis.FakeAsyncRedis(server=server).get("claude_code_session_router:v1:key-hash-a:session-1234") + == b'"premium-router"' + ) @pytest.mark.asyncio - async def test_no_pre_routing_strategies_means_no_session_cache_traffic(self): - from litellm.caching.caching import RedisCache + @patch("time.time", return_value=1800000000.0) + async def test_no_pre_routing_strategies_means_no_session_cache_traffic(self, _clock: object): + import fakeredis + + from tests._support.fake_redis import fake_redis_cache router = self._router() router.complexity_routers.clear() - redis_cache = MagicMock(spec=RedisCache) - redis_cache.async_get_cache = AsyncMock(return_value=None) - redis_cache.async_set_cache = AsyncMock() - redis_cache.async_delete_cache = AsyncMock() + server: Final = fakeredis.FakeServer() + raw: Final = fakeredis.FakeAsyncRedis(server=server) + key: Final = "claude_code_session_router:v1:key-hash-a:session-1234" + await raw.set(key, b'"smart-router"', ex=600) + redis_cache = fake_redis_cache(server) router.update_redis_cache(cache=redis_cache) for request_kwargs in (self._request_kwargs(), self._request_kwargs(agent_id="agent-1234")): response = await router.async_pre_routing_hook(model="expensive-model", request_kwargs=request_kwargs) assert response is None - redis_cache.async_get_cache.assert_not_awaited() - redis_cache.async_set_cache.assert_not_awaited() - redis_cache.async_delete_cache.assert_not_awaited() + assert await raw.get(key) == b'"smart-router"' + assert await raw.pttl(key) == 600000 @pytest.mark.asyncio async def test_session_bindings_do_not_evict_router_rate_limit_state(self): @@ -15945,18 +15739,14 @@ class TestAutoRouterSharedModelNameConnectionParams: return httpx.Response( status_code=200, json={ - "candidates": [ - {"content": {"parts": [{"text": "Paris"}], "role": "model"}, "finishReason": "STOP"} - ], + "candidates": [{"content": {"parts": [{"text": "Paris"}], "role": "model"}, "finishReason": "STOP"}], "usageMetadata": {"promptTokenCount": 5, "candidatesTokenCount": 1, "totalTokenCount": 6}, "modelVersion": "gemini-3.6-flash", }, request=httpx.Request("POST", "https://generativelanguage.googleapis.com"), ) - @pytest.mark.parametrize( - "plain_entry_first", [True, False], ids=["plain_entry_first", "marker_entry_first"] - ) + @pytest.mark.parametrize("plain_entry_first", [True, False], ids=["plain_entry_first", "marker_entry_first"]) async def test_routed_tier_call_goes_out_on_its_own_endpoint_and_credentials(self, plain_entry_first): """The outbound provider request for the routed tier hits the tier's own Gemini host with the tier's own key, never the plain sibling's api_base or api_key.""" @@ -16088,9 +15878,7 @@ async def _drive_cyclic_fallback(router, capture, recorder=None, **request_kwarg litellm.callbacks.append(recorder) try: with pytest.raises(litellm.InternalServerError): - await router.acompletion( - model="group-a", messages=[{"role": "user", "content": "hi"}], **request_kwargs - ) + await router.acompletion(model="group-a", messages=[{"role": "user", "content": "hi"}], **request_kwargs) finally: router_logger.removeHandler(capture) router_logger.setLevel(previous_level) @@ -16364,9 +16152,7 @@ async def test_fallback_failure_detail_from_upstream_is_bounded(): await _drive_cyclic_fallback( _cyclic_fallback_router(), capture, - mock_response=litellm.InternalServerError( - message=huge_message, llm_provider="openai", model="group-a" - ), + mock_response=litellm.InternalServerError(message=huge_message, llm_provider="openai", model="group-a"), ) assert capture.messages, "the fallback failure path did not log at ERROR" @@ -16432,9 +16218,7 @@ def test_ensure_deployment_affinity_callback_is_idempotent(): try: router._ensure_deployment_affinity_callback() router._ensure_deployment_affinity_callback() - affinity_callbacks = [ - cb for cb in router.optional_callbacks or [] if isinstance(cb, DeploymentAffinityCheck) - ] + affinity_callbacks = [cb for cb in router.optional_callbacks or [] if isinstance(cb, DeploymentAffinityCheck)] assert len(affinity_callbacks) == 1 finally: for cb in router.optional_callbacks or []: @@ -16576,9 +16360,7 @@ class TestModelGroupAliasReachesPreRoutingStrategies: router = self._router("auto_routers") metadata: dict = {} - response = await router.acompletion( - model="smart-alias", messages=self._messages(), metadata=metadata - ) + response = await router.acompletion(model="smart-alias", messages=self._messages(), metadata=metadata) assert response.choices[0].message.content == "routed by the tier" assert metadata["model_group"] == "smart-alias" @@ -16820,6 +16602,7 @@ class TestTeamPublicNameReachesPreRoutingStrategies: assert response is not None assert response.model == "gemini-flash" + @pytest.mark.asyncio async def test_strategy_resolution_agrees_with_the_deployment_path_for_every_principal(self): router = self._router( @@ -16878,7 +16661,6 @@ class TestTeamPublicNameReachesPreRoutingStrategies: with pytest.raises(litellm.BadRequestError, match="multiple teams"): two_teams._team_deployments_across_teams(self.PUBLIC_NAME) - def test_compression_policy_follows_the_same_resolution_for_every_principal(self): from litellm.proxy.guardrails.auto_router_compression import AutoRouterCompressionPolicy, policy_for_model @@ -17102,7 +16884,6 @@ class TestAutoRouterCompressionDecoupling: @pytest.mark.usefixtures("local_model_cost_map") - @pytest.mark.usefixtures("local_model_cost_map") class TestAzureBaseModelFallbackLogging: """When an azure deployment has no base_model but its model name is a known @@ -17129,17 +16910,14 @@ class TestAzureBaseModelFallbackLogging: def test_map_known_deployment_name_resolves_without_error_log(self): router = self._router_with_azure_deployment("azure/gpt-4o") - with patch( - "litellm.router.verbose_router_logger.error" - ) as mock_error: + with patch("litellm.router.verbose_router_logger.error") as mock_error: model_info = router.get_router_model_info( deployment=None, received_model_name="my-group", id="azure-base-model-test-id" ) - assert not any( - "Could not identify azure model" in str(call) - for call in mock_error.call_args_list - ), f"unexpected error log: {mock_error.call_args_list}" + assert not any("Could not identify azure model" in str(call) for call in mock_error.call_args_list), ( + f"unexpected error log: {mock_error.call_args_list}" + ) # the fallback resolution must actually surface the map values assert model_info["max_input_tokens"] == litellm.model_cost["azure/gpt-4o"]["max_input_tokens"] assert model_info["input_cost_per_token"] == litellm.model_cost["azure/gpt-4o"]["input_cost_per_token"] @@ -17147,17 +16925,14 @@ class TestAzureBaseModelFallbackLogging: def test_unmappable_deployment_name_still_logs_error(self): router = self._router_with_azure_deployment("azure/my-custom-deployment-name") - with patch( - "litellm.router.verbose_router_logger.error" - ) as mock_error: + with patch("litellm.router.verbose_router_logger.error") as mock_error: model_info = router.get_router_model_info( deployment=None, received_model_name="my-group", id="azure-base-model-test-id" ) - assert any( - "Could not identify azure model" in str(call) - for call in mock_error.call_args_list - ), "expected the error log for an unmappable azure deployment name" + assert any("Could not identify azure model" in str(call) for call in mock_error.call_args_list), ( + "expected the error log for an unmappable azure deployment name" + ) # unmappable names resolve to a zeroed stub — unchanged behavior assert model_info.get("max_input_tokens") is None @@ -17184,6 +16959,7 @@ class TestAzureBaseModelFallbackLogging: ) assert model_info["max_input_tokens"] == litellm.model_cost["azure/gpt-4o-mini"]["max_input_tokens"] + def test_model_group_info_intersects_supported_reasoning_efforts(): router = litellm.Router( model_list=[ @@ -17293,11 +17069,15 @@ def test_model_group_info_reasoning_efforts_are_unknown_when_any_deployment_is_o def test_model_group_info_fast_mode_uses_exact_provider_catalog( local_model_cost_map: None, model: str, provider: str | None, expected: bool, operator_flag: bool ) -> None: - router: Final = Router(model_list=[{ - "model_name": "fast-group", - "litellm_params": {"model": model, "custom_llm_provider": provider, "api_key": "fake-key"}, - "model_info": {"supports_fast_mode": operator_flag}, - }]) + router: Final = Router( + model_list=[ + { + "model_name": "fast-group", + "litellm_params": {"model": model, "custom_llm_provider": provider, "api_key": "fake-key"}, + "model_info": {"supports_fast_mode": operator_flag}, + } + ] + ) result: Final = router.get_model_group_info("fast-group") @@ -17309,16 +17089,21 @@ def test_model_group_info_fast_mode_uses_exact_provider_catalog( def test_model_group_info_fast_mode_fails_closed_without_explicit_boolean( local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch, flag: object ) -> None: - entry: Final = {key: value for key, value in litellm.model_cost["claude-opus-5"].items() - if key != "supports_fast_mode"} + entry: Final = { + key: value for key, value in litellm.model_cost["claude-opus-5"].items() if key != "supports_fast_mode" + } if flag is not None: entry["supports_fast_mode"] = flag monkeypatch.setitem(litellm.model_cost, "claude-opus-5", entry) - router: Final = Router(model_list=[{ - "model_name": "fast-group", - "litellm_params": {"model": "anthropic/claude-opus-5", "api_key": "fake-key"}, - "model_info": {"supports_fast_mode": True}, - }]) + router: Final = Router( + model_list=[ + { + "model_name": "fast-group", + "litellm_params": {"model": "anthropic/claude-opus-5", "api_key": "fake-key"}, + "model_info": {"supports_fast_mode": True}, + } + ] + ) result: Final = router.get_model_group_info("fast-group") @@ -17326,24 +17111,30 @@ def test_model_group_info_fast_mode_fails_closed_without_explicit_boolean( assert result.supports_fast_mode is False -@pytest.mark.parametrize("other_model,expected", [ - ("anthropic/claude-opus-4-8", True), - ("anthropic/claude-opus-4-7", False), - ("anthropic/off-map-opus", False), - ("vertex_ai/claude-opus-5", False), - ("bedrock/claude-opus-5", False), -]) +@pytest.mark.parametrize( + "other_model,expected", + [ + ("anthropic/claude-opus-4-8", True), + ("anthropic/claude-opus-4-7", False), + ("anthropic/off-map-opus", False), + ("vertex_ai/claude-opus-5", False), + ("bedrock/claude-opus-5", False), + ], +) @pytest.mark.parametrize("reverse", [True, False]) def test_model_group_info_fast_mode_requires_every_deployment( local_model_cost_map: None, other_model: str, expected: bool, reverse: bool ) -> None: - models: Final = (other_model, "anthropic/claude-opus-5") if reverse else ( - "anthropic/claude-opus-5", other_model + models: Final = (other_model, "anthropic/claude-opus-5") if reverse else ("anthropic/claude-opus-5", other_model) + router: Final = Router( + model_list=[ + { + "model_name": "fast-group", + "litellm_params": {"model": model, "api_key": "fake-key"}, + } + for model in models + ] ) - router: Final = Router(model_list=[{ - "model_name": "fast-group", - "litellm_params": {"model": model, "api_key": "fake-key"}, - } for model in models]) result: Final = router.get_model_group_info("fast-group") @@ -17583,6 +17374,7 @@ class TestAddDeploymentApiBaseProviderResolution: assert deployment is not None assert deployment.litellm_params.custom_llm_provider == "openai" + # ===================================================================== # anthropic_messages mid-stream-fallback helpers, added for #24004 # (mid-stream fallback not supported for anthropic_messages route type). @@ -17701,10 +17493,7 @@ class _AnthropicMessagesFallbackByteStream: def _anthropic_messages_overloaded_error_chunk() -> bytes: - return ( - b"event: error\n" - b'data: {"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}\n\n' - ) + return b'event: error\ndata: {"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}\n\n' def _anthropic_messages_invalid_request_error_chunk() -> bytes: @@ -17749,9 +17538,7 @@ async def test_anthropic_messages_streaming_iterator_passthrough(): [_anthropic_messages_content_chunk("hi"), _anthropic_messages_content_chunk(" there")] ) - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, initial_kwargs={"model": "primary"} - ) + wrapped = await router._aanthropic_messages_streaming_iterator(response=source, initial_kwargs={"model": "primary"}) collected = [chunk async for chunk in wrapped] assert collected == [_anthropic_messages_content_chunk("hi"), _anthropic_messages_content_chunk(" there")] @@ -17770,12 +17557,14 @@ async def test_anthropic_messages_streaming_iterator_flushes_buffered_lifecycle_ [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("hi"), message_stop] ) - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, initial_kwargs={"model": "primary"} - ) + wrapped = await router._aanthropic_messages_streaming_iterator(response=source, initial_kwargs={"model": "primary"}) collected = [chunk async for chunk in wrapped] - assert collected == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("hi"), message_stop] + assert collected == [ + _anthropic_messages_message_start_chunk(), + _anthropic_messages_content_chunk("hi"), + message_stop, + ] @pytest.mark.asyncio @@ -17787,9 +17576,7 @@ async def test_anthropic_messages_streaming_iterator_flushes_buffered_frames_on_ message_stop = b'event: message_stop\ndata: {"type": "message_stop"}\n\n' source = _AnthropicMessagesFakeByteStream([_anthropic_messages_message_start_chunk(), message_stop]) - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, initial_kwargs={"model": "primary"} - ) + wrapped = await router._aanthropic_messages_streaming_iterator(response=source, initial_kwargs={"model": "primary"}) collected = [chunk async for chunk in wrapped] assert collected == [_anthropic_messages_message_start_chunk(), message_stop] @@ -17840,7 +17627,9 @@ async def test_anthropic_messages_ping_behind_buffered_lifecycle_frame_is_forwar await content_released.wait() yield _anthropic_messages_content_chunk("hi") - wrapped = await router._aanthropic_messages_streaming_iterator(response=source(), initial_kwargs={"model": "primary"}) + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source(), initial_kwargs={"model": "primary"} + ) assert await asyncio.wait_for(wrapped.__anext__(), timeout=1) == _anthropic_messages_ping_chunk() content_released.set() @@ -17885,7 +17674,9 @@ async def test_anthropic_messages_no_fallback_message_start_reaches_client_befor await content_released.wait() yield _anthropic_messages_content_chunk("hi") - wrapped = await router._aanthropic_messages_streaming_iterator(response=source(), initial_kwargs={"model": "primary"}) + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source(), initial_kwargs={"model": "primary"} + ) assert await asyncio.wait_for(wrapped.__anext__(), timeout=1) == _anthropic_messages_message_start_chunk() content_released.set() @@ -17950,7 +17741,9 @@ async def test_anthropic_messages_default_wildcard_fallback_still_buffers_lifecy await content_released.wait() yield _anthropic_messages_content_chunk("hi") - wrapped = await router._aanthropic_messages_streaming_iterator(response=source(), initial_kwargs={"model": "primary"}) + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source(), initial_kwargs={"model": "primary"} + ) pending = asyncio.ensure_future(wrapped.__anext__()) await asyncio.sleep(0.2) @@ -17990,8 +17783,15 @@ def _anthropic_messages_two_order_primary_model_list() -> list: id="wildcard-overridden-by-request-none", ), pytest.param({"fallbacks": [{"*": ["fallback"]}]}, {"model": "primary"}, True, id="wildcard"), - pytest.param({"fallbacks": None}, {"model": "primary", "fallbacks": [{"model": "fallback"}]}, True, id="request-dict-fallback"), - pytest.param({"fallbacks": None}, {"model": "primary", "fallbacks": ["fallback"]}, True, id="request-list-fallback"), + pytest.param( + {"fallbacks": None}, + {"model": "primary", "fallbacks": [{"model": "fallback"}]}, + True, + id="request-dict-fallback", + ), + pytest.param( + {"fallbacks": None}, {"model": "primary", "fallbacks": ["fallback"]}, True, id="request-list-fallback" + ), pytest.param( {"fallbacks": [{"primary": ["fallback"]}]}, {"model": "primary", "disable_fallbacks": True}, @@ -18004,7 +17804,9 @@ def _anthropic_messages_two_order_primary_model_list() -> list: True, id="content-policy-fallback", ), - pytest.param({"fallbacks": None, "enable_weighted_failover": True}, {"model": "primary"}, True, id="weighted-failover"), + pytest.param( + {"fallbacks": None, "enable_weighted_failover": True}, {"model": "primary"}, True, id="weighted-failover" + ), ], ) def test_stream_can_fall_back_direct_call(router_kwargs, request_kwargs, expected): @@ -18077,7 +17879,9 @@ async def test_anthropic_messages_order_fallback_still_buffers_lifecycle_frames( await content_released.wait() yield _anthropic_messages_content_chunk("hi") - wrapped = await router._aanthropic_messages_streaming_iterator(response=source(), initial_kwargs={"model": "primary"}) + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source(), initial_kwargs={"model": "primary"} + ) pending = asyncio.ensure_future(wrapped.__anext__()) await asyncio.sleep(0.2) @@ -18122,7 +17926,9 @@ async def test_anthropic_messages_leading_ping_keepalive_is_forwarded_live(): yield _anthropic_messages_message_start_chunk() yield _anthropic_messages_content_chunk("hi") - wrapped = await router._aanthropic_messages_streaming_iterator(response=source(), initial_kwargs={"model": "primary"}) + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source(), initial_kwargs={"model": "primary"} + ) assert await asyncio.wait_for(wrapped.__anext__(), timeout=1) == _anthropic_messages_ping_chunk() content_released.set() @@ -18819,7 +18625,9 @@ class _AnthropicMessagesScriptedProvider: async def __call__(self, **kwargs): litellm_metadata = kwargs.get("litellm_metadata") or {} - self.calls.append((kwargs["model"], litellm_metadata.get("attempted_retries"), litellm_metadata.get("max_retries"))) + self.calls.append( + (kwargs["model"], litellm_metadata.get("attempted_retries"), litellm_metadata.get("max_retries")) + ) assert self._streams, "provider called more times than scripted" return self._streams.pop(0)() @@ -18837,7 +18645,9 @@ def _anthropic_messages_transport_drop(original_exception: Exception | None = No def _anthropic_messages_dropped_before_content(): - return _AnthropicMessagesRaisingByteStream([_anthropic_messages_message_start_chunk()], _anthropic_messages_transport_drop()) + return _AnthropicMessagesRaisingByteStream( + [_anthropic_messages_message_start_chunk()], _anthropic_messages_transport_drop() + ) def _anthropic_messages_bridge_error_chunk() -> bytes: @@ -18965,9 +18775,7 @@ async def test_anthropic_messages_retry_opens_without_the_superseded_attempts_pa provider: Final = _AnthropicMessagesScriptedProvider(_anthropic_messages_dropped_before_content, retry_stream) - stream: Final = await _anthropic_messages_stream_through_router( - router, provider, litellm_logging_obj=logging_obj - ) + stream: Final = await _anthropic_messages_stream_through_router(router, provider, litellm_logging_obj=logging_obj) body: Final = tuple([chunk async for chunk in stream]) assert stash_recorder.observed == (None,) @@ -19563,20 +19371,28 @@ async def test_anthropic_messages_retries_keep_the_budget_the_first_drop_committ }, { "model_name": "glm", - "litellm_params": {"model": "anthropic/glm-b", "api_key": "sk-test", "num_retries": sibling_num_retries}, + "litellm_params": { + "model": "anthropic/glm-b", + "api_key": "sk-test", + "num_retries": sibling_num_retries, + }, }, ], num_retries=0, fallbacks=None, ) router.set_custom_routing_strategy(_AnthropicMessagesAlternatingDeployments(router, "glm")) - provider = _AnthropicMessagesScriptedProvider(*[_anthropic_messages_dropped_before_content] * len(expected_counters)) + provider = _AnthropicMessagesScriptedProvider( + *[_anthropic_messages_dropped_before_content] * len(expected_counters) + ) stream = await _anthropic_messages_stream_through_router(router, provider) with pytest.raises(litellm.APIConnectionError): [chunk async for chunk in stream] - assert [model for model, _, _ in provider.calls] == (["anthropic/glm-a", "anthropic/glm-b"] * 2)[: len(expected_counters)] + assert [model for model, _, _ in provider.calls] == (["anthropic/glm-a", "anthropic/glm-b"] * 2)[ + : len(expected_counters) + ] assert [(attempted, budget) for _, attempted, budget in provider.calls] == expected_counters @@ -20911,9 +20727,7 @@ def _router_with_failing_primary(fallback_access_check) -> Router: class _DenyFallbackBudgetCheck: - async def __call__( - self, *, model: str, request_kwargs: Mapping[str, object], llm_router: Router - ) -> bool: + async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], llm_router: Router) -> bool: return False @@ -21206,9 +21020,9 @@ class TestTierParamsTheTargetAccepts: def test_declared_param_allowlist_ignores_malformed_declarations(self): """A str is iterable, so without the type guard a YAML scalar mistake like allowed_openai_params: reasoning_effort would allowlist single characters.""" - assert litellm.Router._declared_param_allowlist({"allowed_openai_params": ["reasoning_effort", 3]}) == frozenset( - {"reasoning_effort"} - ) + assert litellm.Router._declared_param_allowlist( + {"allowed_openai_params": ["reasoning_effort", 3]} + ) == frozenset({"reasoning_effort"}) assert litellm.Router._declared_param_allowlist({"allowed_openai_params": "reasoning_effort"}) == frozenset() assert litellm.Router._declared_param_allowlist({}) == frozenset() @@ -21288,7 +21102,11 @@ class TestTierParamsTheTargetAccepts: @pytest.mark.parametrize( "deployment", - [{"model_name": "x"}, {"model_name": "x", "litellm_params": {}}, {"model_name": "x", "litellm_params": {"model": "not-a-real-provider/nope"}}], + [ + {"model_name": "x"}, + {"model_name": "x", "litellm_params": {}}, + {"model_name": "x", "litellm_params": {"model": "not-a-real-provider/nope"}}, + ], ) def test_deployment_accepts_param_fails_open(self, deployment): """An unresolvable deployment must not be the reason a param is dropped.""" @@ -21593,9 +21411,7 @@ class TestPreRoutingTierDrivesFallbacks: async def test_the_selected_tier_fallback_chain_runs(self): router = self._router([{"tier1": ["backup-a"]}]) - response = await router.acompletion( - model="smart-router", messages=[{"role": "user", "content": "hi"}] - ) + response = await router.acompletion(model="smart-router", messages=[{"role": "user", "content": "hi"}]) assert response.choices[0].message.content == "from backup-a" @@ -21628,9 +21444,7 @@ class TestPreRoutingTierDrivesFallbacks: async def test_a_request_without_a_pre_routing_hook_still_uses_its_own_group(self): router = self._router([{"tier1": ["backup-a"]}]) - response = await router.acompletion( - model="tier1", messages=[{"role": "user", "content": "hi"}] - ) + response = await router.acompletion(model="tier1", messages=[{"role": "user", "content": "hi"}]) assert response.choices[0].message.content == "from backup-a" @@ -22986,7 +22800,9 @@ async def test_router_deployment_slot_rejects_while_held_and_frees_slot_on_exit( async with router._deployment_slot(deployment=deployment.model_dump(), kwargs=kwargs, parent_otel_span=None): with pytest.raises(litellm.RateLimitError) as overflow: - async with router._deployment_slot(deployment=deployment.model_dump(), kwargs=kwargs, parent_otel_span=None): + async with router._deployment_slot( + deployment=deployment.model_dump(), kwargs=kwargs, parent_otel_span=None + ): pass assert overflow.value.status_code == 429 assert "slot-deployment" in overflow.value.message @@ -23130,33 +22946,33 @@ def test_cached_model_info_lookups_match_uncached_and_reset_on_model_list_change assert router.cached_model_group_info("grp").max_output_tokens == 200 -class _OpenBreakerRedis: - def __init__(self) -> None: - from litellm.caching.redis_cache import RedisCircuitBreaker - - self._circuit_breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60) - for _ in range(3): - self._circuit_breaker.record_failure() - - @_redis_circuit_breaker_guard - async def async_get_cache(self, key, **kwargs): - raise AssertionError("never reached") - - @pytest.mark.asyncio async def test_an_open_circuit_breaker_skips_the_session_binding_without_a_warning(caplog): + import fakeredis + + from litellm.redis.transport.contract import RedisFailure + from tests._support.fake_redis import fake_redis_cache + router = litellm.Router( model_list=[{"model_name": "haiku", "litellm_params": {"model": "anthropic/claude-haiku-4-5", "api_key": "k"}}] ) - router._claude_code_session_router_cache = DualCache(redis_cache=_OpenBreakerRedis()) # pyright: ignore[reportArgumentType] # duck-typed Redis double + server: Final = fakeredis.FakeServer() + cache: Final = fake_redis_cache(server) + router.update_redis_cache(cache=cache) + server.connected = False + for _ in range(6): + await cache.transport().get("probe") + result: Final = await cache.transport().get("probe") + assert isinstance(result, RedisFailure) and result.kind == "breaker_open" caplog.clear() with caplog.at_level(logging.DEBUG, logger="LiteLLM Router"): - binding = await router._get_claude_code_session_router_binding("quiet-session") + binding = await router._get_claude_code_session_router_binding( + "claude_code_session_router:v1:key-hash:quiet-session" + ) assert binding is None assert [record.getMessage() for record in caplog.records if record.levelno >= logging.WARNING] == [] - assert any("circuit breaker is open" in record.getMessage() for record in caplog.records) @pytest.mark.asyncio @@ -23221,18 +23037,25 @@ class TestMemberAutoRouterInference: self.cache = UserApiKeyCache() self.team = LiteLLM_TeamTable( - team_id="router-team", models=["member-router", "permitted-model"], + team_id="router-team", + models=["member-router", "permitted-model"], members_with_roles=[Member(user_id="router-member", role="user")], ) self.actor = UserAPIKeyAuth( - user_id="router-member", team_id="router-team", user_role=LitellmUserRoles.INTERNAL_USER, - models=["member-router", "permitted-model"], api_key="test-key-hash", config={"timeout": 60}, + user_id="router-member", + team_id="router-team", + user_role=LitellmUserRoles.INTERNAL_USER, + models=["member-router", "permitted-model"], + api_key="test-key-hash", + config={"timeout": 60}, + ) + self.database = SimpleNamespace( + db=SimpleNamespace( + litellm_teamtable=SimpleNamespace(find_unique=AsyncMock(return_value=self.team)), + litellm_teammembership=SimpleNamespace(find_unique=AsyncMock(return_value=None)), + litellm_accessgrouptable=SimpleNamespace(find_unique=AsyncMock()), + ) ) - self.database = SimpleNamespace(db=SimpleNamespace( - litellm_teamtable=SimpleNamespace(find_unique=AsyncMock(return_value=self.team)), - litellm_teammembership=SimpleNamespace(find_unique=AsyncMock(return_value=None)), - litellm_accessgrouptable=SimpleNamespace(find_unique=AsyncMock()), - )) monkeypatch.setattr(proxy_server, "user_api_key_cache", self.cache) monkeypatch.setattr(proxy_server, "prisma_client", self.database) @@ -23242,44 +23065,71 @@ class TestMemberAutoRouterInference: return { "model_name": "model_name_router-team_member-router", "litellm_params": { - "model": "auto_router/complexity_router", "complexity_router_default_model": target, + "model": "auto_router/complexity_router", + "complexity_router_default_model": target, "complexity_router_config": { - "tiers": dict.fromkeys(("SIMPLE", "MEDIUM", "COMPLEX", "REASONING"), target), "adaptive": False, + "tiers": dict.fromkeys(("SIMPLE", "MEDIUM", "COMPLEX", "REASONING"), target), + "adaptive": False, **({"classifier_type": "llm", "classifier_llm_config": {"model": target}} if classifier else {}), }, - "tags": ["member" if member else "admin"], "timeout": 13.0 if member else 29.0, + "tags": ["member" if member else "admin"], + "timeout": 13.0 if member else 29.0, }, "model_info": { - "team_id": "router-team", "team_public_model_name": "member-router", "member_auto_router": member, + "team_id": "router-team", + "team_public_model_name": "member-router", + "member_auto_router": member, }, } @classmethod def _router(cls, *markers: dict[str, object]) -> Router: - return Router(model_list=[ - *(markers or (cls._marker(),)), - {"model_name": "permitted-model", "litellm_params": { - "model": "openai/gpt-4o-mini", "api_key": "test-key", "api_base": "https://api.openai.com/v1", - }}, - {"model_name": "restricted-model", "litellm_params": {"model": "openai/gpt-4o", "api_key": "test-key"}}, - ]) + return Router( + model_list=[ + *(markers or (cls._marker(),)), + { + "model_name": "permitted-model", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "test-key", + "api_base": "https://api.openai.com/v1", + }, + }, + {"model_name": "restricted-model", "litellm_params": {"model": "openai/gpt-4o", "api_key": "test-key"}}, + ] + ) def _request( - self, *, actor: UserAPIKeyAuth | None = None, metadata_name: str = "metadata", tag: str = "member", + self, + *, + actor: UserAPIKeyAuth | None = None, + metadata_name: str = "metadata", + tag: str = "member", ) -> dict[str, object]: from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup return LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata( - data={metadata_name: {"tags": [tag]}, **({"metadata": {"user_api_key_auth": {"user_role": "proxy_admin"}}} - if metadata_name == "litellm_metadata" else {})}, - user_api_key_dict=actor or self.actor, _metadata_variable_name=metadata_name, + data={ + metadata_name: {"tags": [tag]}, + **( + {"metadata": {"user_api_key_auth": {"user_role": "proxy_admin"}}} + if metadata_name == "litellm_metadata" + else {} + ), + }, + user_api_key_dict=actor or self.actor, + _metadata_variable_name=metadata_name, ) async def _route( - self, router: Router, request: dict[str, object] | None = None, model: str = "member-router", + self, + router: Router, + request: dict[str, object] | None = None, + model: str = "member-router", ) -> PreRoutingHookResponse: response: Final = await router.async_pre_routing_hook( - model=model, request_kwargs=request if request is not None else self._request(), + model=model, + request_kwargs=request if request is not None else self._request(), messages=[{"role": "user", "content": "Hello"}], ) assert response is not None @@ -23288,36 +23138,68 @@ class TestMemberAutoRouterInference: @pytest.mark.asyncio @pytest.mark.parametrize("metadata_name", ("metadata", "litellm_metadata")) async def test_cached_roster_revocation_blocks_classifier_and_session_rebinding( - self, metadata_name: str, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, + self, + metadata_name: str, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, ) -> None: from litellm.proxy.auth.auth_checks import delete_cache_team_object monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") router: Final = self._router(self._marker(classifier=True)) - classify: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").respond(200, json={ - "id": "classifier", "object": "chat.completion", "created": 0, "model": "gpt-4o-mini", - "choices": [{"index": 0, "message": {"content": '{"tier":"SIMPLE"}', "role": "assistant"}, - "finish_reason": "stop"}], - "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, - }) - request: Final = {**self._request(metadata_name=metadata_name), "proxy_server_request": {"headers": { - "x-claude-code-session-id": "member-router-session", "x-app": "cli", - }}} + classify: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").respond( + 200, + json={ + "id": "classifier", + "object": "chat.completion", + "created": 0, + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "message": {"content": '{"tier":"SIMPLE"}', "role": "assistant"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + ) + request: Final = { + **self._request(metadata_name=metadata_name), + "proxy_server_request": { + "headers": { + "x-claude-code-session-id": "member-router-session", + "x-app": "cli", + } + }, + } first: Final = await self._route(router, request) assert first.model == "permitted-model" and first.routing_decision is not None assert first.routing_decision["cause"] == "llm_classifier" assert (await self._route(router, request)).model == "permitted-model" assert self.database.db.litellm_teamtable.find_unique.await_count == 1 assert self.database.db.litellm_teammembership.find_unique.await_count == 1 - self.database.db.litellm_teamtable.find_unique.return_value = self.team.model_copy(update={"members_with_roles": []}) + self.database.db.litellm_teamtable.find_unique.return_value = self.team.model_copy( + update={"members_with_roles": []} + ) await delete_cache_team_object( - team_id=self.team.team_id, team_alias=None, user_api_key_cache=self.cache, proxy_logging_obj=None, + team_id=self.team.team_id, + team_alias=None, + user_api_key_cache=self.cache, + proxy_logging_obj=None, ) with pytest.raises(HTTPException, match="no longer a member"): await self._route(router, request) - rebound: Final = {**request, "proxy_server_request": {"headers": { - "x-claude-code-session-id": "member-router-session", "x-app": "cli", "x-claude-code-agent-id": "subagent", - }}} + rebound: Final = { + **request, + "proxy_server_request": { + "headers": { + "x-claude-code-session-id": "member-router-session", + "x-app": "cli", + "x-claude-code-agent-id": "subagent", + } + }, + } with pytest.raises(HTTPException, match="no longer a member"): await self._route(router, rebound, model="restricted-model") assert classify.call_count == 2 @@ -23327,11 +23209,23 @@ class TestMemberAutoRouterInference: async def test_member_router_fails_closed(self, state: str, monkeypatch: pytest.MonkeyPatch) -> None: from litellm.proxy import proxy_server - request: Final = {"metadata": {"user_api_key_team_id": "router-team", "user_api_key_auth": { - "team_id": "router-team", "user_role": "proxy_admin", - }}} if state == "forged" else self._request(actor=self.actor.model_copy( - update={"user_id": ""} if state == "empty-user" else {}, - )) + request: Final = ( + { + "metadata": { + "user_api_key_team_id": "router-team", + "user_api_key_auth": { + "team_id": "router-team", + "user_role": "proxy_admin", + }, + } + } + if state == "forged" + else self._request( + actor=self.actor.model_copy( + update={"user_id": ""} if state == "empty-user" else {}, + ) + ) + ) self.database.db.litellm_teamtable.find_unique.return_value = ( None if state == "deleted" else self.team.model_copy(update={"blocked": state == "blocked"}) ) @@ -23342,11 +23236,20 @@ class TestMemberAutoRouterInference: assert error.value.status_code == (503 if state == "unavailable" else 403) @pytest.mark.asyncio - @pytest.mark.parametrize("user_id,role", [(None, LitellmUserRoles.INTERNAL_USER), ("admin", LitellmUserRoles.PROXY_ADMIN)]) - async def test_service_key_and_admin_preserve_runtime_access(self, user_id: str | None, role: LitellmUserRoles) -> None: - assert (await self._route(self._router(), self._request( - actor=self.actor.model_copy(update={"user_id": user_id, "user_role": role}), - ))).model == "permitted-model" + @pytest.mark.parametrize( + "user_id,role", [(None, LitellmUserRoles.INTERNAL_USER), ("admin", LitellmUserRoles.PROXY_ADMIN)] + ) + async def test_service_key_and_admin_preserve_runtime_access( + self, user_id: str | None, role: LitellmUserRoles + ) -> None: + assert ( + await self._route( + self._router(), + self._request( + actor=self.actor.model_copy(update={"user_id": user_id, "user_role": role}), + ), + ) + ).model == "permitted-model" @pytest.mark.asyncio @pytest.mark.parametrize("ceiling", ("team", "key", "member", "organization", "project")) @@ -23357,35 +23260,56 @@ class TestMemberAutoRouterInference: from litellm.proxy._types import LiteLLM_ProjectTableCachedObj from litellm.proxy.common_utils.user_api_key_cache import team_membership_reservation_cache_key - self.database.db.litellm_teamtable.find_unique.return_value = self.team.model_copy(update={ - "models": ["member-router"] if ceiling == "team" else self.team.models, - "organization_id": "router-org" if ceiling == "organization" else None, - }) + self.database.db.litellm_teamtable.find_unique.return_value = self.team.model_copy( + update={ + "models": ["member-router"] if ceiling == "team" else self.team.models, + "organization_id": "router-org" if ceiling == "organization" else None, + } + ) if ceiling == "member": await self.cache.async_set_cache( key=team_membership_reservation_cache_key(user_id="router-member", team_id="router-team"), - value=LiteLLM_TeamMembership(user_id="router-member", team_id="router-team", - litellm_budget_table=LiteLLM_BudgetTable(allowed_models=["restricted-model"])), + value=LiteLLM_TeamMembership( + user_id="router-member", + team_id="router-team", + litellm_budget_table=LiteLLM_BudgetTable(allowed_models=["restricted-model"]), + ), model_type=LiteLLM_TeamMembership, ) elif ceiling == "organization": await self.cache.async_set_cache( - key="org_id:router-org", value=LiteLLM_OrganizationTable( - organization_id="router-org", budget_id="org-budget", created_by="admin", updated_by="admin", + key="org_id:router-org", + value=LiteLLM_OrganizationTable( + organization_id="router-org", + budget_id="org-budget", + created_by="admin", + updated_by="admin", models=["restricted-model"], - ), model_type=LiteLLM_OrganizationTable, + ), + model_type=LiteLLM_OrganizationTable, ) elif ceiling == "project": await self.cache.async_set_cache( - key="project_id:router-project", value=LiteLLM_ProjectTableCachedObj( - project_id="router-project", team_id="router-team", models=["restricted-model"], - ), model_type=LiteLLM_ProjectTableCachedObj, + key="project_id:router-project", + value=LiteLLM_ProjectTableCachedObj( + project_id="router-project", + team_id="router-team", + models=["restricted-model"], + ), + model_type=LiteLLM_ProjectTableCachedObj, ) with pytest.raises(ProxyException, match="is not available for this API key"): - await self._route(self._router(), self._request(actor=self.actor.model_copy(update={ - "models": ["member-router"] if ceiling == "key" else self.actor.models, - "project_id": "router-project" if ceiling == "project" else None, - }))) + await self._route( + self._router(), + self._request( + actor=self.actor.model_copy( + update={ + "models": ["member-router"] if ceiling == "key" else self.actor.models, + "project_id": "router-project" if ceiling == "project" else None, + } + ) + ), + ) @pytest.mark.asyncio @pytest.mark.parametrize("group_owner", ("team", "key")) @@ -23393,22 +23317,32 @@ class TestMemberAutoRouterInference: from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast group: Final = LiteLLM_AccessGroupTable( - access_group_id="router-group", access_group_name="Router targets", access_model_names=["permitted-model"], + access_group_id="router-group", + access_group_name="Router targets", + access_model_names=["permitted-model"], ) self.database.db.litellm_accessgrouptable.find_unique.return_value = group - self.database.db.litellm_teamtable.find_unique.return_value = self.team.model_copy(update={ - "models": ["member-router"] if group_owner == "team" else self.team.models, - "access_group_ids": ["router-group"] if group_owner == "team" else [], - }) - request: Final = self._request(actor=self.actor.model_copy(update={ - "models": ["member-router"] if group_owner == "key" else self.actor.models, - "access_group_ids": ["router-group"] if group_owner == "key" else [], - })) + self.database.db.litellm_teamtable.find_unique.return_value = self.team.model_copy( + update={ + "models": ["member-router"] if group_owner == "team" else self.team.models, + "access_group_ids": ["router-group"] if group_owner == "team" else [], + } + ) + request: Final = self._request( + actor=self.actor.model_copy( + update={ + "models": ["member-router"] if group_owner == "key" else self.actor.models, + "access_group_ids": ["router-group"] if group_owner == "key" else [], + } + ) + ) router: Final = self._router() assert (await self._route(router, request)).model == "permitted-model" assert (await self._route(router, request)).model == "permitted-model" assert self.database.db.litellm_accessgrouptable.find_unique.await_count == 1 - self.database.db.litellm_accessgrouptable.find_unique.return_value = group.model_copy(update={"access_model_names": []}) + self.database.db.litellm_accessgrouptable.find_unique.return_value = group.model_copy( + update={"access_model_names": []} + ) await evict_and_broadcast(cache_keys=("access_group_id:router-group",), user_api_key_cache=self.cache) with pytest.raises(ProxyException, match="is not available for this API key"): await self._route(router, request) @@ -23419,13 +23353,16 @@ class TestMemberAutoRouterInference: router: Final = self._router(self._marker(member=False), self._marker()) request: Final = self._request() selected: Final = router._selected_strategy_marker_deployment( - model="model_name_router-team_member-router", strategy_tags=("member",), request_kwargs=request, + model="model_name_router-team_member-router", + strategy_tags=("member",), + request_kwargs=request, ) assert selected is not None and selected["model_info"]["member_auto_router"] is True assert (await self._route(router, request)).model == "permitted-model" assert request["timeout"] == 13.0 await self.cache.async_set_cache( - key="team_id:router-team", model_type=LiteLLM_TeamTable, + key="team_id:router-team", + model_type=LiteLLM_TeamTable, value=self.team.model_copy(update={"models": ["member-router"]}), ) with pytest.raises(ProxyException, match="is not available for this API key"): @@ -23441,7 +23378,9 @@ class TestMemberAutoRouterInference: router: Final = self._router(self._marker(member=False)) monkeypatch.setitem(sys.modules, "fastapi", None) monkeypatch.delitem(sys.modules, "litellm.proxy.auth.auto_router_checks", raising=False) - assert (await self._route(router, {"metadata": {"user_api_key_team_id": "router-team"}})).model == "restricted-model" + assert ( + await self._route(router, {"metadata": {"user_api_key_team_id": "router-team"}}) + ).model == "restricted-model" def _access_window_offsets(start_hours: float, end_hours: float, team_ids: list) -> dict: @@ -23486,10 +23425,12 @@ def test_access_windows_hide_reserved_deployment_from_other_teams(): def test_access_windows_raise_when_only_reserved_deployments_remain(): - router = Router(model_list=_reserved_model_list( - windows_for_reserved=[_access_window_offsets(-1, 1, ["team-a"])], - windows_for_open=[_access_window_offsets(-1, 1, ["team-a"])], - )[:1]) + router = Router( + model_list=_reserved_model_list( + windows_for_reserved=[_access_window_offsets(-1, 1, ["team-a"])], + windows_for_open=[_access_window_offsets(-1, 1, ["team-a"])], + )[:1] + ) for request_kwargs in ({"metadata": {"user_api_key_team_id": "team-b"}}, {}): with pytest.raises(litellm.BadRequestError, match="reserved for another team"): router._common_checks_available_deployment(model="gpt-4o-ptu", request_kwargs=request_kwargs) @@ -23552,9 +23493,11 @@ def test_access_windows_inactive_window_leaves_deployments_available(): def test_access_windows_apply_when_calling_by_model_id(): - router = Router(model_list=_reserved_model_list( - windows_for_reserved=[_access_window_offsets(-1, 1, ["team-a"])], - )) + router = Router( + model_list=_reserved_model_list( + windows_for_reserved=[_access_window_offsets(-1, 1, ["team-a"])], + ) + ) with pytest.raises(litellm.BadRequestError, match="reserved for another team"): router._common_checks_available_deployment( model="reserved-deployment", @@ -23800,7 +23743,9 @@ async def test_bare_model_group_served_by_wildcard_deployment_uses_provider_pref @pytest.mark.asyncio -async def test_bare_model_group_served_by_wildcard_deployment_uses_provider_prefixed_context_window_fallback_key() -> None: +async def test_bare_model_group_served_by_wildcard_deployment_uses_provider_prefixed_context_window_fallback_key() -> ( + None +): """The context-window chain is keyed the same way the ordinary chain is, so a key spelled like the wildcard deployment ("anthropic/claude-sonnet-4-6") must catch the bare group's context-window error too.""" router = litellm.Router( @@ -23982,9 +23927,7 @@ def test_a_config_entry_declaring_shares_without_terms_is_refused_even_while_the [ GuardrailRaisedException(guardrail_name="chunk-scanner", message="blocked"), HTTPException(status_code=403, detail={"error": "blocked", "guardrail_name": "chunk-scanner"}), - ModifyResponseException( - message="blocked", model="primary", request_data={}, guardrail_name="chunk-scanner" - ), + ModifyResponseException(message="blocked", model="primary", request_data={}, guardrail_name="chunk-scanner"), ], ) async def test_a_guardrail_verdict_is_neither_retried_nor_fallen_back(verdict: Exception) -> None: @@ -24185,6 +24128,7 @@ async def test_failure_rpm_increment_declares_the_router_usage_key_family(): assert seen == ["router_usage"] assert current_service_target() is None + class _SpanRecordingInMemoryCache(InMemoryCache): """Records the live OTel span each read runs under, so the test sees what a Redis span would nest in.""" @@ -24379,10 +24323,13 @@ async def test_non_chat_surfaces_mark_their_deployment_pick(monkeypatch: pytest. router.completion(model="gpt-4o", messages=[{"role": "user", "content": "hi"}]) assert events == [_pick("embed", "initial", 1), _pick("gpt-4o", "initial", 1)] + + class TestAutoRouterTraceProvenance: @pytest.fixture def _tracing(self, monkeypatch: pytest.MonkeyPatch) -> "tuple[OpenTelemetryV2, InMemorySpanExporter]": from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + from litellm.integrations.otel import OpenTelemetryV2Config from litellm.integrations.otel.logger import OpenTelemetryV2 from litellm.integrations.otel.plumbing import providers @@ -24390,7 +24337,9 @@ class TestAutoRouterTraceProvenance: config: Final = OpenTelemetryV2Config(exporter="in_memory") exporter: Final = InMemorySpanExporter() - logger: Final = OpenTelemetryV2(config=config, tracer_provider=providers.build_tracer_provider(config, exporter=exporter)) + logger: Final = OpenTelemetryV2( + config=config, tracer_provider=providers.build_tracer_provider(config, exporter=exporter) + ) monkeypatch.setattr(proxy_server, "open_telemetry_logger", logger) return logger, exporter @@ -24400,7 +24349,8 @@ class TestAutoRouterTraceProvenance: "model_name": "traced-router", "model_info": {"id": f"definition-{tag}", "updated_at": "2026-10-01T12:00:00+00:00"}, "litellm_params": { - "model": "auto_router/complexity_router", "tags": [tag], + "model": "auto_router/complexity_router", + "tags": [tag], "complexity_router_config": { "tiers": dict.fromkeys(("SIMPLE", "MEDIUM", "COMPLEX", "REASONING"), target), }, @@ -24409,17 +24359,25 @@ class TestAutoRouterTraceProvenance: def test_definition_fingerprint_is_canonical_and_tracks_config_not_credentials(self): original: Final = Deployment.model_validate(self._marker("blue")) - reordered: Final = Deployment.model_validate({ - **self._marker("blue"), - "litellm_params": { - "api_key": "credential-only-change", "tags": ["blue"], "model": "auto_router/complexity_router", - "complexity_router_config": {"tiers": dict.fromkeys(("REASONING", "COMPLEX", "MEDIUM", "SIMPLE"), "answer")}, - }, - }) + reordered: Final = Deployment.model_validate( + { + **self._marker("blue"), + "litellm_params": { + "api_key": "credential-only-change", + "tags": ["blue"], + "model": "auto_router/complexity_router", + "complexity_router_config": { + "tiers": dict.fromkeys(("REASONING", "COMPLEX", "MEDIUM", "SIMPLE"), "answer") + }, + }, + } + ) fingerprint: Final = Router._routing_definition_fingerprint(original) assert fingerprint is not None assert fingerprint == Router._routing_definition_fingerprint(reordered) - assert fingerprint != Router._routing_definition_fingerprint(Deployment.model_validate(self._marker("blue", "other"))) + assert fingerprint != Router._routing_definition_fingerprint( + Deployment.model_validate(self._marker("blue", "other")) + ) @pytest.mark.asyncio async def test_direct_classifier_and_plugin_prelude_never_inherit_parent_decisions( @@ -24437,7 +24395,9 @@ class TestAutoRouterTraceProvenance: async def classify(self, context: RoutingContext) -> str: await litellm.acompletion( - model="openai/gpt-5.6-luna", api_key="test", mock_response="SIMPLE", + model="openai/gpt-5.6-luna", + api_key="test", + mock_response="SIMPLE", messages=[{"role": "user", "content": "choose a tier"}], metadata={**context.metadata, "internal_call_origin": "autorouter_classifier"}, ) @@ -24451,16 +24411,25 @@ class TestAutoRouterTraceProvenance: monkeypatch.setattr(litellm, "_async_success_callback", [logger, CallbackComplete()]) classifier: Final = ObservedClassifier() - router: Final = Router(model_list=[ - {**self._marker("blue"), "litellm_params": { - "model": "auto_router/complexity_router", - "complexity_router_config": { - "classifier_type": "custom", "classifier_plugin": classifier, - "tiers": {"SIMPLE": "answer"}, "session_affinity": False, "deployment_affinity": False, + router: Final = Router( + model_list=[ + { + **self._marker("blue"), + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "classifier_type": "custom", + "classifier_plugin": classifier, + "tiers": {"SIMPLE": "answer"}, + "session_affinity": False, + "deployment_affinity": False, + }, + }, }, - }}, - {"model_name": "answer", "litellm_params": {"model": "openai/gpt-5.6-luna", "api_key": "test"}}, - ], plugins=[classifier]) + {"model_name": "answer", "litellm_params": {"model": "openai/gpt-5.6-luna", "api_key": "test"}}, + ], + plugins=[classifier], + ) stale: Final = {"routing_decision": {"router_model_name": "previous"}, "_autorouter_baseline_route": "old"} request: Final = {"metadata": {**stale}, "litellm_metadata": {**stale, "user_api_key_team_id": "team"}} with logger.tracer.start_as_current_span("request") as root: @@ -24470,7 +24439,8 @@ class TestAutoRouterTraceProvenance: ) await completed.wait() classifier_span: Final = next( - span for span in exporter.get_finished_spans() + span + for span in exporter.get_finished_spans() if (span.attributes or {}).get("litellm.request.purpose") == "autorouter_classifier" ) assert not any(key.startswith("litellm.routing.") for key in classifier_span.attributes) @@ -24484,36 +24454,66 @@ class TestAutoRouterTraceProvenance: @pytest.mark.asyncio @pytest.mark.parametrize("build_fails", [False, True]) async def test_semantic_internal_failure_keeps_phase_identity_without_parent_retry_attribution( - self, _tracing: "tuple[OpenTelemetryV2, InMemorySpanExporter]", monkeypatch: pytest.MonkeyPatch, - respx_mock: respx.MockRouter, build_fails: bool, + self, + _tracing: "tuple[OpenTelemetryV2, InMemorySpanExporter]", + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + build_fails: bool, ) -> None: logger, exporter = _tracing monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) query_statuses: Final = iter((503, 200)) def embedding_response(request: httpx.Request) -> httpx.Response: - status: Final = 400 if build_fails else ( - next(query_statuses) if json.loads(request.content)["input"] == ["query-retry"] else 200 + status: Final = ( + 400 + if build_fails + else (next(query_statuses) if json.loads(request.content)["input"] == ["query-retry"] else 200) + ) + return httpx.Response( + status, + headers={"x-should-retry": "false"}, + json={ + "object": "list", + "model": "text-embedding-3-small", + "data": [{"object": "embedding", "index": 0, "embedding": [1.0, 0.0]}], + "usage": {"prompt_tokens": 1, "total_tokens": 1}, + }, ) - return httpx.Response(status, headers={"x-should-retry": "false"}, json={ - "object": "list", "model": "text-embedding-3-small", - "data": [{"object": "embedding", "index": 0, "embedding": [1.0, 0.0]}], - "usage": {"prompt_tokens": 1, "total_tokens": 1}, - }) - upstream: Final = respx_mock.post("https://embedding.example/v1/embeddings").mock(side_effect=embedding_response) - router: Final = Router(model_list=[ - {"model_name": "semantic", "model_info": {"id": "semantic-definition"}, "litellm_params": { - "model": "auto_router/semantic-router", "auto_router_default_model": "answer", - "auto_router_embedding_model": "embed", - "auto_router_config": json.dumps({"routes": [{"name": "answer", "utterances": ["hello"]}]}), - }}, - *({"model_name": "embed", "model_info": {"id": f"embed-{index}"}, "litellm_params": { - "model": "openai/text-embedding-3-small", "api_base": "https://embedding.example/v1", - "api_key": "test", "max_retries": 0, - }} for index in range(2)), - {"model_name": "answer", "litellm_params": {"model": "openai/gpt-5.6-luna", "api_key": "test"}}, - ], num_retries=1, retry_after=0) + upstream: Final = respx_mock.post("https://embedding.example/v1/embeddings").mock( + side_effect=embedding_response + ) + router: Final = Router( + model_list=[ + { + "model_name": "semantic", + "model_info": {"id": "semantic-definition"}, + "litellm_params": { + "model": "auto_router/semantic-router", + "auto_router_default_model": "answer", + "auto_router_embedding_model": "embed", + "auto_router_config": json.dumps({"routes": [{"name": "answer", "utterances": ["hello"]}]}), + }, + }, + *( + { + "model_name": "embed", + "model_info": {"id": f"embed-{index}"}, + "litellm_params": { + "model": "openai/text-embedding-3-small", + "api_base": "https://embedding.example/v1", + "api_key": "test", + "max_retries": 0, + }, + } + for index in range(2) + ), + {"model_name": "answer", "litellm_params": {"model": "openai/gpt-5.6-luna", "api_key": "test"}}, + ], + num_retries=1, + retry_after=0, + ) stale: Final = {"routing_decision": {"router_model_name": "previous"}, "_autorouter_baseline_route": "old"} request: Final = {"metadata": {**stale}, "litellm_metadata": {**stale}} with logger.tracer.start_as_current_span("request") as root: @@ -24545,20 +24545,22 @@ class TestAutoRouterTraceProvenance: self, _tracing: "tuple[OpenTelemetryV2, InMemorySpanExporter]" ) -> None: logger, exporter = _tracing - router: Final = Router(model_list=[ - self._marker("blue"), - { - "model_name": "traced-router", - "model_info": {"id": "definition-semantic"}, - "litellm_params": { - "model": "auto_router/semantic-router", - "tags": ["blue"], - "auto_router_config": json.dumps({"routes": [{"name": "answer", "utterances": ["hello"]}]}), - "auto_router_default_model": "answer", - "auto_router_embedding_model": "embed", + router: Final = Router( + model_list=[ + self._marker("blue"), + { + "model_name": "traced-router", + "model_info": {"id": "definition-semantic"}, + "litellm_params": { + "model": "auto_router/semantic-router", + "tags": ["blue"], + "auto_router_config": json.dumps({"routes": [{"name": "answer", "utterances": ["hello"]}]}), + "auto_router_default_model": "answer", + "auto_router_embedding_model": "embed", + }, }, - }, - ]) + ] + ) with logger.start_phase_span("route traced-router"): response: Final = await router.async_pre_routing_hook( model="traced-router", request_kwargs={"metadata": {"tags": ["blue"]}} @@ -24578,26 +24580,42 @@ class TestAutoRouterTraceProvenance: ) -> None: logger, exporter = _tracing monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - respx_mock.post("https://embedding.example/v1/embeddings").respond(200, json={ - "object": "list", - "model": "text-embedding-3-small", - "data": [{"object": "embedding", "index": 0, "embedding": [1.0, 0.0]}], - "usage": {"prompt_tokens": 1, "total_tokens": 1}, - }) - router: Final = Router(model_list=[ - { - "model_name": "internal-semantic-router", - "model_info": {"id": "semantic-definition", "team_id": "trace-team", "team_public_model_name": "public-router"}, - "litellm_params": { - "model": "auto_router/semantic-router", - "auto_router_config": json.dumps({"routes": [{"name": "answer", "utterances": ["hello"]}]}), - "auto_router_default_model": "answer", - "auto_router_embedding_model": "embed", - }, + respx_mock.post("https://embedding.example/v1/embeddings").respond( + 200, + json={ + "object": "list", + "model": "text-embedding-3-small", + "data": [{"object": "embedding", "index": 0, "embedding": [1.0, 0.0]}], + "usage": {"prompt_tokens": 1, "total_tokens": 1}, }, - {"model_name": "embed", "litellm_params": {"model": "openai/text-embedding-3-small", "api_base": "https://embedding.example/v1", "api_key": "test"}}, - {"model_name": "answer", "litellm_params": {"model": "openai/gpt-5.6-luna", "api_key": "test"}}, - ]) + ) + router: Final = Router( + model_list=[ + { + "model_name": "internal-semantic-router", + "model_info": { + "id": "semantic-definition", + "team_id": "trace-team", + "team_public_model_name": "public-router", + }, + "litellm_params": { + "model": "auto_router/semantic-router", + "auto_router_config": json.dumps({"routes": [{"name": "answer", "utterances": ["hello"]}]}), + "auto_router_default_model": "answer", + "auto_router_embedding_model": "embed", + }, + }, + { + "model_name": "embed", + "litellm_params": { + "model": "openai/text-embedding-3-small", + "api_base": "https://embedding.example/v1", + "api_key": "test", + }, + }, + {"model_name": "answer", "litellm_params": {"model": "openai/gpt-5.6-luna", "api_key": "test"}}, + ] + ) request: Final = {"metadata": {"user_api_key_team_id": "trace-team"}} with logger.tracer.start_as_current_span("request"): await router.async_get_available_deployment( @@ -24609,8 +24627,13 @@ class TestAutoRouterTraceProvenance: assert decision["router_config_id"] == router.model_list[0]["model_info"]["id"] assert decision["cause"] == "semantic_match" spans: Final = {span.name: span for span in exporter.get_finished_spans()} - assert spans["route public-router"].attributes["litellm.routing.router_model_name"] == decision["router_model_name"] - assert spans["route public-router"].attributes["litellm.routing.router_config_id"] == decision["router_config_id"] + assert ( + spans["route public-router"].attributes["litellm.routing.router_model_name"] + == decision["router_model_name"] + ) + assert ( + spans["route public-router"].attributes["litellm.routing.router_config_id"] == decision["router_config_id"] + ) retry: Final = next(event for event in spans["request"].events if event.name == "litellm.routing.retry") assert retry.attributes["litellm.routing.router_model_name"] == decision["router_model_name"] @@ -24620,18 +24643,37 @@ class TestAutoRouterTraceProvenance: self, _tracing: "tuple[OpenTelemetryV2, InMemorySpanExporter]", passthrough: bool ) -> None: logger, exporter = _tracing - router: Final = Router(model_list=[ - self._marker("blue"), self._marker("green"), - {"model_name": "answer", "litellm_params": {"model": "openai/gpt-5.6-luna", "api_key": "test", "use_in_pass_through": True, "tags": ["blue", "green"]}}, - ], enable_tag_filtering=True) + router: Final = Router( + model_list=[ + self._marker("blue"), + self._marker("green"), + { + "model_name": "answer", + "litellm_params": { + "model": "openai/gpt-5.6-luna", + "api_key": "test", + "use_in_pass_through": True, + "tags": ["blue", "green"], + }, + }, + ], + enable_tag_filtering=True, + ) request: Final = {"metadata": {"tags": ["green"]}, "litellm_metadata": {"tags": ["green"]}} - pick: Final = router.async_get_available_deployment_for_pass_through if passthrough else router.async_get_available_deployment + pick: Final = ( + router.async_get_available_deployment_for_pass_through + if passthrough + else router.async_get_available_deployment + ) with logger.tracer.start_as_current_span("request"): await pick(model="traced-router", request_kwargs=request, messages=[{"role": "user", "content": "hello"}]) decision: Final = request["litellm_metadata"]["routing_decision"] assert decision["router_config_id"] == "definition-green" assert decision["router_config_updated_at"] == "2026-10-01T12:00:00+00:00" - assert decision["router_config_fingerprint"] == router.complexity_routers["traced-router"][1].definition_fingerprint + assert ( + decision["router_config_fingerprint"] + == router.complexity_routers["traced-router"][1].definition_fingerprint + ) assert "routing_decision" not in request["metadata"] router.log_retry({**request, "model": "answer"}, ValueError("private-error-text")) await pick(model="answer", request_kwargs=request, messages=[{"role": "user", "content": "hello"}]) @@ -24652,6 +24694,7 @@ def _vcr_outcome_gate(request, vcr): yield record_vcr_outcome(request, vcr) + @pytest.fixture(scope="function") def isolate_litellm_state(): """ @@ -24704,6 +24747,7 @@ def isolate_litellm_state(): setattr(litellm, attr, original_value) _invalidate_model_cost_lowercase_map() + _SCALAR_DEFAULTS = { "num_retries": getattr(litellm, "num_retries", None), "num_retries_per_request": getattr(litellm, "num_retries_per_request", None), @@ -24724,6 +24768,7 @@ _SCALAR_DEFAULTS = { "api_key": getattr(litellm, "api_key", None), } + @pytest.fixture(scope="module") def setup_and_teardown(): """ @@ -24746,6 +24791,7 @@ def setup_and_teardown(): litellm.in_memory_llm_clients_cache.flush_cache() yield + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @patch("azure.identity.get_bearer_token_provider") @patch("azure.identity.ClientSecretCredential") @@ -24827,6 +24873,7 @@ def test_router_init_azure_service_principal_with_secret_with_environment_variab # finally verify if the mocked token was used by Azure SDK mocked_func_generating_token.assert_called() + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.asyncio async def test_audio_speech_router(): @@ -24870,6 +24917,7 @@ async def test_audio_speech_router(): client_passed_in_request = mock_aspeech.call_args.kwargs["client"] assert client_passed_in_request == expected_openai_client + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_update_kwargs_before_fallbacks_unit_test(): router = Router( @@ -24895,6 +24943,7 @@ def test_update_kwargs_before_fallbacks_unit_test(): assert kwargs["litellm_trace_id"] is not None + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.parametrize( "call_type", @@ -24954,6 +25003,7 @@ async def test_update_kwargs_before_fallbacks(call_type): print(mock_client.call_args.kwargs) assert mock_client.call_args.kwargs["litellm_trace_id"] is not None + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_router_get_model_info_wildcard_routes(): os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" @@ -24973,6 +25023,7 @@ def test_router_get_model_info_wildcard_routes(): assert model_info["tpm"] is not None assert model_info["rpm"] is not None + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.asyncio @pytest.mark.flaky(retries=3, delay=1) @@ -25003,6 +25054,7 @@ async def test_router_get_model_group_usage_wildcard_routes(): assert tpm is not None, "tpm is None" assert rpm is not None, "rpm is None" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.asyncio async def test_call_router_callbacks_on_success(): @@ -25036,6 +25088,7 @@ async def test_call_router_callbacks_on_success(): assert increment["key"].startswith("global_router:1:gemini/gemini-2.5-flash:rpm") assert increment["increment_value"] == 1 + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.serial @pytest.mark.asyncio @@ -25064,6 +25117,7 @@ async def test_call_router_callbacks_on_failure(): assert mock_callback.call_args_list[0].kwargs["key"].startswith("global_router:1:gemini/gemini-2.5-flash:rpm") + @pytest.mark.asyncio async def test_rate_limit_error_callback(): """ @@ -25109,13 +25163,8 @@ async def test_rate_limit_error_callback(): except Exception: pass - with patch.object( - customHandler, "log_model_group_rate_limit_error", new=AsyncMock() - ) as mock_client: - - print( - f"customHandler.log_model_group_rate_limit_error: {customHandler.log_model_group_rate_limit_error}" - ) + with patch.object(customHandler, "log_model_group_rate_limit_error", new=AsyncMock()) as mock_client: + print(f"customHandler.log_model_group_rate_limit_error: {customHandler.log_model_group_rate_limit_error}") try: _ = await router.acompletion( @@ -25138,7 +25187,6 @@ async def test_rate_limit_error_callback(): async def test_router_model_group_headers(): os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" litellm.model_cost = litellm.get_model_cost_map(url="") - from litellm.types.utils import OPENAI_RESPONSE_HEADERS router = Router( model_list=[ @@ -25163,12 +25211,12 @@ async def test_router_model_group_headers(): assert "x-ratelimit-remaining-requests" in resp._hidden_params["additional_headers"] assert "x-ratelimit-remaining-tokens" in resp._hidden_params["additional_headers"] + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.asyncio async def test_get_remaining_model_group_usage(): os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" litellm.model_cost = litellm.get_model_cost_map(url="") - from litellm.types.utils import OPENAI_RESPONSE_HEADERS router = Router( model_list=[ @@ -25194,6 +25242,7 @@ async def test_get_remaining_model_group_usage(): assert "x-ratelimit-remaining-requests" in remaining_usage assert "x-ratelimit-remaining-tokens" in remaining_usage + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.parametrize( "potential_access_group, expected_result", @@ -25212,6 +25261,7 @@ def test_router_get_model_access_groups(potential_access_group, expected_result) access_groups = router.is_model_access_group_for_wildcard_route(model_access_group=potential_access_group) assert access_groups == expected_result + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_router_redis_cache(): router = Router(model_list=[{"model_name": "gemini/*", "litellm_params": {"model": "gemini/*"}}]) @@ -25222,6 +25272,7 @@ def test_router_redis_cache(): assert router.cache.redis_cache == redis_cache + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_router_handle_clientside_credential(): """A caller-supplied credential must stay scoped to the current call: it must @@ -25250,6 +25301,7 @@ def test_router_handle_clientside_credential(): assert len(router.get_model_list()) == 1 assert router.get_deployment(model_id=new_deployment.model_info.id) is None + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") async def test_router_clientside_credential_not_reused_by_other_callers(respx_mock, monkeypatch: pytest.MonkeyPatch): """End-to-end regression test for LIT-7811. @@ -25307,6 +25359,7 @@ async def test_router_clientside_credential_not_reused_by_other_callers(respx_mo used_auth_headers = {call.request.headers["authorization"] for call in route.calls[1:]} assert used_auth_headers == {"Bearer configured-key"} + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_router_get_async_openai_model_client(): router = Router( @@ -25323,6 +25376,7 @@ def test_router_get_async_openai_model_client(): model_client = router._get_async_openai_model_client(deployment=MagicMock(), kwargs={}) assert model_client is None + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_router_get_deployment_credentials(): router = Router( @@ -25338,6 +25392,7 @@ def test_router_get_deployment_credentials(): assert credentials is not None assert credentials["api_key"] == "123" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_router_get_deployment_credentials_with_provider(): """ @@ -25382,6 +25437,7 @@ def test_router_get_deployment_credentials_with_provider(): credentials3 = router.get_deployment_credentials_with_provider(model_id="non-existent") assert credentials3 is None + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_router_get_deployment_credentials_with_provider_wildcard(): """ @@ -25429,6 +25485,7 @@ def test_router_get_deployment_credentials_with_provider_wildcard(): credentials3 = router.get_deployment_credentials_with_provider(model_id="vertex_ai/gemini-pro") assert credentials3 is None + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_router_get_deployment_model_info(): router = Router( @@ -25443,12 +25500,14 @@ def test_router_get_deployment_model_info(): model_info = router.get_deployment_model_info(model_id="1", model_name="gemini/gemini-1.5-flash") assert model_info is not None + @pytest.fixture() def _vcr_outcome_gate_router_unit(request, vcr): install_live_call_probe(request, vcr) yield record_vcr_outcome(request, vcr) + @pytest.fixture(scope="function") def setup_and_teardown_router_unit(): """ @@ -25470,6 +25529,7 @@ def setup_and_teardown_router_unit(): loop.close() asyncio.set_event_loop(None) + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") @pytest.mark.asyncio async def test_acompletion_deployment_not_mutated(): @@ -25519,6 +25579,7 @@ async def test_acompletion_deployment_not_mutated(): assert deployment_after is not None assert deployment_after.litellm_params.model_dump() == original_params + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") def test_completion_deployment_not_mutated(): """ @@ -25567,6 +25628,7 @@ def test_completion_deployment_not_mutated(): assert deployment_after is not None assert deployment_after.litellm_params.model_dump() == original_params + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") def test_default_deployment_isolation(): """ @@ -25639,10 +25701,12 @@ def test_default_deployment_isolation(): assert deployment2["litellm_params"]["custom_config"]["nested_setting"] == "modified" # type: ignore assert router.default_deployment["litellm_params"]["custom_config"]["nested_setting"] == "modified" # type: ignore + class NoItemsAliasDict(dict): def items(self): raise AssertionError("Unexpected full alias iteration via items()") + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") def test_get_model_list_from_model_alias_should_not_iterate_for_non_alias_lookup(): router = Router( @@ -25659,6 +25723,7 @@ def test_get_model_list_from_model_alias_should_not_iterate_for_non_alias_lookup model_alias_list = router.get_model_list_from_model_alias(model_name="gpt-5-mini") assert model_alias_list == [] + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") def test_map_team_model_should_not_iterate_aliases_for_non_alias_team_model_name(): router = Router( @@ -25887,9 +25952,11 @@ class TestPreCallChecksOptimization: assert deployments[0].get("model_info", {}).get("id") == "small", "First deployment ID changed!" assert deployments[1].get("model_info", {}).get("id") == "large", "Second deployment ID changed!" + if __name__ == "__main__": pytest.main([__file__, "-v"]) + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") def test_is_prompt_management_model_optimization(): """ @@ -25944,6 +26011,7 @@ def test_is_prompt_management_model_optimization(): finally: litellm._known_custom_logger_compatible_callbacks = original_callbacks + @pytest.fixture def router(): """Create a router with a mock deployment""" @@ -25959,6 +26027,7 @@ def router(): ] ) + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") @pytest.mark.asyncio async def test_router_acancel_batch(router): @@ -25982,6 +26051,7 @@ async def test_router_acancel_batch(router): assert response.id == "batch_123" assert response.status == "cancelled" + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") @pytest.mark.asyncio async def test_router_acancel_batch_resolves_credential_name(): @@ -26022,6 +26092,7 @@ async def test_router_acancel_batch_resolves_credential_name(): finally: litellm.credential_list = [] + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") @pytest.mark.asyncio async def test_router_acancel_batch_removes_unresolved_credential_name(): @@ -26054,6 +26125,7 @@ async def test_router_acancel_batch_removes_unresolved_credential_name(): call_kwargs = mock_cancel.call_args.kwargs assert "litellm_credential_name" not in call_kwargs + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") class TestRouterEmbeddingHeaders: """Test that embedding methods properly propagate headers from router configuration.""" @@ -26368,6 +26440,7 @@ class TestRouterEmbeddingHeaders: assert "litellm_trace_id" in completion_kwargs assert "litellm_trace_id" in embedding_kwargs + if __name__ == "__main__": # Run a simple test test = TestRouterEmbeddingHeaders() @@ -26384,6 +26457,7 @@ OPENAI_EMBEDDINGS_URL = "https://api.openai.com/v1/embeddings" STORE_EMBEDDINGS_URL = "https://embedding.example/v1/embeddings" + def _mock_embedding_route(respx_mock: respx.MockRouter, url: str) -> respx.Route: return respx_mock.post(url).mock( return_value=httpx.Response( @@ -26397,11 +26471,13 @@ def _mock_embedding_route(respx_mock: respx.MockRouter, url: str) -> respx.Route ) ) + def _sent(route: respx.Route, index: int) -> tuple[str, str, list[str]]: request = route.calls[index].request body = json.loads(request.read()) return request.headers["authorization"], body["model"], body["input"] + def _alias_router() -> Router: return Router( model_list=[ @@ -26415,6 +26491,7 @@ def _alias_router() -> Router: ] ) + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") class TestRouterEmbeddingIntegration: """Integration tests for embedding with router configuration.""" @@ -26901,10 +26978,12 @@ class TestRouterEmbeddingIntegration: assert "headers" in call_kwargs assert call_kwargs["headers"]["X-Custom-Azure-Header"] == "azure-value" + if __name__ == "__main__": # Run tests pytest.main([__file__, "-v"]) + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") class TestRouterIndexManagement: """Test cases for router index management functions""" @@ -27665,7 +27744,6 @@ def test_router_context_window_check_pre_call_check_in_group_custom_model_info() """ import os - litellm.set_verbose = False print(f"len(text): {len(text)}") @@ -27707,7 +27785,7 @@ def test_router_context_window_check_pre_call_check_in_group_custom_model_info() assert response.choices[0].message.content == "Hello world 1!" except Exception as e: - pytest.fail(f"Got unexpected exception on router! - {str(e)}") + pytest.fail(f"Got unexpected exception on router! - {e!s}") @pytest.mark.usefixtures("local_model_cost_map", "isolate_router_model_cost_state") @@ -27719,7 +27797,6 @@ def test_router_context_window_check_pre_call_check(): """ import os - litellm.set_verbose = False print(f"len(text): {len(text)}") @@ -27761,7 +27838,7 @@ def test_router_context_window_check_pre_call_check(): assert response.choices[0].message.content == "Hello world 2!" except Exception as e: - pytest.fail(f"Got unexpected exception on router! - {str(e)}") + pytest.fail(f"Got unexpected exception on router! - {e!s}") @pytest.mark.usefixtures("local_model_cost_map", "isolate_router_model_cost_state") @@ -27773,7 +27850,6 @@ def test_router_context_window_check_pre_call_check_out_group(): """ import os - litellm.set_verbose = False print(f"len(text): {len(text)}") @@ -27799,7 +27875,13 @@ def test_router_context_window_check_pre_call_check_out_group(): }, ] - router = Router(model_list=model_list, set_verbose=True, enable_pre_call_checks=True, num_retries=0, context_window_fallbacks=[{"gpt-3.5-turbo-small": ["gpt-3.5-turbo-large"]}]) # type: ignore + router = Router( + model_list=model_list, + set_verbose=True, + enable_pre_call_checks=True, + num_retries=0, + context_window_fallbacks=[{"gpt-3.5-turbo-small": ["gpt-3.5-turbo-large"]}], + ) # type: ignore response = router.completion( model="gpt-3.5-turbo-small", @@ -27812,7 +27894,7 @@ def test_router_context_window_check_pre_call_check_out_group(): print(f"response: {response}") assert response.choices[0].message.content == "Alexander was a great conqueror." except Exception as e: - pytest.fail(f"Got unexpected exception on router! - {str(e)}") + pytest.fail(f"Got unexpected exception on router! - {e!s}") def test_filter_invalid_params_pre_call_check(): @@ -27850,7 +27932,7 @@ def test_filter_invalid_params_pre_call_check(): ) assert len(filtered_deployments) == 1 except Exception as e: - pytest.fail(f"Got unexpected exception on router! - {str(e)}") + pytest.fail(f"Got unexpected exception on router! - {e!s}") @pytest.mark.parametrize("allowed_model_region", ["eu", None, "us"]) @@ -27895,13 +27977,9 @@ def test_router_region_pre_call_check(allowed_model_region): if allowed_model_region is None: assert len(_healthy_deployments) == 2 else: - assert len(_healthy_deployments) == 1, "{} models selected as healthy".format( - len(_healthy_deployments) - ) - assert ( - _healthy_deployments[0]["model_info"]["id"] == "1" - ), "Incorrect model id picked. Got id={}, expected id=1".format( - _healthy_deployments[0]["model_info"]["id"] + assert len(_healthy_deployments) == 1, f"{len(_healthy_deployments)} models selected as healthy" + assert _healthy_deployments[0]["model_info"]["id"] == "1", ( + "Incorrect model id picked. Got id={}, expected id=1".format(_healthy_deployments[0]["model_info"]["id"]) ) @@ -28072,14 +28150,12 @@ def test_router_context_window_pre_call_check(model, base_model, llm_provider): if llm_provider == "azure" and base_model is None: assert len(updated_list) == 1 else: - pytest.fail("Expected to raise an error. Got={}".format(updated_list)) - except Exception as e: - if ( - llm_provider == "azure" and base_model is not None - ) or llm_provider == "openai": + pytest.fail(f"Expected to raise an error. Got={updated_list}") + except Exception: + if (llm_provider == "azure" and base_model is not None) or llm_provider == "openai": pass except Exception as e: - pytest.fail(f"Got unexpected exception on router! - {str(e)}") + pytest.fail(f"Got unexpected exception on router! - {e!s}") @pytest.fixture @@ -28153,9 +28229,7 @@ def test_router_dynamic_cooldown_correct_retry_after_time(): ): new_retry_after_mock_client = MagicMock(return_value=-1) - litellm.utils._get_retry_after_from_exception_header = ( - new_retry_after_mock_client - ) + litellm.utils._get_retry_after_from_exception_header = new_retry_after_mock_client try: router.embedding( @@ -28241,9 +28315,7 @@ async def test_aaarouter_dynamic_cooldown_message_retry_time(sync_mode): ) if sync_mode: - cooldown_deployments = get_cooldown_deployments( - litellm_router_instance=router, parent_otel_span=None - ) + cooldown_deployments = get_cooldown_deployments(litellm_router_instance=router, parent_otel_span=None) else: cooldown_deployments = await async_get_cooldown_deployments( litellm_router_instance=router, parent_otel_span=None @@ -28642,8 +28714,8 @@ class CompletionCustomHandler( # Class variables or attributes def __init__(self): self.errors = [] - self.states: Optional[ - List[ + self.states: ( + list[ Literal[ "sync_pre_api_call", "async_pre_api_call", @@ -28656,7 +28728,8 @@ class CompletionCustomHandler( "async_failure", ] ] - ] = [] + | None + ) = [] def log_pre_api_call(self, model, messages, kwargs): try: @@ -28680,14 +28753,10 @@ class CompletionCustomHandler( assert isinstance(kwargs["litellm_params"]["metadata"]["deployment"], str) assert isinstance(kwargs["litellm_params"]["model_info"], dict) assert isinstance(kwargs["litellm_params"]["model_info"]["id"], str) - assert isinstance( - kwargs["litellm_params"]["proxy_server_request"], (str, type(None)) - ) - assert isinstance( - kwargs["litellm_params"]["preset_cache_key"], (str, type(None)) - ) + assert isinstance(kwargs["litellm_params"]["proxy_server_request"], (str, type(None))) + assert isinstance(kwargs["litellm_params"]["preset_cache_key"], (str, type(None))) assert isinstance(kwargs["litellm_params"]["stream_response"], dict) - except Exception as e: + except Exception: print(f"Assertion Error: {traceback.format_exc()}") self.errors.append(traceback.format_exc()) @@ -28711,9 +28780,7 @@ class CompletionCustomHandler( assert isinstance(kwargs["input"], (list, dict, str)) assert isinstance(kwargs["api_key"], (str, type(None))) assert ( - isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) - ) + isinstance(kwargs["original_response"], (str, litellm.CustomStreamWrapper)) or inspect.iscoroutine(kwargs["original_response"]) or inspect.isasyncgen(kwargs["original_response"]) ) @@ -28725,12 +28792,8 @@ class CompletionCustomHandler( assert isinstance(kwargs["litellm_params"]["metadata"]["deployment"], str) assert isinstance(kwargs["litellm_params"]["model_info"], dict) assert isinstance(kwargs["litellm_params"]["model_info"]["id"], str) - assert isinstance( - kwargs["litellm_params"]["proxy_server_request"], (str, type(None)) - ) - assert isinstance( - kwargs["litellm_params"]["preset_cache_key"], (str, type(None)) - ) + assert isinstance(kwargs["litellm_params"]["proxy_server_request"], (str, type(None))) + assert isinstance(kwargs["litellm_params"]["preset_cache_key"], (str, type(None))) assert isinstance(kwargs["litellm_params"]["stream_response"], dict) except Exception: print(f"Assertion Error: {traceback.format_exc()}") @@ -28747,23 +28810,18 @@ class CompletionCustomHandler( assert isinstance(response_obj, litellm.ModelResponseStream) ## KWARGS assert isinstance(kwargs["model"], str) - assert isinstance(kwargs["messages"], list) and isinstance( - kwargs["messages"][0], dict - ) + assert isinstance(kwargs["messages"], list) and isinstance(kwargs["messages"][0], dict) assert isinstance(kwargs["optional_params"], dict) assert isinstance(kwargs["litellm_params"], dict) assert isinstance(kwargs["start_time"], (datetime, type(None))) assert isinstance(kwargs["stream"], bool) assert isinstance(kwargs["user"], (str, type(None))) - assert ( - isinstance(kwargs["input"], list) - and isinstance(kwargs["input"][0], dict) - ) or isinstance(kwargs["input"], (dict, str)) + assert (isinstance(kwargs["input"], list) and isinstance(kwargs["input"][0], dict)) or isinstance( + kwargs["input"], (dict, str) + ) assert isinstance(kwargs["api_key"], (str, type(None))) assert ( - isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) - ) + isinstance(kwargs["original_response"], (str, litellm.CustomStreamWrapper)) or inspect.isasyncgen(kwargs["original_response"]) or inspect.iscoroutine(kwargs["original_response"]) ) @@ -28784,22 +28842,17 @@ class CompletionCustomHandler( assert isinstance(response_obj, litellm.ModelResponse) ## KWARGS assert isinstance(kwargs["model"], str) - assert isinstance(kwargs["messages"], list) and isinstance( - kwargs["messages"][0], dict - ) + assert isinstance(kwargs["messages"], list) and isinstance(kwargs["messages"][0], dict) assert isinstance(kwargs["optional_params"], dict) assert isinstance(kwargs["litellm_params"], dict) assert isinstance(kwargs["start_time"], (datetime, type(None))) assert isinstance(kwargs["stream"], bool) assert isinstance(kwargs["user"], (str, type(None))) - assert ( - isinstance(kwargs["input"], list) - and isinstance(kwargs["input"][0], dict) - ) or isinstance(kwargs["input"], (dict, str)) - assert isinstance(kwargs["api_key"], (str, type(None))) - assert isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) + assert (isinstance(kwargs["input"], list) and isinstance(kwargs["input"][0], dict)) or isinstance( + kwargs["input"], (dict, str) ) + assert isinstance(kwargs["api_key"], (str, type(None))) + assert isinstance(kwargs["original_response"], (str, litellm.CustomStreamWrapper)) assert isinstance(kwargs["additional_args"], (dict, type(None))) assert isinstance(kwargs["log_event_type"], str) assert kwargs["cache_hit"] is None or isinstance(kwargs["cache_hit"], bool) @@ -28818,23 +28871,18 @@ class CompletionCustomHandler( assert response_obj == None ## KWARGS assert isinstance(kwargs["model"], str) - assert isinstance(kwargs["messages"], list) and isinstance( - kwargs["messages"][0], dict - ) + assert isinstance(kwargs["messages"], list) and isinstance(kwargs["messages"][0], dict) assert isinstance(kwargs["optional_params"], dict) assert isinstance(kwargs["litellm_params"], dict) assert isinstance(kwargs["start_time"], (datetime, type(None))) assert isinstance(kwargs["stream"], bool) assert isinstance(kwargs["user"], (str, type(None))) - assert ( - isinstance(kwargs["input"], list) - and isinstance(kwargs["input"][0], dict) - ) or isinstance(kwargs["input"], (dict, str)) + assert (isinstance(kwargs["input"], list) and isinstance(kwargs["input"][0], dict)) or isinstance( + kwargs["input"], (dict, str) + ) assert isinstance(kwargs["api_key"], (str, type(None))) assert ( - isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) - ) + isinstance(kwargs["original_response"], (str, litellm.CustomStreamWrapper)) or kwargs["original_response"] == None ) assert isinstance(kwargs["additional_args"], (dict, type(None))) @@ -28849,8 +28897,7 @@ class CompletionCustomHandler( No-op. Not implemented yet. """ - pass - except Exception as e: + except Exception: print(f"Assertion Error: {traceback.format_exc()}") self.errors.append(traceback.format_exc()) @@ -28867,29 +28914,19 @@ class CompletionCustomHandler( ## END TIME assert isinstance(end_time, datetime) ## RESPONSE OBJECT - assert isinstance( - response_obj, (litellm.ModelResponse, litellm.EmbeddingResponse) - ) + assert isinstance(response_obj, (litellm.ModelResponse, litellm.EmbeddingResponse)) ## KWARGS assert isinstance(kwargs["model"], str) # checking we use base_model for azure cost calculation - base_model = litellm.utils.get_base_model_from_metadata( - model_call_details=kwargs - ) + base_model = litellm.utils.get_base_model_from_metadata(model_call_details=kwargs) - if ( - kwargs["model"] == "chatgpt-v-3" - and base_model is not None - and kwargs["stream"] != True - ): + if kwargs["model"] == "chatgpt-v-3" and base_model is not None and kwargs["stream"] != True: # when base_model is set for azure, we should use pricing for the base_model # this checks response_cost == litellm.cost_per_token(model=base_model) assert isinstance(kwargs["response_cost"], float) response_cost = kwargs["response_cost"] - print( - f"response_cost: {response_cost}, for model: {kwargs['model']} and base_model: {base_model}" - ) + print(f"response_cost: {response_cost}, for model: {kwargs['model']} and base_model: {base_model}") prompt_tokens = response_obj.usage.prompt_tokens completion_tokens = response_obj.usage.completion_tokens # ensure the pricing is based on the base_model here @@ -28900,9 +28937,9 @@ class CompletionCustomHandler( ) expected_price = prompt_price + completion_price print(f"expected price: {expected_price}") - assert ( - response_cost == expected_price - ), f"response_cost: {response_cost} != expected_price: {expected_price}. For model: {kwargs['model']} and base_model: {base_model}. should have used base_model for price" + assert response_cost == expected_price, ( + f"response_cost: {response_cost} != expected_price: {expected_price}. For model: {kwargs['model']} and base_model: {base_model}. should have used base_model for price" + ) assert isinstance(kwargs["messages"], list) assert isinstance(kwargs["optional_params"], dict) @@ -28913,9 +28950,7 @@ class CompletionCustomHandler( assert isinstance(kwargs["input"], (list, dict, str)) assert isinstance(kwargs["api_key"], (str, type(None))) assert ( - isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) - ) + isinstance(kwargs["original_response"], (str, litellm.CustomStreamWrapper)) or inspect.isasyncgen(kwargs["original_response"]) or inspect.iscoroutine(kwargs["original_response"]) ) @@ -28928,12 +28963,8 @@ class CompletionCustomHandler( assert isinstance(kwargs["litellm_params"]["metadata"]["deployment"], str) assert isinstance(kwargs["litellm_params"]["model_info"], dict) assert isinstance(kwargs["litellm_params"]["model_info"]["id"], str) - assert isinstance( - kwargs["litellm_params"]["proxy_server_request"], (str, type(None)) - ) - assert isinstance( - kwargs["litellm_params"]["preset_cache_key"], (str, type(None)) - ) + assert isinstance(kwargs["litellm_params"]["proxy_server_request"], (str, type(None))) + assert isinstance(kwargs["litellm_params"]["preset_cache_key"], (str, type(None))) assert isinstance(kwargs["litellm_params"]["stream_response"], dict) except Exception: print(f"Assertion Error: {traceback.format_exc()}") @@ -28960,9 +28991,7 @@ class CompletionCustomHandler( assert isinstance(kwargs["input"], (list, str, dict)) assert isinstance(kwargs["api_key"], (str, type(None))) assert ( - isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) - ) + isinstance(kwargs["original_response"], (str, litellm.CustomStreamWrapper)) or inspect.isasyncgen(kwargs["original_response"]) or inspect.iscoroutine(kwargs["original_response"]) or kwargs["original_response"] == None @@ -29105,7 +29134,6 @@ class SpecialEndpointRoutingStrategy(CustomRoutingStrategyBase): if isinstance(model, dict): if model["litellm_params"]["model"] == "openai/very-special-endpoint": return model - pass def get_available_deployment( self, @@ -29238,9 +29266,7 @@ def test_multiple_deployments_route_to_a_configured_model( }, ) - route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( - side_effect=_deployment_response - ) + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock(side_effect=_deployment_response) router: Final = Router( model_list=[ { @@ -29389,9 +29415,7 @@ def test_router_retries_transient_server_error( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - route: Final = respx_mock.post( - "https://api.openai.com/v1/chat/completions" - ).mock( + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( side_effect=[ httpx.Response(500, json={"error": {"message": "temporary failure"}}), _openai_chat_response("recovered"), @@ -29424,9 +29448,7 @@ def test_router_authentication_error_is_not_swallowed( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - route: Final = respx_mock.post( - "https://api.openai.com/v1/chat/completions" - ).mock( + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( side_effect=[ httpx.Response( 401, @@ -29524,9 +29546,9 @@ def test_router_reads_api_key_from_model_list( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - route: Final = respx_mock.post( - "https://api.openai.com/v1/chat/completions" - ).mock(side_effect=[_openai_chat_response(), _openai_chat_stream_response()]) + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + side_effect=[_openai_chat_response(), _openai_chat_stream_response()] + ) router: Final = Router( model_list=[ { @@ -29549,9 +29571,7 @@ def test_router_reads_api_key_from_model_list( messages=[{"role": "user", "content": "stream key from model list"}], stream=True, ) - streamed_content: Final = "".join( - chunk.choices[0].delta.content or "" for chunk in stream - ) + streamed_content: Final = "".join(chunk.choices[0].delta.content or "" for chunk in stream) assert response.choices[0].message.content == "scripted reply" assert streamed_content == "streamed reply" @@ -29567,9 +29587,9 @@ def test_router_calls_requested_specific_deployment( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - route: Final = respx_mock.post( - "https://api.openai.com/v1/chat/completions" - ).mock(return_value=_openai_chat_response()) + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + return_value=_openai_chat_response() + ) router: Final = Router( model_list=[ { @@ -29611,9 +29631,7 @@ async def test_router_context_window_error_uses_configured_fallback( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - route: Final = respx_mock.post( - "https://api.openai.com/v1/chat/completions" - ).mock( + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( side_effect=[ httpx.Response( 400, @@ -29656,9 +29674,7 @@ async def test_router_context_window_error_uses_configured_fallback( assert response.choices[0].message.content == "fallback response" assert route.call_count == 2 - assert tuple( - json.loads(call.request.content)["model"] for call in route.calls - ) == ("gpt-4o-mini", "gpt-4.1-mini") + assert tuple(json.loads(call.request.content)["model"] for call in route.calls) == ("gpt-4o-mini", "gpt-4.1-mini") def test_router_function_calling_sends_tools_and_returns_tool_call( @@ -29666,9 +29682,7 @@ def test_router_function_calling_sends_tools_and_returns_tool_call( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - route: Final = respx_mock.post( - "https://api.openai.com/v1/chat/completions" - ).mock( + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( return_value=httpx.Response( 200, json={ @@ -29747,9 +29761,7 @@ async def test_router_openai_completion_supports_sync_async_and_streaming( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - route: Final = respx_mock.post( - "https://api.openai.com/v1/chat/completions" - ).mock( + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( side_effect=[ _openai_chat_response("async reply"), _openai_chat_response("sync reply"), @@ -29784,17 +29796,13 @@ async def test_router_openai_completion_supports_sync_async_and_streaming( stream=True, ) async_stream_chunks: Final = tuple([chunk async for chunk in async_stream]) - async_stream_content: Final = "".join( - chunk.choices[0].delta.content or "" for chunk in async_stream_chunks - ) + async_stream_content: Final = "".join(chunk.choices[0].delta.content or "" for chunk in async_stream_chunks) sync_stream: Final = router.completion( model="openai-model", messages=[{"role": "user", "content": "sync stream"}], stream=True, ) - sync_stream_content: Final = "".join( - chunk.choices[0].delta.content or "" for chunk in sync_stream - ) + sync_stream_content: Final = "".join(chunk.choices[0].delta.content or "" for chunk in sync_stream) assert async_response.choices[0].message.content == "async reply" assert sync_response.choices[0].message.content == "sync reply" @@ -29838,9 +29846,7 @@ async def test_router_moderation_returns_provider_result( num_retries=0, ) - response: Final = await router.amoderation( - model="moderation-model", input="safe scripted text" - ) + response: Final = await router.amoderation(model="moderation-model", input="safe scripted text") assert response.results[0].flagged is False assert json.loads(route.calls[0].request.content)["input"] == "safe scripted text" @@ -30009,9 +30015,7 @@ async def test_batch_completion_multiple_models_returns_each_model_response( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - route: Final = respx_mock.post("https://batch-router.test/v1/chat/completions").mock( - side_effect=_batch_response - ) + route: Final = respx_mock.post("https://batch-router.test/v1/chat/completions").mock(side_effect=_batch_response) router: Final = _batch_router() responses: Final = cast( @@ -30034,9 +30038,7 @@ async def test_batch_completion_multiple_models_and_messages_preserves_each_pair respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - route: Final = respx_mock.post("https://batch-router.test/v1/chat/completions").mock( - side_effect=_batch_response - ) + route: Final = respx_mock.post("https://batch-router.test/v1/chat/completions").mock(side_effect=_batch_response) router: Final = _batch_router() responses: Final = cast( @@ -30117,9 +30119,7 @@ async def test_batch_completion_fastest_response_streams_mocked_chunks() -> None ) chunks: Final = tuple([chunk async for chunk in response]) text: Final = "".join( - chunk.choices[0].delta.content or "" - for chunk in chunks - if chunk.choices and chunk.choices[0].delta is not None + chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices and chunk.choices[0].delta is not None ) assert text == "stream answer" @@ -30275,7 +30275,9 @@ def _first_chunk_cases(): from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices chunk = ModelResponseStream(choices=[StreamingChoices(delta=Delta(content="hi"))]) - provider_error = litellm.InternalServerError(message="primary died at stream start", llm_provider="openai", model="m") + provider_error = litellm.InternalServerError( + message="primary died at stream start", llm_provider="openai", model="m" + ) wrapped = MidStreamFallbackError( message=str(provider_error), model="m", @@ -30561,7 +30563,9 @@ def test_pre_call_checks_keeps_deployment_whose_unsupported_param_is_allowed_or_ assert [d["model_info"]["id"] for d in filtered] == ["only"] -def test_pre_call_checks_keeps_retryable_429_when_a_cooling_down_deployment_supports_the_param(monkeypatch: pytest.MonkeyPatch): +def test_pre_call_checks_keeps_retryable_429_when_a_cooling_down_deployment_supports_the_param( + monkeypatch: pytest.MonkeyPatch, +): monkeypatch.setattr(litellm, "drop_params", False) model_list: Final = [ { @@ -30584,7 +30588,9 @@ def test_pre_call_checks_keeps_retryable_429_when_a_cooling_down_deployment_supp assert filtered == [] -def test_completion_uses_context_window_fallback_when_the_small_deployment_also_lacks_the_param(monkeypatch: pytest.MonkeyPatch): +def test_completion_uses_context_window_fallback_when_the_small_deployment_also_lacks_the_param( + monkeypatch: pytest.MonkeyPatch, +): monkeypatch.setattr(litellm, "drop_params", False) router: Final = Router( model_list=[ @@ -30653,7 +30659,9 @@ def test_pre_call_checks_prefers_a_deployment_that_takes_the_param_over_dropping assert [d["model_info"]["id"] for d in filtered] == ["gpt-4o"] -def test_pre_call_checks_drops_the_param_when_the_capable_deployment_is_outside_the_region(monkeypatch: pytest.MonkeyPatch): +def test_pre_call_checks_drops_the_param_when_the_capable_deployment_is_outside_the_region( + monkeypatch: pytest.MonkeyPatch, +): monkeypatch.setattr(litellm, "drop_params", False) model_list: Final = [ { @@ -30898,7 +30906,9 @@ def test_update_kwargs_with_deployment_reads_stream_options_from_pydantic_litell Router()._update_kwargs_with_deployment( deployment={ "model_name": "gpt-5-mini", - "litellm_params": LiteLLM_Params(model="openai/gpt-5-mini", api_key="fake-key", stream_options=stream_options), + "litellm_params": LiteLLM_Params( + model="openai/gpt-5-mini", api_key="fake-key", stream_options=stream_options + ), "model_info": {"id": "pydantic-params"}, }, kwargs=kwargs, @@ -31296,7 +31306,10 @@ _IMAGE_RESPONSES_INPUT = [ ], } ] -_VISION_LEG: Final[dict[str, object]] = {"litellm_params": {"model": "openai/gpt-5.6-terra"}, "model_info": {"id": "vision"}} +_VISION_LEG: Final[dict[str, object]] = { + "litellm_params": {"model": "openai/gpt-5.6-terra"}, + "model_info": {"id": "vision"}, +} _DECLARED_TEXT_LEG: Final[dict[str, object]] = { "litellm_params": {"model": "bedrock/openai.gpt-oss-120b-1:0"}, "model_info": {"id": "declared-text", "supports_vision": False}, @@ -31359,7 +31372,10 @@ def test_router_sends_an_image_to_the_text_only_leg_when_the_vision_leg_is_cooli def test_without_declared_text_only_deployments_passes_a_chosen_deployment_through(): - assert litellm.Router._without_declared_text_only_deployments(_DECLARED_TEXT_LEG, "mixed", _IMAGE_CHAT, None) is _DECLARED_TEXT_LEG + assert ( + litellm.Router._without_declared_text_only_deployments(_DECLARED_TEXT_LEG, "mixed", _IMAGE_CHAT, None) + is _DECLARED_TEXT_LEG + ) def test_image_routing_drop_declared_text_only_legs_for_image_requests(): @@ -31412,16 +31428,29 @@ def test_image_routing_resolve_vision_through_the_cost_map_and_base_model(): ("deployment", "expected"), [ (_DECLARED_TEXT_LEG, True), - ({"litellm_params": {"model": "bedrock/openai.gpt-oss-120b-1:0"}, "model_info": {"supports_vision": True}}, False), + ( + {"litellm_params": {"model": "bedrock/openai.gpt-oss-120b-1:0"}, "model_info": {"supports_vision": True}}, + False, + ), ({"litellm_params": {"model": "bedrock/deepseek.v3.2"}, "model_info": {}}, True), ({"litellm_params": {"model": "azure/my-deploy"}, "model_info": {"base_model": "azure/gpt-audio"}}, True), ({"litellm_params": {"model": "hosted_vllm/Qwen2.5-VL-7B-Instruct"}, "model_info": {}}, False), ({"litellm_params": {"model": "no-such-model-anywhere"}, "model_info": {}}, False), ({"litellm_params": {}, "model_info": {}}, False), ], - ids=["override-false", "override-true", "cost-map-false", "base-model-false", "undeclared", "unresolvable", "no-params"], + ids=[ + "override-false", + "override-true", + "cost-map-false", + "base-model-false", + "undeclared", + "unresolvable", + "no-params", + ], ) -def test_deployment_declares_no_vision_reads_the_override_then_the_cost_map(deployment: dict[str, object], expected: bool): +def test_deployment_declares_no_vision_reads_the_override_then_the_cost_map( + deployment: dict[str, object], expected: bool +): assert litellm.Router._deployment_declares_no_vision(deployment, "mixed") is expected @@ -31453,19 +31482,31 @@ async def test_anthropic_caller_rate_limit_does_not_block_other_credentials( headers={"retry-after": "168307"}, json={"type": "error", "error": {"type": "rate_limit_error", "message": "usage limit"}}, ) - return httpx.Response(200, json={ - "id": "msg-healthy", "type": "message", "role": "assistant", "model": "claude-sonnet-5-5", - "content": [{"type": "text", "text": "healthy account"}], "stop_reason": "end_turn", - "usage": {"input_tokens": 2, "output_tokens": 2}, - }) + return httpx.Response( + 200, + json={ + "id": "msg-healthy", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [{"type": "text", "text": "healthy account"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 2, "output_tokens": 2}, + }, + ) route: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock(side_effect=upstream) router: Final = Router( - model_list=[{ - "model_name": "claude", "litellm_params": { - "model": "anthropic/claude-sonnet-5-5", "api_key": "fake-gateway-key", - }, "model_info": {"id": "shared-anthropic"}, - }], + model_list=[ + { + "model_name": "claude", + "litellm_params": { + "model": "anthropic/claude-sonnet-5-5", + "api_key": "fake-gateway-key", + }, + "model_info": {"id": "shared-anthropic"}, + } + ], num_retries=0, allowed_fails_policy=AllowedFailsPolicy(RateLimitErrorAllowedFails=1), ) @@ -31473,22 +31514,36 @@ async def test_anthropic_caller_rate_limit_does_not_block_other_credentials( async def call(account: str | None) -> object: credentials: Final = ( - {"provider_specific_header": { - "custom_llm_provider": "anthropic", "extra_headers": { - "Authorization": f"Bearer sk-ant-oat-fake-{account}", - }, - }} if credential_kind == "oauth" else {"api_key": f"fake-{account}"} - ) if account is not None else {} + ( + { + "provider_specific_header": { + "custom_llm_provider": "anthropic", + "extra_headers": { + "Authorization": f"Bearer sk-ant-oat-fake-{account}", + }, + } + } + if credential_kind == "oauth" + else {"api_key": f"fake-{account}"} + ) + if account is not None + else {} + ) if surface == "messages": max_tokens: Final = litellm.get_max_tokens("anthropic/claude-sonnet-5-5") assert max_tokens is not None return await router.aanthropic_messages( - model="claude", messages=[{"role": "user", "content": "hello"}], - client=client, max_tokens=max_tokens, **credentials, + model="claude", + messages=[{"role": "user", "content": "hello"}], + client=client, + max_tokens=max_tokens, + **credentials, ) return await router.acompletion( - model="claude", messages=[{"role": "user", "content": "hello"}], - client=client, **credentials, + model="claude", + messages=[{"role": "user", "content": "hello"}], + client=client, + **credentials, ) try: @@ -31726,7 +31781,9 @@ async def test_a_client_model_id_cannot_move_a_sagemaker_embedding_call_off_the_ served: Final = component_embeddings[str(request["InferenceComponentName"])] return { "Body": io.BytesIO( - json.dumps({"object": "list", "data": [{"object": "embedding", "index": 0, "embedding": served}]}).encode() + json.dumps( + {"object": "list", "data": [{"object": "embedding", "index": 0, "embedding": served}]} + ).encode() ) } @@ -31759,7 +31816,9 @@ async def test_a_header_pin_holds_when_the_deployment_names_sagemaker_through_cu async def test_a_component_header_on_a_non_sagemaker_deployment_leaves_the_clients_model_id_alone( respx_mock: respx.MockRouter, _router_httpx_transport: None ) -> None: - route: Final = respx_mock.post("https://gateway.local/v1/chat/completions").respond(json=dict(_CHAT_COMPLETION_REPLY)) + route: Final = respx_mock.post("https://gateway.local/v1/chat/completions").respond( + json=dict(_CHAT_COMPLETION_REPLY) + ) router: Final = Router( model_list=[ { @@ -31774,7 +31833,9 @@ async def test_a_component_header_on_a_non_sagemaker_deployment_leaves_the_clien ] ) - await router.acompletion(model="gateway-model", messages=[{"role": "user", "content": "hi"}], model_id="client-value") + await router.acompletion( + model="gateway-model", messages=[{"role": "user", "content": "hi"}], model_id="client-value" + ) assert json.loads(route.calls.last.request.content)["model_id"] == "client-value" From 961290fc46429d48de4664d49fb3d64dfa62467a Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 21:28:31 +0000 Subject: [PATCH 10/17] refactor(redis): move search free-tier allocation into a typed store --- .../proxy/hooks/proxy_track_cost_callback.py | 7 +- litellm/proxy/hooks/search_free_tier.py | 33 ++- .../proxy/spend_tracking/cost_finalization.py | 12 +- litellm/redis/stores/search_free_tier.py | 41 +++ tests/_support/fake_redis.py | 30 ++- .../hooks/test_proxy_track_cost_callback.py | 59 +++-- .../unit/proxy/hooks/test_search_free_tier.py | 192 +++++++++++---- .../test_success_handler.py | 45 ++-- .../spend_tracking/test_cost_finalization.py | 194 ++++++++++----- .../redis/stores/test_search_free_tier.py | 233 ++++++++++++++++++ 10 files changed, 670 insertions(+), 176 deletions(-) create mode 100644 litellm/redis/stores/search_free_tier.py create mode 100644 tests/unit/redis/stores/test_search_free_tier.py diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 0d333d4fd08..bca9f8fb843 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -68,6 +68,7 @@ from litellm.proxy.utils import ( BilledFailedAttempt, ProxyUpdateSpend, ) +from litellm.redis.stores.search_free_tier import TransportSearchFreeTierStore from litellm.types.router import Deployment from litellm.types.utils import ( CallTypes, @@ -711,7 +712,11 @@ class _ProxyDBLogger(CustomLogger): sl_object=sl_object, # pyright: ignore[reportUnknownArgumentType] # payload carries untyped lists response_cost=priced_cost, completion_response=completion_response, - counter=proxy_logging_obj.internal_usage_cache.dual_cache.redis_cache, + counter=( + TransportSearchFreeTierStore(search_redis.transport()) + if (search_redis := proxy_logging_obj.internal_usage_cache.dual_cache.redis_cache) is not None + else None + ), now=datetime.now(timezone.utc), ) if litellm.apply_search_free_tier diff --git a/litellm/proxy/hooks/search_free_tier.py b/litellm/proxy/hooks/search_free_tier.py index 57123efbcec..76edc8f0d2d 100644 --- a/litellm/proxy/hooks/search_free_tier.py +++ b/litellm/proxy/hooks/search_free_tier.py @@ -1,8 +1,8 @@ from collections.abc import MutableMapping from dataclasses import dataclass -from datetime import datetime +from datetime import datetime, timedelta from types import MappingProxyType -from typing import Final, Protocol +from typing import Final from pydantic import BaseModel, ConfigDict, ValidationError @@ -12,6 +12,8 @@ from litellm.litellm_core_utils.duration_parser import get_next_standardized_res from litellm.litellm_core_utils.llm_cost_calc.finalization import finalize_cost from litellm.litellm_core_utils.llm_cost_calc.utils import get_web_search_requests_from_usage from litellm.llms.gemini.cost_calculator import cost_per_web_search_request +from litellm.redis.stores.search_free_tier import SearchFreeTierStore, allocation_key +from litellm.redis.transport.contract import RedisFailure from litellm.types.utils import ModelInfo, SearchFreeRequests, StandardLoggingPayload, Usage _PERIOD_DURATION: Final = MappingProxyType({"day": "1d", "month": "1mo"}) @@ -20,10 +22,6 @@ _TTL_GRACE_SECONDS: Final = 3600 _COUNTED_KWARG: Final = "search_free_tier_counted" -class SharedCounter(Protocol): - async def async_increment(self, key: str, value: float, ttl: int | None = None) -> float: ... - - class _UsageCarrier(BaseModel): model_config = ConfigDict(frozen=True, extra="ignore", from_attributes=True) @@ -37,7 +35,7 @@ def freed_units(allowance: int, units: int, total_after: int) -> int: def search_free_tier_key(provider: str, allowance: SearchFreeRequests, now: datetime) -> str: period_stamp: Final = now.strftime(_PERIOD_STAMP[allowance["period"]]) - return f"litellm:search_free_tier:{provider}:{allowance['shared_group']}:{period_stamp}" + return allocation_key(provider, allowance["shared_group"], period_stamp) def search_free_tier_ttl(allowance: SearchFreeRequests, now: datetime) -> int: @@ -91,7 +89,7 @@ async def apply_search_free_tier( sl_object: StandardLoggingPayload, response_cost: float, completion_response: object, - counter: SharedCounter | None, + counter: SearchFreeTierStore | None, now: datetime, ) -> float: """Zero the grounding queries that landed under the model's shared free allowance, everything else is unchanged.""" @@ -109,7 +107,7 @@ async def allocate_search_free_tier( sl_object: StandardLoggingPayload, response_cost: float, completion_response: object, - counter: SharedCounter | None, + counter: SearchFreeTierStore | None, now: datetime, ) -> SearchAllowanceAllocation | None: model_info: Final = sl_object["model_map_information"]["model_map_value"] @@ -126,14 +124,15 @@ async def allocate_search_free_tier( ) return None provider: Final = sl_object["custom_llm_provider"] or model_info["litellm_provider"] - key: Final = search_free_tier_key(provider, allowance, now) - try: - total_after: Final = await counter.async_increment(key, units, ttl=search_free_tier_ttl(allowance, now)) - except Exception: # noqa: BLE001 # any counter failure means full price, never a lost spend row - verbose_proxy_logger.warning( - "search free tier counter %s is unavailable, grounding is billed at full price", key - ) + total_after: Final = await counter.allocate_units( + provider, + allowance["shared_group"], + now.strftime(_PERIOD_STAMP[allowance["period"]]), + units, + ttl=timedelta(seconds=search_free_tier_ttl(allowance, now)), + ) + if isinstance(total_after, RedisFailure): return None freed: Final = freed_units(allowance["count"], units, int(total_after)) freed_cost: Final = min(response_cost, cost_per_web_search_request(usage, model_info) * freed / units) - return SearchAllowanceAllocation(key, units, freed, freed_cost) + return SearchAllowanceAllocation(search_free_tier_key(provider, allowance, now), units, freed, freed_cost) diff --git a/litellm/proxy/spend_tracking/cost_finalization.py b/litellm/proxy/spend_tracking/cost_finalization.py index b260412331e..7dbf1e42689 100644 --- a/litellm/proxy/spend_tracking/cost_finalization.py +++ b/litellm/proxy/spend_tracking/cost_finalization.py @@ -7,7 +7,8 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.llm_cost_calc.finalization import FinalizedCost, PreparedCost, finalize_cost -from litellm.proxy.hooks.search_free_tier import SharedCounter, allocate_search_free_tier +from litellm.proxy.hooks.search_free_tier import allocate_search_free_tier +from litellm.redis.stores.search_free_tier import SearchFreeTierStore, TransportSearchFreeTierStore from litellm.types.utils import CostBreakdown, StandardLoggingPayload @@ -22,7 +23,7 @@ def prepare_response_cost(logging_obj: Logging, response: object, now: datetime) async def finalize_response_cost( logging_obj: Logging, response: object, - counter: SharedCounter | None, + counter: SearchFreeTierStore | None, now: datetime, guardrail_cost: float | None = None, ) -> FinalizedCost | None: @@ -40,7 +41,7 @@ async def finalize_response_cost( async def _finalize_response_cost( logging_obj: Logging, response: object, - counter: SharedCounter | None, + counter: SearchFreeTierStore | None, now: datetime, guardrail_cost: float | None, ) -> FinalizedCost | None: @@ -80,7 +81,7 @@ async def _finish_finalization( payload: StandardLoggingPayload, amount: float, breakdown: CostBreakdown | None, - counter: SharedCounter | None, + counter: SearchFreeTierStore | None, now: datetime, ) -> FinalizedCost: if logging_obj.cost_allowance_discount is None: @@ -102,10 +103,11 @@ async def finalize_proxy_response_cost( ) -> FinalizedCost | None: from litellm.proxy.proxy_server import proxy_logging_obj + redis_cache: Final = proxy_logging_obj.internal_usage_cache.dual_cache.redis_cache return await finalize_response_cost( logging_obj, response, - proxy_logging_obj.internal_usage_cache.dual_cache.redis_cache, + TransportSearchFreeTierStore(redis_cache.transport()) if redis_cache is not None else None, datetime.now(timezone.utc), guardrail_cost, ) diff --git a/litellm/redis/stores/search_free_tier.py b/litellm/redis/stores/search_free_tier.py new file mode 100644 index 00000000000..8fd0cf40a62 --- /dev/null +++ b/litellm/redis/stores/search_free_tier.py @@ -0,0 +1,41 @@ +from dataclasses import dataclass +from datetime import timedelta +from typing import Final, Protocol + +from litellm.redis.transport.contract import RedisFailure, RedisTransport + + +def allocation_key(provider: str, shared_group: str, period_stamp: str) -> str: + return f"litellm:search_free_tier:{provider}:{shared_group}:{period_stamp}" + + +class SearchFreeTierStore(Protocol): + async def allocate_units( + self, provider: str, shared_group: str, period_stamp: str, units: int, *, ttl: timedelta | None + ) -> float | RedisFailure: ... + + +@dataclass(frozen=True, slots=True) +class TransportSearchFreeTierStore: + transport: RedisTransport + + async def allocate_units( + self, provider: str, shared_group: str, period_stamp: str, units: int, *, ttl: timedelta | None + ) -> float | RedisFailure: + key: Final = allocation_key(provider, shared_group, period_stamp) + if ttl is None: + return await self.transport.incrbyfloat(key, units) + batch: Final = self.transport.post_call_batch() or self.transport.batch() + increment: Final = batch.incrbyfloat(key, units) + expiry: Final = batch.ttl(key) + total: Final = await increment + current_ttl: Final = await expiry + if isinstance(total, RedisFailure): + return total + if isinstance(current_ttl, RedisFailure): + return current_ttl + if current_ttl == -1: + expiration: Final = await self.transport.expire_seconds(key, ttl) + if isinstance(expiration, RedisFailure): + return expiration + return total diff --git a/tests/_support/fake_redis.py b/tests/_support/fake_redis.py index cc5fd18a305..79a6ef3eb20 100644 --- a/tests/_support/fake_redis.py +++ b/tests/_support/fake_redis.py @@ -1,7 +1,8 @@ from collections.abc import Awaitable, Callable -from typing import TypeVar +from typing import Final, TypeVar import fakeredis +from pydantic import TypeAdapter import litellm from litellm._service_logger import ServiceLogging @@ -55,3 +56,30 @@ def fake_redis_transport( request_engine=lambda: None, post_call_engine=lambda: None, ) + + +_FLOAT: Final = TypeAdapter(float) +_INT: Final = TypeAdapter(int) +_KEYS: Final = TypeAdapter(list[bytes]) + + +def set_reply_callback(cache: RedisCache, command: str, callback: Callable[..., object]) -> None: + cache.init_async_client().set_response_callback(command, callback) # pyright: ignore[reportUnknownMemberType, reportArgumentType] # redis-py types the callback for decoding only, tests also raise or await in it + + +async def legacy_increment(cache: RedisCache, key: str, value: float, *, ttl: int | None) -> float: + return _FLOAT.validate_python(await cache.async_increment(key, value, ttl=ttl)) # pyright: ignore[reportUnknownMemberType] # legacy RedisCache API is untyped + + +def legacy_default_ttl(cache: RedisCache) -> int: + return _INT.validate_python(cache.get_ttl(ttl=None)) # pyright: ignore[reportUnknownMemberType] # legacy RedisCache API is untyped + + +async def keys_matching(client: fakeredis.FakeAsyncRedis, pattern: str) -> tuple[bytes, ...]: + return tuple(_KEYS.validate_python(await client.keys(pattern))) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # fakeredis is untyped + + +def trip_breaker(cache: RedisCache, failures: int | None = None) -> None: + breaker: Final = cache._circuit_breaker # pyright: ignore[reportPrivateUsage] # tests open the breaker the way RedisCache does after hard failures + for _ in range(breaker.failure_threshold if failures is None else failures): + breaker.record_failure() diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index 8d4f8cd87ad..2d754373f4a 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -8,6 +8,7 @@ from types import SimpleNamespace from typing import Final, cast from unittest.mock import AsyncMock, MagicMock, create_autospec, patch +import fakeredis import httpx import pytest from fastapi import HTTPException @@ -40,6 +41,10 @@ from litellm.proxy.utils import ( ) from litellm.types.router import Deployment, LiteLLM_Params from litellm.types.utils import CallTypes, LiteLLMBatch, ModelResponse, PromptTokensDetailsWrapper, Usage +from tests._support.fake_redis import ( + fake_redis_cache, + keys_matching, +) @pytest.mark.asyncio @@ -536,7 +541,10 @@ async def test_async_post_call_failure_hook_carries_used_client_oauth_token_from @pytest.mark.parametrize( "metadata_buckets, expected", [ - ({"metadata": {"used_client_oauth_token": True}, "litellm_metadata": {"used_client_oauth_token": False}}, False), + ( + {"metadata": {"used_client_oauth_token": True}, "litellm_metadata": {"used_client_oauth_token": False}}, + False, + ), ({"metadata": {"used_client_oauth_token": True}, "litellm_metadata": {"user_id": "caller"}}, None), ({"metadata": {"used_client_oauth_token": "yes"}}, None), ], @@ -623,7 +631,6 @@ async def test_async_post_call_failure_hook_reads_used_client_oauth_token_from_t assert json.loads(payload["metadata"])["used_client_oauth_token"] is expected - @pytest.mark.asyncio @pytest.mark.parametrize( "request_route, body", @@ -743,7 +750,10 @@ async def test_async_post_call_failure_hook_adds_guardrail_cost_to_recovered_str [ (None, ["detect-only-scanner"]), ([], ["detect-only-scanner"]), - ([{"guardrail_name": "logged-before-the-failure", "guardrail_status": "success"}], ["logged-before-the-failure"]), + ( + [{"guardrail_name": "logged-before-the-failure", "guardrail_status": "success"}], + ["logged-before-the-failure"], + ), ], ) async def test_async_post_call_failure_hook_keeps_a_guardrail_verdict_recorded_after_the_failure_was_logged( @@ -3368,7 +3378,9 @@ async def test_track_cost_callback_failure_alert_never_carries_request_metadata_ @pytest.mark.asyncio @pytest.mark.parametrize("identity_field", ["agent_id", "billing_agent_id"]) -async def test_autonomous_llm_callback_persists_without_human_or_key(identity_field: str) -> None: # test-quality-ok: verifies anonymous-agent charges reach the persistence boundary; no injection seam +async def test_autonomous_llm_callback_persists_without_human_or_key( # test-quality-ok: verifies anonymous-agent charges reach the persistence boundary; no injection seam + identity_field: str, +) -> None: kwargs: Final = { "call_type": "acompletion", "model": "test-model", @@ -3392,9 +3404,12 @@ async def test_autonomous_llm_callback_persists_without_human_or_key(identity_fi @pytest.mark.parametrize("agent_id,expected", [(None, False), ("autonomous-agent", True)]) def test_autonomous_agent_cost_tracking_needs_no_human_or_virtual_key(agent_id: str | None, expected: bool) -> None: - assert _should_track_cost_callback( - user_api_key=None, user_id=None, team_id=None, end_user_id=None, call_type="acompletion", agent_id=agent_id - ) is expected + assert ( + _should_track_cost_callback( + user_api_key=None, user_id=None, team_id=None, end_user_id=None, call_type="acompletion", agent_id=agent_id + ) + is expected + ) _CALL_START: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) @@ -3435,14 +3450,6 @@ async def test_track_cost_callback_enqueue_emits_no_service_span(): # test-qual assert update_database.await_count == 1, "the spend enqueue itself must still run" assert success_hook.await_count == 0, [call.kwargs for call in success_hook.await_args_list] -class _SharedCounter: - def __init__(self) -> None: - self.totals: dict[str, float] = {} - - async def async_increment(self, key: str, value: float, ttl: int | None = None) -> float: - self.totals[key] = self.totals.get(key, 0) + value - return self.totals[key] - def _grounded_gemini_kwargs(response_cost: float) -> dict: return { @@ -3478,9 +3485,13 @@ def _grounded_gemini_kwargs(response_cost: float) -> dict: @pytest.mark.asyncio @pytest.mark.parametrize(("opted_in", "finalized"), [(True, False), (False, False), (True, True)]) -async def test_track_cost_callback_frees_grounding_under_the_allowance_only_when_opted_in(monkeypatch, opted_in, finalized): +async def test_track_cost_callback_frees_grounding_under_the_allowance_only_when_opted_in( + monkeypatch, opted_in, finalized +): monkeypatch.setattr(litellm, "apply_search_free_tier", opted_in, raising=False) - counter = _SharedCounter() + server: Final = fakeredis.FakeServer() + counter: Final = fake_redis_cache(server, namespace="search-callback") + raw: Final = fakeredis.FakeAsyncRedis(server=server) kwargs = _grounded_gemini_kwargs(response_cost=1.0) kwargs["standard_logging_object"]["cost_finalized"] = finalized response = ModelResponse( @@ -3520,10 +3531,12 @@ async def test_track_cost_callback_frees_grounding_under_the_allowance_only_when if opted_in and not finalized: assert kwargs["response_cost"] == pytest.approx(1.0 - 5 * 0.1) assert breakdown["original_cost"] == 1.0 and breakdown["discount_amount"] == pytest.approx(5 * 0.1) - assert sum(counter.totals.values()) == 5 + keys: Final = await keys_matching(raw, "search-callback:litellm:search_free_tier:*") + assert len(keys) == 1 + assert await raw.get(keys[0]) == b"5" else: assert kwargs["response_cost"] == 1.0 and "discount_amount" not in breakdown - assert counter.totals == {} + assert await keys_matching(raw, "*") == () written_cost = mock_proxy_logging.db_spend_update_writer.update_database.await_args.kwargs["response_cost"] assert written_cost == pytest.approx(kwargs["standard_logging_object"]["response_cost"]) @@ -3603,7 +3616,9 @@ async def test_batch_accounting_does_not_retry_after_a_search_allowance_update(m from litellm.proxy.hooks.proxy_track_cost_callback import track_batch_accounting monkeypatch.setattr(litellm, "apply_search_free_tier", True) - counter: Final = _SharedCounter() + server: Final = fakeredis.FakeServer() + counter: Final = fake_redis_cache(server, namespace="search-callback") + raw: Final = fakeredis.FakeAsyncRedis(server=server) dependency: Final = SimpleNamespace( failed_tracking_alert=AsyncMock(), internal_usage_cache=SimpleNamespace(dual_cache=SimpleNamespace(redis_cache=counter)), @@ -3621,7 +3636,9 @@ async def test_batch_accounting_does_not_retry_after_a_search_allowance_update(m with track_batch_accounting("grounded-batch") as outcome: await ProxyDBLogger().async_log_success_event(kwargs, response, _CALL_START, _CALL_START) - assert sum(counter.totals.values()) == 5 + keys: Final = await keys_matching(raw, "search-callback:litellm:search_free_tier:*") + assert len(keys) == 1 + assert await raw.get(keys[0]) == b"5" assert outcome.failed is True assert outcome.retryable_failure is False diff --git a/tests/unit/proxy/hooks/test_search_free_tier.py b/tests/unit/proxy/hooks/test_search_free_tier.py index b57908378e5..a6f9f9ea836 100644 --- a/tests/unit/proxy/hooks/test_search_free_tier.py +++ b/tests/unit/proxy/hooks/test_search_free_tier.py @@ -1,7 +1,11 @@ from datetime import datetime, timedelta, timezone from typing import Final +from unittest.mock import patch +import fakeredis import pytest +import pytest_asyncio +from redis.exceptions import ConnectionError as RedisConnectionError from litellm.proxy.hooks.search_free_tier import ( apply_search_free_tier, @@ -9,29 +13,34 @@ from litellm.proxy.hooks.search_free_tier import ( search_free_tier_key, search_free_tier_ttl, ) -from litellm.types.utils import ModelResponse, PromptTokensDetailsWrapper, Usage +from litellm.redis.stores.search_free_tier import TransportSearchFreeTierStore +from litellm.types.utils import CostBreakdown, ModelResponse, PromptTokensDetailsWrapper, StandardLoggingPayload, Usage +from tests._support.fake_redis import ( + fake_redis_cache, + keys_matching, + set_reply_callback, +) NOW: Final = datetime(2026, 10, 20, 12, 30, tzinfo=timezone.utc) RATE: Final = 0.01 -class _Counter: - def __init__(self) -> None: - self.totals: dict[str, float] = {} - self.ttls: dict[str, int | None] = {} - - async def async_increment(self, key: str, value: float, ttl: int | None = None) -> float: - self.totals[key] = self.totals.get(key, 0) + value - self.ttls[key] = ttl - return self.totals[key] +@pytest.fixture +def server() -> fakeredis.FakeServer: + return fakeredis.FakeServer() -class _BrokenCounter: - async def async_increment(self, key: str, value: float, ttl: int | None = None) -> float: - raise ConnectionError("redis is down") +@pytest.fixture +def raw(server: fakeredis.FakeServer) -> fakeredis.FakeAsyncRedis: + return fakeredis.FakeAsyncRedis(server=server) -def _model_info(billing_unit: str, count: int, period: str = "month", provider: str = "vertex_ai") -> dict: +@pytest_asyncio.fixture(loop_scope="function") +async def counter(server: fakeredis.FakeServer) -> TransportSearchFreeTierStore: + return TransportSearchFreeTierStore(fake_redis_cache(server, namespace="search-test").transport()) + + +def _model_info(billing_unit: str, count: int, period: str = "month", provider: str = "vertex_ai") -> dict[str, object]: return { "litellm_provider": provider, "mode": "chat", @@ -48,8 +57,16 @@ def _response(queries: int | None) -> ModelResponse: ) -def _call(model_info: dict, response_cost: float, cache_hit: bool | None = None) -> tuple[dict, dict]: - sl_object = { +def _breakdown(sl_object: StandardLoggingPayload) -> CostBreakdown: + breakdown: Final = sl_object["cost_breakdown"] + assert breakdown is not None + return breakdown + + +def _call( + model_info: dict[str, object], response_cost: float, cache_hit: bool | None = None +) -> tuple[dict[str, object], StandardLoggingPayload]: + sl_object: StandardLoggingPayload = { # pyright: ignore[reportAssignmentType] # the discount path reads only these keys "response_cost": response_cost, "cache_hit": cache_hit, "custom_llm_provider": model_info["litellm_provider"], @@ -93,8 +110,31 @@ def test_counter_ttl_outlives_the_period_boundary_by_under_two_hours(period, bou @pytest.mark.asyncio -async def test_per_query_allowance_is_shared_across_calls_and_split_on_the_crossing_call(): - counter = _Counter() +async def test_search_discount_works_without_lua_millisecond_or_transaction_permissions( + counter: TransportSearchFreeTierStore, raw: fakeredis.FakeAsyncRedis +) -> None: + await raw.acl_setuser( + "default", + enabled=True, + nopass=True, + reset=True, + keys=["*"], + commands=["+incrbyfloat", "+ttl", "+expire", "+get", "+pttl", "+ping", "+info", "+client"], + ) + kwargs, sl_object = _call(_model_info("per_query", count=7), response_cost=1.0) + key: Final = "search-test:litellm:search_free_tier:vertex_ai:test-group:2026-10" + + with patch("time.time", return_value=NOW.timestamp()): + cost: Final = await apply_search_free_tier(kwargs, sl_object, 1.0, _response(5), counter, NOW) + assert cost == kwargs["response_cost"] == sl_object["response_cost"] == 0.95 + assert await raw.get(key) == b"5" + assert await raw.pttl(key) == 995_400_000 + + +@pytest.mark.asyncio +async def test_per_query_allowance_is_shared_across_calls_and_split_on_the_crossing_call( + counter: TransportSearchFreeTierStore, raw: fakeredis.FakeAsyncRedis +): model_info = _model_info("per_query", count=7) reduced = [] for _ in range(3): @@ -108,13 +148,12 @@ async def test_per_query_allowance_is_shared_across_calls_and_split_on_the_cross assert c2 == k2 == s2 == pytest.approx(1.0 - 2 * RATE) assert b2["discount_amount"] == pytest.approx(2 * RATE) assert c3 == k3 == s3 == 1.0 and "discount_amount" not in b3 - assert counter.totals == {"litellm:search_free_tier:vertex_ai:test-group:2026-10": 15} - assert counter.ttls["litellm:search_free_tier:vertex_ai:test-group:2026-10"] > 0 + assert await raw.get("search-test:litellm:search_free_tier:vertex_ai:test-group:2026-10") == b"15" + assert await raw.pttl("search-test:litellm:search_free_tier:vertex_ai:test-group:2026-10") > 0 @pytest.mark.asyncio -async def test_gemini_api_and_vertex_ai_grounding_draw_on_separate_allowances(): - counter = _Counter() +async def test_gemini_api_and_vertex_ai_grounding_draw_on_separate_allowances(counter: TransportSearchFreeTierStore): gemini_api = _model_info("per_query", count=3, provider="gemini") vertex_ai = _model_info("per_query", count=3, provider="vertex_ai") kwargs, sl_object = _call(gemini_api, response_cost=1.0) @@ -128,8 +167,9 @@ async def test_gemini_api_and_vertex_ai_grounding_draw_on_separate_allowances(): @pytest.mark.asyncio -async def test_per_prompt_allowance_counts_a_grounded_prompt_once_whatever_the_query_count(): - counter = _Counter() +async def test_per_prompt_allowance_counts_a_grounded_prompt_once_whatever_the_query_count( + counter: TransportSearchFreeTierStore, raw: fakeredis.FakeAsyncRedis +): model_info = _model_info("per_prompt", count=1, period="day") kwargs, sl_object = _call(model_info, response_cost=1.0) first = await apply_search_free_tier(kwargs, sl_object, 1.0, _response(5), counter, NOW) @@ -137,12 +177,11 @@ async def test_per_prompt_allowance_counts_a_grounded_prompt_once_whatever_the_q second = await apply_search_free_tier(kwargs, sl_object, 1.0, _response(5), counter, NOW) assert first == pytest.approx(1.0 - RATE) assert second == 1.0 - assert counter.totals == {"litellm:search_free_tier:vertex_ai:test-group:2026-10-20": 2} + assert await raw.get("search-test:litellm:search_free_tier:vertex_ai:test-group:2026-10-20") == b"2" @pytest.mark.asyncio -async def test_sidecar_shaped_response_mapping_is_counted_like_a_model_response(): - counter = _Counter() +async def test_sidecar_shaped_response_mapping_is_counted_like_a_model_response(counter: TransportSearchFreeTierStore): kwargs, sl_object = _call(_model_info("per_query", count=100), response_cost=1.0) response = { "usage": { @@ -156,55 +195,112 @@ async def test_sidecar_shaped_response_mapping_is_counted_like_a_model_response( @pytest.mark.asyncio -@pytest.mark.parametrize("counter", [None, _BrokenCounter()]) -async def test_full_price_when_no_counter_is_available(counter): +@pytest.mark.parametrize("connected", [True, False]) +async def test_full_price_when_no_counter_is_available( + counter: TransportSearchFreeTierStore, server: fakeredis.FakeServer, connected: bool +): + server.connected = connected kwargs, sl_object = _call(_model_info("per_query", count=100), response_cost=1.0) - assert await apply_search_free_tier(kwargs, sl_object, 1.0, _response(5), counter, NOW) == 1.0 - assert kwargs["response_cost"] == 1.0 and "discount_amount" not in sl_object["cost_breakdown"] + assert ( + await apply_search_free_tier(kwargs, sl_object, 1.0, _response(5), None if connected else counter, NOW) == 1.0 + ) + assert kwargs["response_cost"] == 1.0 and "discount_amount" not in _breakdown(sl_object) @pytest.mark.asyncio @pytest.mark.parametrize("response", [_response(0), _response(None), ModelResponse()]) -async def test_requests_without_grounding_never_touch_the_counter(response): - counter = _Counter() +async def test_requests_without_grounding_never_touch_the_counter( + response, counter: TransportSearchFreeTierStore, raw: fakeredis.FakeAsyncRedis +): kwargs, sl_object = _call(_model_info("per_query", count=100), response_cost=1.0) assert await apply_search_free_tier(kwargs, sl_object, 1.0, response, counter, NOW) == 1.0 - assert counter.totals == {} + assert await keys_matching(raw, "*") == () @pytest.mark.asyncio -async def test_models_without_an_allowance_never_touch_the_counter(): - counter = _Counter() +async def test_models_without_an_allowance_never_touch_the_counter( + counter: TransportSearchFreeTierStore, raw: fakeredis.FakeAsyncRedis +): model_info = {k: v for k, v in _model_info("per_query", count=100).items() if k != "search_free_requests"} kwargs, sl_object = _call(model_info, response_cost=1.0) assert await apply_search_free_tier(kwargs, sl_object, 1.0, _response(5), counter, NOW) == 1.0 - assert counter.totals == {} + assert await keys_matching(raw, "*") == () @pytest.mark.asyncio -async def test_a_second_callback_run_on_the_same_request_counts_it_once(): - counter = _Counter() +async def test_a_second_callback_run_on_the_same_request_counts_it_once( + counter: TransportSearchFreeTierStore, raw: fakeredis.FakeAsyncRedis +): kwargs, sl_object = _call(_model_info("per_query", count=100), response_cost=1.0) first = await apply_search_free_tier(kwargs, sl_object, 1.0, _response(5), counter, NOW) second = await apply_search_free_tier(kwargs, sl_object, first, _response(5), counter, NOW) assert first == second == pytest.approx(1.0 - 5 * RATE) - assert counter.totals == {"litellm:search_free_tier:vertex_ai:test-group:2026-10": 5} + assert await raw.get("search-test:litellm:search_free_tier:vertex_ai:test-group:2026-10") == b"5" @pytest.mark.asyncio -async def test_a_cache_hit_with_grounded_usage_is_not_counted_and_stays_at_zero(): - counter = _Counter() +async def test_lost_increment_reply_keeps_full_price_and_only_recounts_on_an_explicit_callback_retry( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + cache: Final = fake_redis_cache(server, namespace="search-test") + store: Final = TransportSearchFreeTierStore(cache.transport()) + kwargs, sl_object = _call(_model_info("per_query", count=7), response_cost=1.0) + + def lost_reply(_: object) -> None: + raise RedisConnectionError("reply lost after allocation") + + set_reply_callback(cache, "INCRBYFLOAT", lost_reply) + assert await apply_search_free_tier(kwargs, sl_object, 1.0, _response(5), store, NOW) == 1.0 + assert kwargs["response_cost"] == sl_object["response_cost"] == 1.0 + assert not kwargs.get("search_free_tier_counted", False) + assert await raw.get("search-test:litellm:search_free_tier:vertex_ai:test-group:2026-10") == b"5" + + set_reply_callback(cache, "INCRBYFLOAT", float) + assert await apply_search_free_tier(kwargs, sl_object, 1.0, _response(5), store, NOW) == pytest.approx(0.98) + assert kwargs["search_free_tier_counted"] is True + assert await raw.get("search-test:litellm:search_free_tier:vertex_ai:test-group:2026-10") == b"10" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("period", "stamp", "boundary"), + [ + ("month", "2026-10", datetime(2026, 11, 1, tzinfo=timezone.utc)), + ("day", "2026-10-20", datetime(2026, 10, 21, tzinfo=timezone.utc)), + ], +) +async def test_caller_places_expiry_at_period_boundary_plus_one_hour( + server: fakeredis.FakeServer, + raw: fakeredis.FakeAsyncRedis, + period: str, + stamp: str, + boundary: datetime, +) -> None: + store: Final = TransportSearchFreeTierStore(fake_redis_cache(server, namespace="search-test").transport()) + kwargs, sl_object = _call(_model_info("per_query", count=7, period=period), response_cost=1.0) + with patch("time.time", return_value=NOW.timestamp()): + assert await apply_search_free_tier(kwargs, sl_object, 1.0, _response(5), store, NOW) == pytest.approx(0.95) # pyright: ignore[reportUnknownMemberType] # pytest.approx is untyped + assert await raw.pttl(f"search-test:litellm:search_free_tier:vertex_ai:test-group:{stamp}") == ( + (boundary - NOW + timedelta(hours=1)) // timedelta(milliseconds=1) + ) + + +@pytest.mark.asyncio +async def test_a_cache_hit_with_grounded_usage_is_not_counted_and_stays_at_zero( + counter: TransportSearchFreeTierStore, raw: fakeredis.FakeAsyncRedis +): kwargs, sl_object = _call(_model_info("per_query", count=100), response_cost=0.0, cache_hit=True) assert await apply_search_free_tier(kwargs, sl_object, 0.0, _response(2), counter, NOW) == 0.0 - assert counter.totals == {} - assert "discount_amount" not in sl_object["cost_breakdown"] + assert await keys_matching(raw, "*") == () + assert "discount_amount" not in _breakdown(sl_object) @pytest.mark.asyncio -async def test_freed_cost_never_exceeds_the_billed_cost(): - counter = _Counter() +async def test_freed_cost_never_exceeds_the_billed_cost( + counter: TransportSearchFreeTierStore, raw: fakeredis.FakeAsyncRedis +): billed = 3 * RATE kwargs, sl_object = _call(_model_info("per_query", count=100), response_cost=billed) assert await apply_search_free_tier(kwargs, sl_object, billed, _response(5), counter, NOW) == 0.0 - assert sl_object["cost_breakdown"]["discount_amount"] == pytest.approx(billed) - assert counter.totals == {"litellm:search_free_tier:vertex_ai:test-group:2026-10": 5} + assert _breakdown(sl_object).get("discount_amount") == pytest.approx(billed) # pyright: ignore[reportUnknownMemberType] # pytest.approx is untyped + assert await raw.get("search-test:litellm:search_free_tier:vertex_ai:test-group:2026-10") == b"5" diff --git a/tests/unit/proxy/pass_through_endpoints/test_success_handler.py b/tests/unit/proxy/pass_through_endpoints/test_success_handler.py index c721e808ed5..ea881d5cac2 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_success_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/test_success_handler.py @@ -659,6 +659,8 @@ async def test_completed_passthrough_cancellation_preserves_accounting( from unittest.mock import create_autospec import anyio + import fakeredis + from redis.asyncio.client import Pipeline from litellm.caching.dual_cache import DualCache from litellm.proxy import proxy_server @@ -667,20 +669,29 @@ async def test_completed_passthrough_cancellation_preserves_accounting( from litellm.proxy.spend_tracking.budget_reservation import release_unbound_budget_reservation from litellm.proxy.utils import ProxyLogging from litellm.types.utils import ModelResponse, PromptTokensDetailsWrapper, Usage + from tests._support.fake_redis import fake_redis_cache - class AllowanceCounter: - def __init__(self) -> None: - self.started = asyncio.Event() - self.release = asyncio.Event() - self.calls = 0 - self.units = 0.0 + allowance_started: Final = asyncio.Event() + allowance_release: Final = asyncio.Event() + executions: Final[list[int]] = [] # mutable-ok: counts pipeline flushes observed by the gated fake client - async def async_increment(self, key: str, value: float, ttl: int | None = None) -> float: - self.started.set() - await self.release.wait() - self.calls += 1 - self.units += value - return self.units + class GatedClient(fakeredis.FakeAsyncRedis): + def pipeline(self, transaction: bool = True, shard_hint: str | None = None) -> "Pipeline[bytes]": + pipe: Final = super().pipeline(transaction=transaction, shard_hint=shard_hint) # pyright: ignore[reportUnknownVariableType] # fakeredis pipeline is untyped + native_execute: Final = pipe.execute # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # redis-py pipeline.execute reply is untyped + + async def gated_execute(raise_on_error: bool = True) -> list[object]: + allowance_started.set() + await allowance_release.wait() + executions.append(1) + return await native_execute(raise_on_error=raise_on_error) # pyright: ignore[reportUnknownVariableType] # redis-py erases the reply type + + monkeypatch.setattr(pipe, "execute", gated_execute) # pyright: ignore[reportUnknownArgumentType] # fakeredis pipeline is untyped + return pipe # pyright: ignore[reportUnknownVariableType] # fakeredis pipeline is untyped + + server: Final = fakeredis.FakeServer() + allowance: Final = fake_redis_cache(server, client=GatedClient) # pyright: ignore[reportUnknownArgumentType] # fakeredis client class is partially untyped + verifier: Final = fakeredis.FakeAsyncRedis(server=server) now: Final = datetime(2026, 10, 20, tzinfo=timezone.utc) model: Final = "passthrough-cancellation-test" @@ -714,7 +725,6 @@ async def test_completed_passthrough_cancellation_preserves_accounting( writer: Final = create_autospec(DBSpendUpdateWriter, instance=True) proxy_logging: Final = ProxyLogging(user_api_key_cache=DualCache()) proxy_logging.db_spend_update_writer = writer - allowance: Final = AllowanceCounter() monkeypatch.setattr(proxy_logging.internal_usage_cache.dual_cache, "redis_cache", allowance) monkeypatch.setattr(proxy_server, "proxy_logging_obj", proxy_logging) logging_obj: Final = LiteLLMLoggingObj( @@ -771,7 +781,7 @@ async def test_completed_passthrough_cancellation_preserves_accounting( pending: Final = asyncio.create_task(dispatch()) await entered.wait() if cancel_phase == "allocation": - await allowance.started.wait() + await allowance_started.wait() if cancel_mode == "scope": scope.cancel() elif cancel_mode != "none": @@ -781,7 +791,7 @@ async def test_completed_passthrough_cancellation_preserves_accounting( pending.cancel() if cancel_phase == "lock": logging_obj.cost_finalization_lock.release() - allowance.release.set() + allowance_release.set() if cancel_mode in ("once", "twice"): with pytest.raises(asyncio.CancelledError): await pending @@ -791,8 +801,9 @@ async def test_completed_passthrough_cancellation_preserves_accounting( assert reservation["callback_bound"] is True await release_unbound_budget_reservation(reservation) await GLOBAL_LOGGING_WORKER.flush() - assert allowance.calls == 1 - assert allowance.units == 3 + assert executions == [1] + (allocation_key,) = await verifier.keys("*") # pyright: ignore[reportUnknownVariableType] # fakeredis replies are untyped + assert await verifier.get(allocation_key) == b"3" # pyright: ignore[reportUnknownArgumentType] # fakeredis replies are untyped assert writer.update_database.await_count == 1 recorded: Final = writer.update_database.call_args.kwargs assert recorded["response_cost"] == pytest.approx(0.48) diff --git a/tests/unit/proxy/spend_tracking/test_cost_finalization.py b/tests/unit/proxy/spend_tracking/test_cost_finalization.py index aefc4a96a28..47eb37c414f 100644 --- a/tests/unit/proxy/spend_tracking/test_cost_finalization.py +++ b/tests/unit/proxy/spend_tracking/test_cost_finalization.py @@ -1,44 +1,61 @@ import asyncio +from dataclasses import dataclass from datetime import datetime, timezone from typing import Final +import fakeredis import pytest +import pytest_asyncio +from redis.exceptions import ConnectionError as RedisConnectionError import litellm +from litellm.caching.redis_cache import RedisCache from litellm.litellm_core_utils.litellm_logging import Logging from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.proxy.spend_tracking.cost_finalization import finalize_response_cost +from litellm.redis.stores.search_free_tier import TransportSearchFreeTierStore from litellm.types.utils import ModelResponse, PromptTokensDetailsWrapper, Usage +from tests._support.fake_redis import ( + fake_redis_cache, + set_reply_callback, +) NOW: Final = datetime(2026, 10, 20, 12, 30, tzinfo=timezone.utc) MODEL: Final = "gemini-finalization-test" -class Counter: - def __init__(self, total: float = 0.0, fail: bool = False) -> None: - self.total = total - self.calls = 0 - self.fail = fail - - async def async_increment(self, key: str, value: float, ttl: int | None = None) -> float: - self.calls += 1 - if self.fail: - raise ConnectionError("counter unavailable") - self.total += value - return self.total +_ALLOCATION_KEY: Final = "ns:litellm:search_free_tier:vertex_ai:finalization-test:2026-10" -class BlockingCounter(Counter): - def __init__(self) -> None: - super().__init__() - self.started = asyncio.Event() - self.release = asyncio.Event() +@dataclass(frozen=True, slots=True) +class AllocationRedis: + server: fakeredis.FakeServer + raw: fakeredis.FakeAsyncRedis + cache: RedisCache + store: TransportSearchFreeTierStore - async def async_increment(self, key: str, value: float, ttl: int | None = None) -> float: - self.started.set() - await self.release.wait() - return await super().async_increment(key, value, ttl) + +@pytest_asyncio.fixture(loop_scope="function") +async def allocation_redis() -> AllocationRedis: + server: Final = fakeredis.FakeServer() + cache: Final = fake_redis_cache(server, namespace="ns") + return AllocationRedis( + server, fakeredis.FakeAsyncRedis(server=server), cache, TransportSearchFreeTierStore(cache.transport()) + ) + + +def block_increment(cache: RedisCache) -> tuple[asyncio.Event, asyncio.Event]: + started: Final = asyncio.Event() + release: Final = asyncio.Event() + + async def receive(value: bytes) -> float: + started.set() + await release.wait() + return float(value) + + set_reply_callback(cache, "INCRBYFLOAT", receive) + return started, release @pytest.fixture @@ -89,15 +106,16 @@ def response(cost: float = 0.5) -> ModelResponse: @pytest.mark.asyncio @pytest.mark.parametrize(("used", "expected"), [(0, 0.48), (1, 0.49), (2, 0.5)]) async def test_header_and_logging_share_one_allowance_allocation( - logging_obj: Logging, used: int, expected: float + logging_obj: Logging, allocation_redis: AllocationRedis, used: int, expected: float ) -> None: - counter: Final = Counter(used) + await allocation_redis.raw.set(_ALLOCATION_KEY, used) + counter: Final = allocation_redis.store result: Final = response() costs: Final = await asyncio.gather(*(finalize_response_cost(logging_obj, result, counter, NOW) for _ in range(3))) assert costs[0] is costs[1] is costs[2] assert costs[0] is not None assert costs[0].amount == pytest.approx(expected) - assert counter.calls == 1 + assert await allocation_redis.raw.get(_ALLOCATION_KEY) == str(used + 3).encode() logging_obj._success_handler_helper_fn(result=result, start_time=NOW, end_time=NOW) payload: Final = logging_obj.model_call_details["standard_logging_object"] headers: Final = ProxyBaseLLMRequestProcessing.get_custom_headers( @@ -109,19 +127,47 @@ async def test_header_and_logging_share_one_allowance_allocation( @pytest.mark.asyncio -async def test_failed_counter_freezes_full_price_without_a_retry(logging_obj: Logging) -> None: - counter: Final = Counter(fail=True) +async def test_failed_counter_freezes_full_price_without_a_retry( + logging_obj: Logging, allocation_redis: AllocationRedis +) -> None: + allocation_redis.server.connected = False + counter: Final = allocation_redis.store result: Final = response() finalized: Final = await finalize_response_cost(logging_obj, result, counter, NOW) assert finalized is not None and finalized.amount == 0.5 + allocation_redis.server.connected = True assert await finalize_response_cost(logging_obj, result, counter, NOW) is finalized - assert counter.calls == 1 + assert await allocation_redis.raw.get(_ALLOCATION_KEY) is None + assert logging_obj.cost_allowance_discount == 0.0 @pytest.mark.asyncio -async def test_guardrail_cost_is_added_once_after_allowance(logging_obj: Logging) -> None: +async def test_lost_allocation_reply_freezes_full_price_after_allowance_was_consumed( + logging_obj: Logging, allocation_redis: AllocationRedis +) -> None: + def lost_reply(_: object) -> None: + raise RedisConnectionError("allocation response lost") + + set_reply_callback(allocation_redis.cache, "INCRBYFLOAT", lost_reply) result: Final = response() - finalized: Final = await finalize_response_cost(logging_obj, result, Counter(), NOW, guardrail_cost=0.2) + finalized: Final = await finalize_response_cost(logging_obj, result, allocation_redis.store, NOW) + assert finalized is not None and finalized.amount == 0.5 + assert logging_obj.cost_allowance_discount == 0.0 + assert await allocation_redis.raw.get(_ALLOCATION_KEY) == b"3" + + set_reply_callback(allocation_redis.cache, "INCRBYFLOAT", float) + assert await finalize_response_cost(logging_obj, result, allocation_redis.store, NOW) is finalized + assert await allocation_redis.raw.get(_ALLOCATION_KEY) == b"3" + + +@pytest.mark.asyncio +async def test_guardrail_cost_is_added_once_after_allowance( + logging_obj: Logging, allocation_redis: AllocationRedis +) -> None: + result: Final = response() + finalized: Final = await finalize_response_cost( + logging_obj, result, allocation_redis.store, NOW, guardrail_cost=0.2 + ) assert finalized is not None and finalized.amount == pytest.approx(0.68) logging_obj._success_handler_helper_fn(result=result, start_time=NOW, end_time=NOW) payload: Final = logging_obj.model_call_details["standard_logging_object"] @@ -130,18 +176,22 @@ async def test_guardrail_cost_is_added_once_after_allowance(logging_obj: Logging @pytest.mark.asyncio -async def test_cache_hit_finalizes_zero_without_allocating(logging_obj: Logging) -> None: +async def test_cache_hit_finalizes_zero_without_allocating( + logging_obj: Logging, allocation_redis: AllocationRedis +) -> None: result: Final = response() result._hidden_params["cache_hit"] = True result._hidden_params["response_cost"] = 0.0 - counter: Final = Counter() + counter: Final = allocation_redis.store finalized: Final = await finalize_response_cost(logging_obj, result, counter, NOW) assert finalized is not None and finalized.amount == 0.0 - assert counter.calls == 0 + assert await allocation_redis.raw.get(_ALLOCATION_KEY) is None @pytest.mark.asyncio -async def test_both_callback_paths_wait_for_finalization(logging_obj: Logging) -> None: +async def test_both_callback_paths_wait_for_finalization( + logging_obj: Logging, allocation_redis: AllocationRedis +) -> None: from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.utils import dispatch_success_logging @@ -169,18 +219,19 @@ async def test_both_callback_paths_wait_for_finalization(logging_obj: Logging) - assert not logging_obj._async_success_enqueued assert logging_obj.deferred_success_result is not None assert logging_obj.deferred_success_result[0] is inner - counter: Final = BlockingCounter() + counter: Final = allocation_redis.store + started, release = block_increment(allocation_redis.cache) pending: Final = asyncio.create_task(finalize_response_cost(logging_obj, outer, counter, NOW)) - await counter.started.wait() + await started.wait() dispatch_success_logging(logging_obj, outer, NOW, NOW, False, False) assert not logging_obj._sync_success_scheduled assert not logging_obj._async_success_enqueued concurrent: Final = asyncio.create_task(finalize_response_cost(logging_obj, outer, counter, NOW)) - counter.release.set() + release.set() finalized: Final = await pending assert finalized is not None assert await concurrent is finalized - assert counter.calls == 1 + assert await allocation_redis.raw.get(_ALLOCATION_KEY) == b"3" ProxyBaseLLMRequestProcessing._flush_deferred_async_logging(logging_obj, exception_raised=False) observed: Final = await asyncio.wait_for(asyncio.gather(sync_cost, async_cost), timeout=10) await GLOBAL_LOGGING_WORKER.flush() @@ -201,12 +252,14 @@ async def test_rejected_response_does_not_release_success_callbacks(logging_obj: @pytest.mark.asyncio -async def test_prebuilt_payload_is_reconciled_with_final_cost(logging_obj: Logging) -> None: +async def test_prebuilt_payload_is_reconciled_with_final_cost( + logging_obj: Logging, allocation_redis: AllocationRedis +) -> None: result: Final = response() prepared: Final = logging_obj.prepare_success_logging(result, NOW, NOW, None, None, True, False) original: Final = prepared[3] assert original is not None and original["response_cost"] == 0.5 - finalized: Final = await finalize_response_cost(logging_obj, result, Counter(), NOW) + finalized: Final = await finalize_response_cost(logging_obj, result, allocation_redis.store, NOW) assert finalized is not None logging_obj._success_handler_helper_fn(result, NOW, NOW, standard_logging_object=original) payload: Final = logging_obj.model_call_details["standard_logging_object"] @@ -215,7 +268,9 @@ async def test_prebuilt_payload_is_reconciled_with_final_cost(logging_obj: Loggi @pytest.mark.asyncio -async def test_cached_response_waits_for_guardrail_cost(logging_obj: Logging) -> None: +async def test_cached_response_waits_for_guardrail_cost( + logging_obj: Logging, allocation_redis: AllocationRedis +) -> None: from litellm.caching.caching_handler import LLMCachingHandler result: Final = response(0.0) @@ -225,17 +280,19 @@ async def test_cached_response_waits_for_guardrail_cost(logging_obj: Logging) -> handler._async_log_cache_hit_on_callbacks(logging_obj, result, NOW, NOW, True) assert not logging_obj._sync_success_scheduled assert not logging_obj._async_success_enqueued - counter: Final = Counter() + counter: Final = allocation_redis.store finalized: Final = await finalize_response_cost(logging_obj, result, counter, NOW, guardrail_cost=0.2) assert finalized is not None and finalized.amount == 0.2 - assert counter.calls == 0 + assert await allocation_redis.raw.get(_ALLOCATION_KEY) is None logging_obj._success_handler_helper_fn(result, NOW, NOW) assert logging_obj.model_call_details["standard_logging_object"]["response_cost"] == 0.2 ProxyBaseLLMRequestProcessing._flush_deferred_async_logging(logging_obj, exception_raised=True) @pytest.mark.asyncio -async def test_collector_event_preserves_finalized_cost(logging_obj: Logging) -> None: +async def test_collector_event_preserves_finalized_cost( + logging_obj: Logging, allocation_redis: AllocationRedis +) -> None: from litellm.proxy.spend_tracking.spend_event import ( SpendEvent, build_spend_event, @@ -244,7 +301,7 @@ async def test_collector_event_preserves_finalized_cost(logging_obj: Logging) -> ) result: Final = response() - await finalize_response_cost(logging_obj, result, Counter(), NOW) + await finalize_response_cost(logging_obj, result, allocation_redis.store, NOW) logging_obj._success_handler_helper_fn(result, NOW, NOW) event: Final = build_spend_event(logging_obj.model_call_details, result, NOW, NOW, store_bodies=False) assert isinstance(event, bytes) @@ -259,30 +316,33 @@ async def test_collector_event_preserves_finalized_cost(logging_obj: Logging) -> @pytest.mark.asyncio @pytest.mark.parametrize("cancel_count", [1, 2]) async def test_cancellation_finishes_an_inflight_allocation_before_releasing_callbacks( - logging_obj: Logging, cancel_count: int + logging_obj: Logging, allocation_redis: AllocationRedis, cancel_count: int ) -> None: - counter: Final = BlockingCounter() + counter: Final = allocation_redis.store + started, release = block_increment(allocation_redis.cache) result: Final = response() pending: Final = asyncio.create_task(finalize_response_cost(logging_obj, result, counter, NOW)) - await counter.started.wait() + await started.wait() for _ in range(cancel_count): pending.cancel() await asyncio.sleep(0) assert not pending.done() - counter.release.set() + release.set() with pytest.raises(asyncio.CancelledError): await pending assert logging_obj.finalized_cost is not None assert logging_obj.finalized_cost.amount == pytest.approx(0.48) assert await finalize_response_cost(logging_obj, result, counter, NOW) is logging_obj.finalized_cost - assert counter.calls == 1 + assert await allocation_redis.raw.get(_ALLOCATION_KEY) == b"3" @pytest.mark.asyncio -async def test_existing_calculator_prices_usage_before_finalizing(logging_obj: Logging) -> None: +async def test_existing_calculator_prices_usage_before_finalizing( + logging_obj: Logging, allocation_redis: AllocationRedis +) -> None: result: Final = response() result._hidden_params.pop("response_cost") - finalized: Final = await finalize_response_cost(logging_obj, result, Counter(), NOW) + finalized: Final = await finalize_response_cost(logging_obj, result, allocation_redis.store, NOW) assert finalized is not None and finalized.amount == pytest.approx(0.21) logging_obj._success_handler_helper_fn(result, NOW, NOW) assert logging_obj.model_call_details["standard_logging_object"]["response_cost"] == pytest.approx(0.21) @@ -371,20 +431,22 @@ async def test_malformed_passthrough_logging_preserves_response_dispatch_contrac @pytest.mark.asyncio async def test_finalization_failure_leaves_the_request_and_callback_billing_intact( - logging_obj: Logging, monkeypatch: pytest.MonkeyPatch + logging_obj: Logging, allocation_redis: AllocationRedis, monkeypatch: pytest.MonkeyPatch ) -> None: def explode(**kwargs: object) -> object: raise RuntimeError("payload build failed") monkeypatch.setattr(logging_obj, "prepare_success_logging", explode) - counter: Final = Counter() + counter: Final = allocation_redis.store assert await finalize_response_cost(logging_obj, response(), counter, NOW) is None assert logging_obj.finalized_cost is None - assert counter.calls == 0 + assert await allocation_redis.raw.get(_ALLOCATION_KEY) is None @pytest.mark.asyncio -async def test_deferred_cache_hit_finalizes_zero_for_results_without_hidden_params(logging_obj: Logging) -> None: +async def test_deferred_cache_hit_finalizes_zero_for_results_without_hidden_params( + logging_obj: Logging, allocation_redis: AllocationRedis +) -> None: from litellm import completion from litellm.caching.caching_handler import LLMCachingHandler @@ -395,10 +457,10 @@ async def test_deferred_cache_hit_finalizes_zero_for_results_without_hidden_para )._async_log_cache_hit_on_callbacks( logging_obj=logging_obj, cached_result=cached_result, start_time=NOW, end_time=NOW, cache_hit=True ) - counter: Final = Counter() + counter: Final = allocation_redis.store finalized: Final = await finalize_response_cost(logging_obj, cached_result, counter, NOW) assert finalized is not None and finalized.amount == 0.0 - assert counter.calls == 0 + assert await allocation_redis.raw.get(_ALLOCATION_KEY) is None logging_obj._success_handler_helper_fn(cached_result, NOW, NOW) payload: Final = logging_obj.model_call_details["standard_logging_object"] assert payload["cache_hit"] is True @@ -423,18 +485,17 @@ async def test_completed_call_keeps_its_charge_when_post_call_processing_fails( disable_spend_updates: bool, ) -> None: from unittest.mock import create_autospec + import anyio - from fastapi import Request, Response - from litellm.proxy._types import ProxyException - from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger + from litellm.proxy import proxy_server + from litellm.proxy._types import ProxyException from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger from litellm.proxy.proxy_server import ProxyConfig - from litellm.proxy import proxy_server from litellm.proxy.utils import ProxyLogging scope: Final = anyio.CancelScope() @@ -545,7 +606,7 @@ async def test_completed_call_keeps_its_charge_when_post_call_processing_fails( @pytest.mark.asyncio async def test_redacted_completed_call_preserves_usage_and_search_allowance( - logging_obj: Logging, monkeypatch: pytest.MonkeyPatch + logging_obj: Logging, allocation_redis: AllocationRedis, monkeypatch: pytest.MonkeyPatch ) -> None: from litellm.proxy.spend_tracking.cost_finalization import prepare_response_cost @@ -554,11 +615,11 @@ async def test_redacted_completed_call_preserves_usage_and_search_allowance( prepared: Final = prepare_response_cost(logging_obj, result, NOW) assert prepared is not None result.usage = Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0) - counter: Final = Counter() + counter: Final = allocation_redis.store finalized: Final = await finalize_response_cost(logging_obj, result, counter, NOW) assert finalized is not None assert finalized.amount == pytest.approx(0.48) - assert counter.total == 3 + assert await allocation_redis.raw.get(_ALLOCATION_KEY) == b"3" logging_obj._failure_handler_helper_fn(RuntimeError("post-call failed"), "", NOW, NOW) payload: Final = logging_obj.model_call_details["standard_logging_object"] assert (payload["prompt_tokens"], payload["completion_tokens"]) == (10, 5) @@ -568,9 +629,10 @@ async def test_redacted_completed_call_preserves_usage_and_search_allowance( @pytest.mark.asyncio async def test_guardrail_json_rejection_retains_completed_provider_usage( - logging_obj: Logging, monkeypatch: pytest.MonkeyPatch + logging_obj: Logging, allocation_redis: AllocationRedis, monkeypatch: pytest.MonkeyPatch ) -> None: import httpx + from litellm.caching.dual_cache import DualCache from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.proxy.utils import ProxyLogging @@ -599,7 +661,7 @@ async def test_guardrail_json_rejection_retains_completed_provider_usage( ) assert logging_obj.prepared_cost is not None normalized.usage = Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0) - finalized: Final = await finalize_response_cost(logging_obj, upstream, Counter(), NOW) + finalized: Final = await finalize_response_cost(logging_obj, upstream, allocation_redis.store, NOW) assert finalized is not None and finalized.amount == 0.5 logging_obj._failure_handler_helper_fn(RuntimeError("blocked JSON response"), "", NOW, NOW) payload: Final = logging_obj.model_call_details["standard_logging_object"] diff --git a/tests/unit/redis/stores/test_search_free_tier.py b/tests/unit/redis/stores/test_search_free_tier.py new file mode 100644 index 00000000000..993e6d45fd1 --- /dev/null +++ b/tests/unit/redis/stores/test_search_free_tier.py @@ -0,0 +1,233 @@ +from collections.abc import Iterator +from datetime import timedelta +from typing import Final +from unittest.mock import patch + +import fakeredis +import pytest +from redis.exceptions import ConnectionError as RedisConnectionError +from redis.exceptions import ResponseError + +from litellm.caching.redis_batch import request_redis_batch_scope +from litellm.redis.stores.search_free_tier import TransportSearchFreeTierStore +from litellm.redis.transport.contract import RedisFailure +from tests._support.fake_redis import ( + fake_redis_cache, + keys_matching, + legacy_default_ttl, + legacy_increment, + set_reply_callback, +) + +_KEY: Final = "ns:litellm:search_free_tier:vertex_ai:shared:2026-10" +_LOGICAL_KEY: Final = "litellm:search_free_tier:vertex_ai:shared:2026-10" + + +def _not_a_number(_: object) -> bytes: + return b"not-a-number" + + +@pytest.fixture(autouse=True) +def frozen_clock() -> Iterator[None]: + with patch("time.time", return_value=1792497600.0): + yield + + +@pytest.mark.asyncio +async def test_old_and_new_writers_share_numeric_bytes_and_do_not_refresh_expiry( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + cache: Final = fake_redis_cache(server, namespace="ns") + store: Final = TransportSearchFreeTierStore(cache.transport()) + + assert await legacy_increment(cache, _LOGICAL_KEY, 2.5, ttl=600) == 2.5 + assert await raw.get(_KEY) == b"2.5" + assert await raw.pttl(_KEY) == 600_000 + assert await store.allocate_units("vertex_ai", "shared", "2026-10", 3, ttl=timedelta(seconds=900)) == 5.5 + assert await raw.get(_KEY) == b"5.5" + assert await raw.pttl(_KEY) == 600_000 + assert await legacy_increment(cache, _LOGICAL_KEY, 1.5, ttl=1200) == 7.0 + assert await raw.get(_KEY) == b"7" + assert await raw.pttl(_KEY) == 600_000 + assert await keys_matching(raw, "*") == (_KEY.encode(),) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ttl", [None, timedelta(seconds=60)]) +async def test_new_writer_is_readable_by_legacy_increment( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis, ttl: timedelta | None +) -> None: + cache: Final = fake_redis_cache(server, namespace="ns") + store: Final = TransportSearchFreeTierStore(cache.transport()) + + assert await store.allocate_units("vertex_ai", "shared", "2026-10", 4, ttl=ttl) == 4.0 + assert await raw.get(_KEY) == b"4" + assert await raw.pttl(_KEY) == (-1 if ttl is None else 60_000) + assert await legacy_increment(cache, _LOGICAL_KEY, 2.5, ttl=None) == 6.5 + assert await raw.get(_KEY) == b"6.5" + assert await raw.pttl(_KEY) == (legacy_default_ttl(cache) * 1000 if ttl is None else 60_000) + + +@pytest.mark.asyncio +async def test_provider_group_and_period_have_independent_allowance_counters( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + store: Final = TransportSearchFreeTierStore(fake_redis_cache(server, namespace="ns").transport()) + identities: Final = ( + ("vertex_ai", "shared", "2026-10"), + ("gemini", "shared", "2026-10"), + ("vertex_ai", "other", "2026-10"), + ("vertex_ai", "shared", "2026-11"), + ("vertex_ai", "shared", "2026-10-20"), + ("vertex_ai", "shared 雪", "2026-10"), + ) + for provider, group, period in identities: + assert await store.allocate_units(provider, group, period, 2, ttl=None) == 2.0 + assert sorted(await keys_matching(raw, "*")) == sorted( + f"ns:litellm:search_free_tier:{provider}:{group}:{period}".encode() for provider, group, period in identities + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("seconds", [-1, 0, 300]) +async def test_unexpiring_legacy_counter_gets_seconds_based_expiry( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis, seconds: int +) -> None: + store: Final = TransportSearchFreeTierStore(fake_redis_cache(server, namespace="ns").transport()) + legacy: Final = fake_redis_cache(server, namespace="legacy") + legacy_key: Final = f"legacy:{_LOGICAL_KEY}" + await raw.set(_KEY, b"5") + await raw.set(legacy_key, b"5") + assert await legacy_increment(legacy, _LOGICAL_KEY, 2, ttl=seconds) == 7.0 + assert await store.allocate_units("vertex_ai", "shared", "2026-10", 2, ttl=timedelta(seconds=seconds)) == 7.0 + assert await raw.get(_KEY) == await raw.get(legacy_key) + assert await raw.pttl(_KEY) == await raw.pttl(legacy_key) + + +@pytest.mark.asyncio +async def test_expiry_remains_a_separate_write_after_increment( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + cache: Final = fake_redis_cache(server, namespace="ns") + store: Final = TransportSearchFreeTierStore(cache.transport()) + + async def ttl_reply(value: int) -> int: + assert value == -1 + assert await raw.get(_KEY) == b"3" + assert await raw.pttl(_KEY) == -1 + await raw.pexpire(_KEY, 900_000) + return value + + set_reply_callback(cache, "TTL", ttl_reply) + assert await store.allocate_units("vertex_ai", "shared", "2026-10", 3, ttl=timedelta(seconds=60)) == 3.0 + assert await raw.pttl(_KEY) == 60_000 + + +@pytest.mark.asyncio +async def test_counter_disappearing_before_expire_keeps_successful_allocation_result( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + cache: Final = fake_redis_cache(server, namespace="ns") + store: Final = TransportSearchFreeTierStore(cache.transport()) + + async def ttl_reply(value: int) -> int: + await raw.delete(_KEY) + return value + + set_reply_callback(cache, "TTL", ttl_reply) + assert await store.allocate_units("vertex_ai", "shared", "2026-10", 3, ttl=timedelta(seconds=60)) == 3.0 + assert await raw.get(_KEY) is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("command,remaining_ttl", [("INCRBYFLOAT", -1), ("TTL", -1), ("EXPIRE", 60_000)]) +async def test_malformed_replies_are_failures_after_the_counter_was_consumed( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis, command: str, remaining_ttl: int +) -> None: + cache: Final = fake_redis_cache(server, namespace="ns") + store: Final = TransportSearchFreeTierStore(cache.transport()) + set_reply_callback(cache, command, _not_a_number) + + outcome: Final = await store.allocate_units("vertex_ai", "shared", "2026-10", 3, ttl=timedelta(seconds=60)) + assert isinstance(outcome, RedisFailure) and outcome.kind == "decode" + assert await raw.get(_KEY) == b"3" + assert await raw.pttl(_KEY) == remaining_ttl + + +@pytest.mark.asyncio +@pytest.mark.parametrize("command", ["INCRBYFLOAT", "TTL", "EXPIRE"]) +async def test_lost_reply_reports_unknown_applied_write_without_retry( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis, command: str +) -> None: + cache: Final = fake_redis_cache(server, namespace="ns") + store: Final = TransportSearchFreeTierStore(cache.transport()) + + def lose_reply(_: object) -> None: + raise RedisConnectionError("allocation reply lost") + + set_reply_callback(cache, command, lose_reply) + outcome: Final = await store.allocate_units("vertex_ai", "shared", "2026-10", 3, ttl=timedelta(seconds=60)) + assert isinstance(outcome, RedisFailure) and outcome.kind == "transport" + assert await raw.get(_KEY) == b"3" + assert await raw.pttl(_KEY) == (60_000 if command == "EXPIRE" else -1) + + +@pytest.mark.asyncio +async def test_corrupt_counter_is_not_reset_or_expired( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + store: Final = TransportSearchFreeTierStore(fake_redis_cache(server, namespace="ns").transport()) + await raw.set(_KEY, b"corrupt") + outcome: Final = await store.allocate_units("vertex_ai", "shared", "2026-10", 3, ttl=timedelta(seconds=60)) + assert isinstance(outcome, RedisFailure) and outcome.kind == "server" + assert await raw.get(_KEY) == b"corrupt" + assert await raw.pttl(_KEY) == -1 + + +@pytest.mark.asyncio +async def test_outage_opens_existing_cache_breaker_and_rejects_writes_after_reconnection( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + store: Final = TransportSearchFreeTierStore(fake_redis_cache(server, namespace="ns").transport()) + server.connected = False + for _ in range(10): + outcome = await store.allocate_units("vertex_ai", "shared", "2026-10", 3, ttl=None) + assert isinstance(outcome, RedisFailure) and outcome.kind in ("transport", "breaker_open") + server.connected = True + blocked: Final = await store.allocate_units("vertex_ai", "shared", "2026-10", 3, ttl=None) + assert isinstance(blocked, RedisFailure) and blocked.kind == "breaker_open" + assert await raw.get(_KEY) is None + + +@pytest.mark.asyncio +async def test_expiry_server_failure_reports_consumed_allocation_without_retry( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + cache: Final = fake_redis_cache(server, namespace="ns") + store: Final = TransportSearchFreeTierStore(cache.transport()) + + def denied(_: object) -> None: + raise ResponseError("NOPERM expiry denied") + + set_reply_callback(cache, "EXPIRE", denied) + outcome: Final = await store.allocate_units("vertex_ai", "shared", "2026-10", 3, ttl=timedelta(seconds=60)) + assert isinstance(outcome, RedisFailure) and outcome.kind == "server" + assert await raw.get(_KEY) == b"3" + + +@pytest.mark.asyncio +async def test_allocation_flushes_pending_post_call_work_on_the_same_transport( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + cache: Final = fake_redis_cache(server, namespace="ns") + transport: Final = cache.transport() + store: Final = TransportSearchFreeTierStore(transport) + with request_redis_batch_scope(): + batch: Final = transport.post_call_batch() + assert batch is not None + pending: Final = batch.set("other-keyspace", b"queued", ttl=None) + assert await store.allocate_units("vertex_ai", "shared", "2026-10", 3, ttl=timedelta(seconds=60)) == 3.0 + assert pending.done and await pending is True + assert await raw.get("ns:other-keyspace") == b"queued" + assert await raw.get(_KEY) == b"3" + assert await raw.pttl(_KEY) == 60_000 From dd35d6f0d3d280039b7f8ce8b3e02757743a3747 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 21:28:32 +0000 Subject: [PATCH 11/17] refactor(redis): move user provider credentials into a typed store --- .../user_provider_credentials.py | 173 ++++---------- .../redis/stores/user_provider_credentials.py | 121 ++++++++++ tests/_support/fake_redis.py | 22 ++ .../credential_endpoints/test_endpoints.py | 95 ++++---- .../test_user_provider_credentials.py | 53 ++--- .../stores/test_user_provider_credentials.py | 218 ++++++++++++++++++ 6 files changed, 478 insertions(+), 204 deletions(-) create mode 100644 litellm/redis/stores/user_provider_credentials.py create mode 100644 tests/unit/redis/stores/test_user_provider_credentials.py diff --git a/litellm/proxy/credential_endpoints/user_provider_credentials.py b/litellm/proxy/credential_endpoints/user_provider_credentials.py index 22ee271410e..70950fd340e 100644 --- a/litellm/proxy/credential_endpoints/user_provider_credentials.py +++ b/litellm/proxy/credential_endpoints/user_provider_credentials.py @@ -7,16 +7,9 @@ payload encrypted; the request-time cache carries only ciphertext plus a negative marker so an unconnected user does not hit the DB per request. """ -import json from collections.abc import Mapping, Sequence -from typing import ( - TYPE_CHECKING, - Final, - Protocol, - cast, # noqa: TID251 # narrows the untyped Redis cache to the str-keyed Protocol -) - -from pydantic import BaseModel, ConfigDict, Field, ValidationError +from datetime import timedelta +from typing import TYPE_CHECKING, Final from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger @@ -25,15 +18,19 @@ from litellm.caching.dual_cache import DualCache if TYPE_CHECKING: from litellm.caching.redis_cache import RedisCache from litellm.constants import ( - GITHUB_COPILOT_DEFAULT_HOST, GITHUB_COPILOT_USER_CREDENTIAL_CACHE_TTL_SECONDS, - USER_PROVIDER_CREDENTIAL_CACHE_PREFIX, USER_PROVIDER_CREDENTIAL_NOT_CONNECTED, ) -from litellm.proxy.common_utils.encrypt_decrypt_utils import ( - decrypt_value_helper, - encrypt_value_helper, +from litellm.redis.stores.user_provider_credentials import ( + GithubCopilotUserConnectionPayload, + InvalidCredentialValue, + TransportUserProviderCredentialStore, + UserProviderCredentialStore, + decode_credential, + encode_credential, + user_provider_credential_key, ) +from litellm.redis.transport.contract import RedisFailure from litellm.repositories.chunked_in import find_many_in from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.table_repositories import PrismaTableRepository @@ -44,14 +41,7 @@ if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient _NOT_CONNECTED: Final = USER_PROVIDER_CREDENTIAL_NOT_CONNECTED - - -class GithubCopilotUserConnectionPayload(BaseModel): - model_config = ConfigDict(extra="ignore", frozen=True) - - access_token: str = Field(repr=False) - github_login: str - github_host: str = GITHUB_COPILOT_DEFAULT_HOST +_cache_key: Final = user_provider_credential_key class _UserProviderCredentialsRepository(PrismaTableRepository["prisma_models.LiteLLM_UserProviderCredentials"]): @@ -64,28 +54,12 @@ def _table( return _UserProviderCredentialsRepository(prisma_client, use_writer=True).table -def _cache_key(user_id: str, credential_name: str) -> str: - pair: Final = json.dumps([user_id, credential_name], separators=(",", ":")) - return f"{USER_PROVIDER_CREDENTIAL_CACHE_PREFIX}:{pair}" - - def _encode(payload: GithubCopilotUserConnectionPayload) -> str: - return encrypt_value_helper(payload.model_dump_json()) + return encode_credential(payload) def decode_user_provider_credential(stored: str) -> GithubCopilotUserConnectionPayload | None: - decrypted: Final = decrypt_value_helper( - value=stored, - key="user_provider_credential", - exception_type="debug", - return_original_value=False, - ) - if decrypted is None: - return None - try: - return GithubCopilotUserConnectionPayload.model_validate_json(decrypted) - except ValidationError: - return None + return decode_credential(stored) async def upsert_user_provider_credential( @@ -158,47 +132,15 @@ async def list_user_provider_credentials_for_credential( return await _table(prisma_client).find_many(where={"credential_name": credential_name}) -class _StringKeyCache(Protocol): - async def async_get_cache( - self, - key: str, - **kwargs: object, # kwargs-ok: RedisCache accepts optional cache kwargs - ) -> object: ... - - async def async_set_cache( - self, - key: str, - value: str, - **kwargs: object, # kwargs-ok: RedisCache accepts ttl/nx kwargs - ) -> None: ... - - async def async_delete_cache( - self, - key: str, - **kwargs: object, # kwargs-ok: RedisCache accepts optional cache kwargs - ) -> None: ... +def _store(cache: "RedisCache") -> UserProviderCredentialStore: + return TransportUserProviderCredentialStore(cache.transport()) -def _string_cache(cache: "RedisCache") -> _StringKeyCache: - return cast(_StringKeyCache, cache) # cast-ok: RedisCache exposes the str-keyed cache protocol - - -async def _try_cache_get(token_cache: "RedisCache", key: str) -> object: - """Redis errors must read as a plain miss: the DB is the source of truth.""" - try: - return await _string_cache(token_cache).async_get_cache(key) - except Exception: # noqa: BLE001 # a Redis outage must fall back to the database, never reject a connected user - verbose_proxy_logger.warning("aget_user_provider_credentials: Redis get failed; falling back to the database") - return None - - -async def _try_cache_set(token_cache: "RedisCache", key: str, value: str, nx: bool = False) -> None: - try: - await _string_cache(token_cache).async_set_cache( - key, value, nx=nx, ttl=GITHUB_COPILOT_USER_CREDENTIAL_CACHE_TTL_SECONDS - ) - except Exception: # noqa: BLE001 # caching is best-effort; a Redis outage is not worth failing the request - verbose_proxy_logger.warning("aget_user_provider_credentials: Redis set failed; skipping the cache write") +async def _try_cache_get( + token_cache: "RedisCache", user_id: str, credential_name: str +) -> str | None | InvalidCredentialValue: + result: Final = await _store(token_cache).read(user_id, credential_name) + return None if isinstance(result, RedisFailure) else result @with_service_target("user_provider_connections") @@ -214,23 +156,11 @@ async def invalidate_user_provider_credential_cache( token_cache: Final = cache.redis_cache if token_cache is None: return True - try: - await _string_cache(token_cache).async_set_cache( - _cache_key(user_id, credential_name), - _NOT_CONNECTED, - ttl=GITHUB_COPILOT_USER_CREDENTIAL_CACHE_TTL_SECONDS, - ) - except Exception: # noqa: BLE001 # caller decides whether a Redis outage aborts the disconnect - verbose_proxy_logger.warning("invalidate_user_provider_credential_cache: Redis tombstone failed") - return False - # async_set_cache swallows client errors internally, so a write that never - # landed looks identical to a success. The tombstone is only trusted when the - # key reads back as _NOT_CONNECTED; anything else means revoke failed. - readback: Final = await _try_cache_get(token_cache, _cache_key(user_id, credential_name)) - if readback != _NOT_CONNECTED: - verbose_proxy_logger.warning("invalidate_user_provider_credential_cache: tombstone not visible after write") - return False - return True + await _store(token_cache).disconnect( + user_id, credential_name, ttl=timedelta(seconds=GITHUB_COPILOT_USER_CREDENTIAL_CACHE_TTL_SECONDS) + ) + readback: Final = await _try_cache_get(token_cache, user_id, credential_name) + return readback == _NOT_CONNECTED @with_service_target("user_provider_connections") @@ -244,10 +174,7 @@ async def drop_user_provider_credential_cache( token_cache: Final = cache.redis_cache if token_cache is None: return - try: - await _string_cache(token_cache).async_delete_cache(_cache_key(user_id, credential_name)) - except Exception: # noqa: BLE001 # a Redis outage must not fail a connect - verbose_proxy_logger.warning("drop_user_provider_credential_cache: Redis delete failed") + await _store(token_cache).delete(user_id, credential_name) @with_service_target("user_provider_connections") @@ -265,32 +192,21 @@ async def set_user_provider_credential_cache( token_cache: Final = cache.redis_cache if token_cache is None: return True - key: Final = _cache_key(user_id, credential_name) - try: - await _string_cache(token_cache).async_set_cache( - key, - _encode(payload), - ttl=GITHUB_COPILOT_USER_CREDENTIAL_CACHE_TTL_SECONDS, - ) - readback: Final = await _try_cache_get(token_cache, key) + write: Final = await _store(token_cache).connect( + user_id, + credential_name, + payload, + ttl=timedelta(seconds=GITHUB_COPILOT_USER_CREDENTIAL_CACHE_TTL_SECONDS), + ) + if not (isinstance(write, RedisFailure) and write.kind == "decode"): + readback: Final = await _try_cache_get(token_cache, user_id, credential_name) if isinstance(readback, str) and decode_user_provider_credential(readback) == payload: return True - verbose_proxy_logger.warning( - "set_user_provider_credential_cache: Redis overwrite not visible; falling back to delete" - ) - except Exception: # noqa: BLE001 # a Redis outage must not fail a connect - verbose_proxy_logger.warning("set_user_provider_credential_cache: Redis set failed; falling back to delete") - try: - await _string_cache(token_cache).async_delete_cache(key) - except Exception: # noqa: BLE001 # a Redis outage must not fail a connect - verbose_proxy_logger.warning("set_user_provider_credential_cache: Redis delete failed") - after_delete: Final = await _try_cache_get(token_cache, key) + await _store(token_cache).delete(user_id, credential_name) + after_delete: Final = await _try_cache_get(token_cache, user_id, credential_name) if after_delete is None: return True - if isinstance(after_delete, str) and decode_user_provider_credential(after_delete) == payload: - return True - verbose_proxy_logger.warning("set_user_provider_credential_cache: stale entry survived delete") - return False + return isinstance(after_delete, str) and decode_user_provider_credential(after_delete) == payload @with_service_target("user_provider_connections") @@ -313,8 +229,9 @@ async def aget_user_provider_credentials( cached: Final[dict[str, str]] = {} # mutable-ok: accumulates hits and DB reads misses: Final[list[str]] = [] # mutable-ok: accumulates cache misses for name in names: - value = await _try_cache_get(token_cache, _cache_key(user_id, name)) if token_cache is not None else None - if value == _NOT_CONNECTED: + if ( + value := await _try_cache_get(token_cache, user_id, name) if token_cache is not None else None + ) == _NOT_CONNECTED: continue if isinstance(value, str) and value: cached[name] = value @@ -340,13 +257,15 @@ async def aget_user_provider_credentials( found: Final = {row.credential_name: row.credential_b64 for row in rows} for name in misses: if token_cache is not None: - await _try_cache_set(token_cache, _cache_key(user_id, name), found.get(name, _NOT_CONNECTED), nx=True) + await _store(token_cache).fill( + user_id, name, found.get(name), ttl=timedelta(seconds=GITHUB_COPILOT_USER_CREDENTIAL_CACHE_TTL_SECONDS) + ) if name in found: cached[name] = found[name] return { - name: connection + name: found_connection for name, ciphertext in cached.items() - if (connection := _connection_from_ciphertext(ciphertext)) is not None + if (found_connection := _connection_from_ciphertext(ciphertext)) is not None } diff --git a/litellm/redis/stores/user_provider_credentials.py b/litellm/redis/stores/user_provider_credentials.py new file mode 100644 index 00000000000..ee7d8772ad0 --- /dev/null +++ b/litellm/redis/stores/user_provider_credentials.py @@ -0,0 +1,121 @@ +import ast +import json +from dataclasses import dataclass +from datetime import timedelta +from typing import Final, Protocol + +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError + +from litellm.constants import ( + GITHUB_COPILOT_DEFAULT_HOST, + USER_PROVIDER_CREDENTIAL_CACHE_PREFIX, + USER_PROVIDER_CREDENTIAL_NOT_CONNECTED, +) +from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper, encrypt_value_helper +from litellm.redis.transport.contract import RedisFailure, RedisTransport + +_JSON: Final = TypeAdapter(object) + + +@dataclass(frozen=True, slots=True) +class InvalidCredentialValue: + pass + + +class GithubCopilotUserConnectionPayload(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + access_token: str = Field(repr=False) + github_login: str + github_host: str = GITHUB_COPILOT_DEFAULT_HOST + + +def user_provider_credential_key(user_id: str, credential_name: str) -> str: + pair: Final = json.dumps([user_id, credential_name], separators=(",", ":")) + return f"{USER_PROVIDER_CREDENTIAL_CACHE_PREFIX}:{pair}" + + +def encode_credential(payload: GithubCopilotUserConnectionPayload) -> str: + return encrypt_value_helper(payload.model_dump_json()) + + +def decode_credential(stored: str) -> GithubCopilotUserConnectionPayload | None: + decrypted: Final = decrypt_value_helper( + value=stored, key="user_provider_credential", exception_type="debug", return_original_value=False + ) + if decrypted is None: + return None + try: + return GithubCopilotUserConnectionPayload.model_validate_json(decrypted) + except ValidationError: + return None + + +def _decode_reply(reply: bytes | None | RedisFailure) -> str | None | InvalidCredentialValue | RedisFailure: + if not isinstance(reply, bytes): + return reply + try: + decoded: Final = reply.decode("utf-8") + try: + value: Final = _JSON.validate_json(decoded) + except ValidationError: + legacy: Final = _JSON.validate_python(ast.literal_eval(decoded)) + return legacy if isinstance(legacy, str) or legacy is None else InvalidCredentialValue() + return value if isinstance(value, str) or value is None else InvalidCredentialValue() + except (ValueError, SyntaxError): + return RedisFailure(kind="decode", message="Invalid user provider credential cache value") + + +class UserProviderCredentialStore(Protocol): + async def read(self, user_id: str, credential_name: str) -> str | None | InvalidCredentialValue | RedisFailure: ... + + async def fill( + self, user_id: str, credential_name: str, ciphertext: str | None, *, ttl: timedelta + ) -> bool | RedisFailure: ... + + async def connect( + self, user_id: str, credential_name: str, payload: GithubCopilotUserConnectionPayload, *, ttl: timedelta + ) -> bool | RedisFailure: ... + + async def disconnect(self, user_id: str, credential_name: str, *, ttl: timedelta) -> bool | RedisFailure: ... + + async def delete(self, user_id: str, credential_name: str) -> int | RedisFailure: ... + + +@dataclass(frozen=True, slots=True) +class TransportUserProviderCredentialStore: + transport: RedisTransport + + async def read(self, user_id: str, credential_name: str) -> str | None | InvalidCredentialValue | RedisFailure: + return _decode_reply(await self.transport.get(user_provider_credential_key(user_id, credential_name))) + + async def fill( + self, user_id: str, credential_name: str, ciphertext: str | None, *, ttl: timedelta + ) -> bool | RedisFailure: + return await self.transport.set_seconds( + user_provider_credential_key(user_id, credential_name), + json.dumps(ciphertext if ciphertext is not None else USER_PROVIDER_CREDENTIAL_NOT_CONNECTED), + ttl=ttl, + nx=True, + ) + + async def connect( + self, user_id: str, credential_name: str, payload: GithubCopilotUserConnectionPayload, *, ttl: timedelta + ) -> bool | RedisFailure: + try: + ciphertext: Final = encode_credential(payload) + except Exception: # noqa: BLE001 # the caller handles encryption failures like failed writes + return RedisFailure(kind="decode", message="Cannot encode user provider credential") + return await self.transport.set_seconds( + user_provider_credential_key(user_id, credential_name), json.dumps(ciphertext), ttl=ttl + ) + + async def disconnect(self, user_id: str, credential_name: str, *, ttl: timedelta) -> bool | RedisFailure: + return await self.transport.set_seconds( + user_provider_credential_key(user_id, credential_name), + json.dumps(USER_PROVIDER_CREDENTIAL_NOT_CONNECTED), + ttl=ttl, + ) + + async def delete(self, user_id: str, credential_name: str) -> int | RedisFailure: + return await self.transport.delete(user_provider_credential_key(user_id, credential_name)) diff --git a/tests/_support/fake_redis.py b/tests/_support/fake_redis.py index 79a6ef3eb20..79f7e271ec6 100644 --- a/tests/_support/fake_redis.py +++ b/tests/_support/fake_redis.py @@ -83,3 +83,25 @@ def trip_breaker(cache: RedisCache, failures: int | None = None) -> None: breaker: Final = cache._circuit_breaker # pyright: ignore[reportPrivateUsage] # tests open the breaker the way RedisCache does after hard failures for _ in range(breaker.failure_threshold if failures is None else failures): breaker.record_failure() + + +def async_client(cache: RedisCache) -> fakeredis.FakeAsyncRedis: + client: Final = cache.init_async_client() # pyright: ignore[reportUnknownMemberType] # legacy RedisCache API is untyped + assert isinstance(client, fakeredis.FakeAsyncRedis) + return client + + +def pin_async_client(cache: RedisCache) -> Callable[[], None]: + client: Final = async_client(cache) + + def pin() -> None: + litellm.in_memory_llm_clients_cache.set_cache( + key=cache._get_async_client_cache_key(), # pyright: ignore[reportPrivateUsage] # re-seeds the per-loop client slot RedisCache reads + value=client, + ) + + return pin + + +async def legacy_set(cache: RedisCache, key: str, value: object, *, ttl: int | None = None) -> None: + await cache.async_set_cache(key, value, ttl=ttl) # pyright: ignore[reportUnknownMemberType] # legacy RedisCache API is untyped diff --git a/tests/unit/proxy/credential_endpoints/test_endpoints.py b/tests/unit/proxy/credential_endpoints/test_endpoints.py index 3d8e92988d9..46e99d48d4f 100644 --- a/tests/unit/proxy/credential_endpoints/test_endpoints.py +++ b/tests/unit/proxy/credential_endpoints/test_endpoints.py @@ -1,12 +1,14 @@ """Tests for the credential management endpoints.""" import json +from collections.abc import Awaitable, Callable from contextlib import contextmanager from functools import reduce from types import SimpleNamespace from typing import Final, Protocol from unittest.mock import AsyncMock, MagicMock, call, patch +import fakeredis import httpx import pytest from cryptography.hazmat.primitives import serialization @@ -16,6 +18,7 @@ from fastapi.testclient import TestClient from pydantic import TypeAdapter import litellm +from litellm.caching.redis_cache import RedisCache from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.litellm_logging import get_masked_values from litellm.llms.github_copilot.per_user_auth import _SESSION_CACHE, _session_cache_key @@ -26,10 +29,21 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helpe from litellm.proxy.credential_endpoints.endpoints import _decode_flow_handle, get_llm_router from litellm.proxy.proxy_server import app from litellm.types.utils import CredentialItem +from tests._support.fake_redis import async_client, fake_redis_cache, pin_async_client client = TestClient(app) +def _auth_with_redis(cache: RedisCache, auth: Callable[[], object]) -> Callable[[], Awaitable[object]]: + pin: Final = pin_async_client(cache) + + async def authenticate() -> object: + pin() + return auth() + + return authenticate + + def _as_admin(): return UserAPIKeyAuth(api_key="test-key", user_role="proxy_admin") @@ -929,8 +943,8 @@ class TestNonAdminCannotPersistWifFieldsOnCredential: ids=["non_admin", "proxy_admin"], ) def test_only_a_proxy_admin_can_create_a_credential_with_an_oidc_api_key( - self, restore_credential_list, auth, expected_status: int - ): + self, restore_credential_list: None, auth: Callable[[], object], expected_status: int + ) -> None: """The call path reads an oidc/ api_key from the proxy's own filesystem, so it is a server-owned value the same way a federation field is.""" with _repository_holding(None) as repository: @@ -1799,17 +1813,8 @@ def test_update_credential_purges_connections_when_github_host_changes(monkeypat proxy_server.prisma_client.db.litellm_userprovidercredentials = table proxy_server.prisma_client.writer_db = proxy_server.prisma_client.db - cached_values: Final[dict[str, object]] = {} # mutable-ok: fake Redis stores invalidation writes - - async def set_cache(key: str, value: str, **kwargs: object) -> None: - cached_values[key] = value - - async def get_cache(key: str, **kwargs: object) -> object: - return cached_values.get(key) - - redis_cache: Final = MagicMock() - redis_cache.async_set_cache = AsyncMock(side_effect=set_cache) - redis_cache.async_get_cache = AsyncMock(side_effect=get_cache) + server: Final = fakeredis.FakeServer() + redis_cache: Final = fake_redis_cache(server) monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache(redis_cache=redis_cache)) session_key: Final = _session_cache_key("user-a", "copilot", "github.com") @@ -1825,6 +1830,7 @@ def test_update_credential_purges_connections_when_github_host_changes(monkeypat "credential_values": {"github_copilot_host": "octo.ghe.com"}, "credential_info": {}, }, + auth=_auth_with_redis(redis_cache, _as_admin), ) assert response.status_code == 200, response.text @@ -1835,7 +1841,10 @@ def test_update_credential_purges_connections_when_github_host_changes(monkeypat ] ) table.delete_many.assert_awaited_once_with(where={"credential_name": "copilot"}) - assert cached_values[upc._cache_key("user-a", "copilot")] == upc._NOT_CONNECTED + assert ( + fakeredis.FakeRedis(server=server).get(upc._cache_key("user-a", "copilot")) + == json.dumps(upc._NOT_CONNECTED).encode() + ) assert _SESSION_CACHE.get_cache(session_key) is None @@ -2125,7 +2134,7 @@ def _github_http(mapping): """Patch the shared async HTTP client to an httpx MockTransport that answers GitHub endpoints from {url: response-json dict or callable}; anything else fails.""" - def respond(request): + def respond(request: httpx.Request) -> httpx.Response: entry = mapping.get(str(request.url)) if entry is None: raise AssertionError(f"unexpected request to {request.url}") @@ -2146,7 +2155,7 @@ _DEVICE_FLOW_START = { } -def _patch_user_connection_env(monkeypatch, credentials, rows=()): +def _patch_user_connection_env(monkeypatch: pytest.MonkeyPatch, credentials, rows=()): """Per-user env: credential in memory, prisma table answering find_many with rows, and a real in-memory cache for device-flow + credential caching.""" monkeypatch.setattr(litellm, "credential_list", list(credentials)) @@ -2465,7 +2474,7 @@ def test_user_connection_poll_reports_no_copilot_seat(monkeypatch, seat_status): nothing, and clears the in-flight device code.""" table = _patch_user_connection_env(monkeypatch, [_per_user_credential()]) - def respond(request): + def respond(request: httpx.Request) -> httpx.Response: url = str(request.url) if url == _DEVICE_CODE_URL: return httpx.Response(200, json=_DEVICE_FLOW_START, request=request) @@ -2493,7 +2502,7 @@ def test_user_connection_poll_reports_no_copilot_seat(monkeypatch, seat_status): def test_user_connection_poll_rate_limit_returns_429(monkeypatch): table = _patch_user_connection_env(monkeypatch, [_per_user_credential()]) - def respond(request): + def respond(request: httpx.Request) -> httpx.Response: url = str(request.url) if url == _DEVICE_CODE_URL: return httpx.Response(200, json=_DEVICE_FLOW_START, request=request) @@ -2809,7 +2818,7 @@ def test_user_connection_route_check_allows_slashed_names_for_internal_users(): ) -def _flow_handle_for(user_id="user-a", credential_name="copilot-cred", expires_at=None): +def _flow_handle_for(user_id: str = "user-a", credential_name: str = "copilot-cred", expires_at: float | None = None): import time as _time from litellm.models.credentials import UserConnectionFlowHandle @@ -2901,18 +2910,16 @@ def test_delete_user_connection_fails_closed_when_the_tombstone_write_fails(monk """Redis attached but its set() raises: the disconnect must not delete the row, or the cached token would keep working until TTL with nothing left to revoke.""" table = _patch_user_connection_env(monkeypatch, [_per_user_credential()], rows=[_connection_row()]) - from types import SimpleNamespace - from litellm.proxy import proxy_server - failing_redis = SimpleNamespace( - async_set_cache=AsyncMock(side_effect=ConnectionError("redis down")), - async_get_cache=AsyncMock(return_value=None), - async_delete_cache=AsyncMock(return_value=None), - ) + server: Final = fakeredis.FakeServer() + server.connected = False + failing_redis: Final = fake_redis_cache(server) monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache(redis_cache=failing_redis)) - response = _call_as("DELETE", "/credentials/copilot-cred/user_connection", auth=_as_user()) + response = _call_as( + "DELETE", "/credentials/copilot-cred/user_connection", auth=_auth_with_redis(failing_redis, _as_user()) + ) assert response.status_code == 503, response.text table.delete_many.assert_not_awaited() @@ -2922,38 +2929,34 @@ def test_delete_user_connection_fails_when_the_tombstone_write_silently_noops(mo without writing still looks successful. The disconnect must verify the tombstone by reading the key back and refuse to delete the row.""" table = _patch_user_connection_env(monkeypatch, [_per_user_credential()], rows=[_connection_row()]) - from types import SimpleNamespace - from litellm.proxy import proxy_server - store: dict = {} - silent_redis = SimpleNamespace( - async_set_cache=AsyncMock(return_value=None), # reports success, writes nothing - async_get_cache=AsyncMock(side_effect=lambda key: store.get(key)), - async_delete_cache=AsyncMock(return_value=None), - ) + silent_redis: Final = fake_redis_cache(fakeredis.FakeServer()) + monkeypatch.setattr(async_client(silent_redis), "set", AsyncMock(return_value=None)) monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache(redis_cache=silent_redis)) - response = _call_as("DELETE", "/credentials/copilot-cred/user_connection", auth=_as_user()) + response = _call_as( + "DELETE", "/credentials/copilot-cred/user_connection", auth=_auth_with_redis(silent_redis, _as_user()) + ) assert response.status_code == 503, response.text table.delete_many.assert_not_awaited() -def test_user_connection_poll_reports_503_when_the_cache_overwrite_cannot_land(monkeypatch): +def test_user_connection_poll_reports_503_when_the_cache_overwrite_cannot_land(monkeypatch: pytest.MonkeyPatch) -> None: """The connect save is durable, so a stale not-connected marker that survives both the overwrite and the delete must surface as a retryable 503 rather than a false "connected".""" table = _patch_user_connection_env(monkeypatch, [_per_user_credential()]) - from types import SimpleNamespace - from litellm.proxy import proxy_server - stale_redis = SimpleNamespace( - async_set_cache=AsyncMock(return_value=None), - async_get_cache=AsyncMock(return_value="stale-value"), - async_delete_cache=AsyncMock(return_value=None), - ) + server: Final = fakeredis.FakeServer() + raw: Final = fakeredis.FakeRedis(server=server) + raw.set('user_provider_credential:["user-a","copilot-cred"]', json.dumps("stale-value")) + stale_redis: Final = fake_redis_cache(server) + monkeypatch.setattr(async_client(stale_redis), "set", AsyncMock(return_value=None)) + monkeypatch.setattr(async_client(stale_redis), "delete", AsyncMock(return_value=0)) monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache(redis_cache=stale_redis)) + redis_auth: Final = _auth_with_redis(stale_redis, _as_user()) with _github_http( { @@ -2963,12 +2966,12 @@ def test_user_connection_poll_reports_503_when_the_cache_overwrite_cannot_land(m _GITHUB_USER_URL: {"login": "octo"}, } ): - start = _call_as("POST", "/credentials/copilot-cred/user_connection/start", auth=_as_user()) + start = _call_as("POST", "/credentials/copilot-cred/user_connection/start", auth=redis_auth) poll = _call_as( "POST", "/credentials/copilot-cred/user_connection/poll", json_body={"flow_handle": start.json()["flow_handle"]}, - auth=_as_user(), + auth=redis_auth, ) assert poll.status_code == 503, poll.text assert "cache could not be refreshed" in poll.text diff --git a/tests/unit/proxy/credential_endpoints/test_user_provider_credentials.py b/tests/unit/proxy/credential_endpoints/test_user_provider_credentials.py index dae3a0459e2..8c65dcb04cf 100644 --- a/tests/unit/proxy/credential_endpoints/test_user_provider_credentials.py +++ b/tests/unit/proxy/credential_endpoints/test_user_provider_credentials.py @@ -3,8 +3,9 @@ import json from types import SimpleNamespace from typing import Final -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, patch +import fakeredis import pytest from litellm.caching.dual_cache import DualCache @@ -22,6 +23,7 @@ from litellm.proxy.credential_endpoints.user_provider_credentials import ( set_user_provider_credential_cache, upsert_user_provider_credential, ) +from tests._support.fake_redis import async_client, fake_redis_cache, legacy_set def _prisma(table=None): @@ -41,15 +43,7 @@ def _row(user_id="user-a", credential_name="copilot-cred", credential_b64="ciphe def _fake_redis(monkeypatch): - """A RedisCache backed by a fakeredis client; two DualCaches can share it.""" - import fakeredis - - from litellm.caching.redis_cache import RedisCache - - fake = fakeredis.FakeAsyncRedis() - redis = RedisCache(host="localhost", port=6379) - monkeypatch.setattr(redis, "init_async_client", lambda: fake) - return redis + return fake_redis_cache(fakeredis.FakeServer()) @pytest.fixture(autouse=True) @@ -279,8 +273,9 @@ async def test_redis_read_failure_falls_back_to_the_database(monkeypatch): """A broken Redis must not 401 a connected user: reads degrade to a cache miss.""" from litellm.proxy.credential_endpoints import user_provider_credentials as upc - redis = _fake_redis(monkeypatch) - monkeypatch.setattr(redis, "async_get_cache", AsyncMock(side_effect=RuntimeError("redis down"))) + server: Final = fakeredis.FakeServer() + server.connected = False + redis = fake_redis_cache(server) cache = DualCache(redis_cache=redis) payload = GithubCopilotUserConnectionPayload(access_token="gho_secret", github_login="octo") @@ -295,9 +290,9 @@ async def test_redis_read_failure_falls_back_to_the_database(monkeypatch): @pytest.mark.asyncio async def test_redis_write_and_delete_failures_do_not_fail_the_request(monkeypatch): - redis = _fake_redis(monkeypatch) - monkeypatch.setattr(redis, "async_set_cache", AsyncMock(side_effect=RuntimeError("redis down"))) - monkeypatch.setattr(redis, "async_delete_cache", AsyncMock(side_effect=RuntimeError("redis down"))) + server: Final = fakeredis.FakeServer() + server.connected = False + redis = fake_redis_cache(server) cache = DualCache(redis_cache=redis) payload = GithubCopilotUserConnectionPayload(access_token="gho_secret", github_login="octo") @@ -342,7 +337,6 @@ async def test_without_redis_no_worker_local_entries_are_kept(): async def test_connect_overwrites_a_not_connected_tombstone(monkeypatch): """Connect writes the new connection over any cached tombstone so the next read sees the token immediately, no TTL wait.""" - from litellm.proxy.credential_endpoints import user_provider_credentials as upc redis = _fake_redis(monkeypatch) cache = DualCache(redis_cache=redis) @@ -365,7 +359,6 @@ async def test_stale_not_connected_fill_cannot_overwrite_a_connect_write(monkeyp """A read that fetched "not connected" before the poll saved the row fills with set-if-absent; the connect's plain overwrite must win so the marker never lingers for the TTL.""" - from litellm.proxy.credential_endpoints import user_provider_credentials as upc redis = _fake_redis(monkeypatch) cache = DualCache(redis_cache=redis) @@ -399,15 +392,13 @@ async def test_set_reports_connected_when_delete_clears_a_dropped_write(): """async_set_cache swallows errors: a set that reports success but writes nothing is verified by read-back, and a working delete still resolves the stale entry.""" - store: dict = {} - silent_redis = SimpleNamespace( - async_set_cache=AsyncMock(return_value=None), - async_get_cache=AsyncMock(side_effect=lambda key: store.get(key)), - async_delete_cache=AsyncMock(side_effect=lambda key: store.pop(key, None)), - ) + silent_redis: Final = fake_redis_cache(fakeredis.FakeServer()) cache = DualCache(redis_cache=silent_redis) payload = GithubCopilotUserConnectionPayload(access_token="gho_secret", github_login="octo") - assert await set_user_provider_credential_cache(cache, "user-a", "copilot-cred", payload) + await legacy_set(silent_redis, 'user_provider_credential:["user-a","copilot-cred"]', "stale-value") + with patch.object(async_client(silent_redis), "set", AsyncMock(return_value=None)): + assert await set_user_provider_credential_cache(cache, "user-a", "copilot-cred", payload) + assert await silent_redis.async_get_cache('user_provider_credential:["user-a","copilot-cred"]') is None @pytest.mark.asyncio @@ -415,15 +406,15 @@ async def test_set_reports_failure_when_a_stale_entry_survives_both_attempts(): """A write silently dropped while a not-connected marker stays cached means the caller cannot trust the connect: the helper reports failure instead of claiming success.""" - stale = "stale-value" - silent_redis = SimpleNamespace( - async_set_cache=AsyncMock(return_value=None), - async_get_cache=AsyncMock(return_value=stale), - async_delete_cache=AsyncMock(return_value=None), - ) + silent_redis: Final = fake_redis_cache(fakeredis.FakeServer()) cache = DualCache(redis_cache=silent_redis) payload = GithubCopilotUserConnectionPayload(access_token="gho_secret", github_login="octo") - assert not await set_user_provider_credential_cache(cache, "user-a", "copilot-cred", payload) + await legacy_set(silent_redis, 'user_provider_credential:["user-a","copilot-cred"]', "stale-value") + with ( + patch.object(async_client(silent_redis), "set", AsyncMock(return_value=None)), + patch.object(async_client(silent_redis), "delete", AsyncMock(return_value=0)), + ): + assert not await set_user_provider_credential_cache(cache, "user-a", "copilot-cred", payload) @pytest.mark.asyncio diff --git a/tests/unit/redis/stores/test_user_provider_credentials.py b/tests/unit/redis/stores/test_user_provider_credentials.py new file mode 100644 index 00000000000..27447620cbf --- /dev/null +++ b/tests/unit/redis/stores/test_user_provider_credentials.py @@ -0,0 +1,218 @@ +import json +from datetime import timedelta +from types import SimpleNamespace +from typing import Final +from unittest.mock import AsyncMock, MagicMock, patch + +import fakeredis +import pytest +from redis.exceptions import ConnectionError, ResponseError + +from litellm.caching.dual_cache import DualCache +from litellm.constants import GITHUB_COPILOT_USER_CREDENTIAL_CACHE_TTL_SECONDS +from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper +from litellm.proxy.credential_endpoints.user_provider_credentials import ( + GithubCopilotUserConnectionPayload, + aget_user_provider_credentials, + decode_user_provider_credential, + drop_user_provider_credential_cache, + invalidate_user_provider_credential_cache, + set_user_provider_credential_cache, +) +from litellm.redis.stores.user_provider_credentials import InvalidCredentialValue, TransportUserProviderCredentialStore +from litellm.redis.transport.contract import RedisFailure +from tests._support.fake_redis import fake_redis_cache, trip_breaker + + +@pytest.fixture(autouse=True) +def encryption_key(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "sk-test-master") + monkeypatch.setattr("time.time", lambda: 1_800_000_000.0) + + +@pytest.mark.asyncio +async def test_old_and_new_connections_share_ciphertext_key_tombstones_and_ttl() -> None: + server: Final = fakeredis.FakeServer() + redis: Final = fake_redis_cache(server, namespace="credentials") + raw: Final = fakeredis.FakeAsyncRedis(server=server) + cache: Final = DualCache(redis_cache=redis) + user: Final = 'user:é"' + name: Final = "provider:東京" + key: Final = "user_provider_credential:" + json.dumps([user, name], separators=(",", ":")) + payload: Final = GithubCopilotUserConnectionPayload(access_token="test-token-é", github_login="octo") + ciphertext: Final = encrypt_value_helper(payload.model_dump_json()) + prisma: Final = MagicMock() + await redis.async_set_cache(key, ciphertext, ttl=60) + assert await raw.get("credentials:" + key) == json.dumps(ciphertext).encode() + assert await aget_user_provider_credentials(prisma, cache, user, [name]) == { + name: {"access_token": payload.access_token, "github_host": "github.com"} + } + assert await invalidate_user_provider_credential_cache(cache, user, name) + assert await raw.get("credentials:" + key) == b'"__not_connected__"' + assert await redis.async_get_cache(key) == "__not_connected__" + assert await aget_user_provider_credentials(prisma, cache, user, [name]) == {} + assert await set_user_provider_credential_cache(cache, user, name, payload) + old_read: Final = await redis.async_get_cache(key) + assert isinstance(old_read, str) + assert decode_user_provider_credential(old_read) == payload + assert await raw.get("credentials:" + key) == json.dumps(old_read).encode() + assert await raw.pttl("credentials:" + key) == GITHUB_COPILOT_USER_CREDENTIAL_CACHE_TTL_SECONDS * 1000 + assert await raw.keys("*") == [("credentials:" + key).encode()] + assert await cache.in_memory_cache.async_get_cache(key) is None + await drop_user_provider_credential_cache(cache, user, name) + assert await raw.get("credentials:" + key) is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("applied", [False, True]) +async def test_ambiguous_tombstone_requires_readback(applied: bool) -> None: + server: Final = fakeredis.FakeServer() + redis: Final = fake_redis_cache(server, namespace="credentials") + raw: Final = fakeredis.FakeAsyncRedis(server=server) + cache: Final = DualCache(redis_cache=redis) + key: Final = 'credentials:user_provider_credential:["u","c"]' + + async def lost_reply(*args: object, **kwargs: object) -> None: + if applied: + await raw.set(key, b'"__not_connected__"', ex=60) + raise ConnectionError("lost reply") + + with patch.object(redis.init_async_client(), "set", lost_reply): + assert await invalidate_user_provider_credential_cache(cache, "u", "c") is applied + assert await raw.get(key) == (b'"__not_connected__"' if applied else None) + + +@pytest.mark.asyncio +async def test_ambiguous_connection_write_is_read_before_delete() -> None: + server: Final = fakeredis.FakeServer() + redis: Final = fake_redis_cache(server) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + payload: Final = GithubCopilotUserConnectionPayload(access_token="test-token", github_login="octo") + key: Final = 'user_provider_credential:["u","c"]' + ciphertext: Final = encrypt_value_helper(payload.model_dump_json()) + + async def lost_reply(*args: object, **kwargs: object) -> None: + await raw.set(key, json.dumps(ciphertext), ex=60) + raise ConnectionError("lost reply") + + with patch.object(redis.init_async_client(), "set", lost_reply): + assert await set_user_provider_credential_cache(DualCache(redis_cache=redis), "u", "c", payload) + assert await raw.get(key) == json.dumps(ciphertext).encode() + + +@pytest.mark.asyncio +async def test_outage_uses_database_but_refuses_disconnect_and_keeps_connect_fallback() -> None: + server: Final = fakeredis.FakeServer() + redis: Final = fake_redis_cache(server) + cache: Final = DualCache(redis_cache=redis) + payload: Final = GithubCopilotUserConnectionPayload(access_token="test-token", github_login="octo") + prisma: Final = MagicMock() + prisma.writer_db.litellm_userprovidercredentials.find_many = AsyncMock( + return_value=[ + SimpleNamespace(credential_name="c", credential_b64=encrypt_value_helper(payload.model_dump_json())) + ] + ) + server.connected = False + assert await aget_user_provider_credentials(prisma, cache, "u", ["c"]) == { + "c": {"access_token": "test-token", "github_host": "github.com"} + } + assert not await invalidate_user_provider_credential_cache(cache, "u", "c") + assert await set_user_provider_credential_cache(cache, "u", "c", payload) + store: Final = TransportUserProviderCredentialStore(redis.transport()) + assert isinstance(await store.read("u", "c"), RedisFailure) + trip_breaker(redis, failures=3) + result: Final = await store.read("u", "c") + assert isinstance(result, RedisFailure) and result.kind == "breaker_open" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("value", [b"garbage", b"\xff"]) +async def test_malformed_values_are_typed_failures_and_caller_database_misses(value: bytes) -> None: + server: Final = fakeredis.FakeServer() + redis: Final = fake_redis_cache(server) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await raw.set('user_provider_credential:["u","c"]', value) + store: Final = TransportUserProviderCredentialStore(redis.transport()) + result: Final = await store.read("u", "c") + assert isinstance(result, RedisFailure) and result.kind == "decode" + prisma: Final = MagicMock() + prisma.writer_db.litellm_userprovidercredentials.find_many = AsyncMock(return_value=[]) + assert await aget_user_provider_credentials(prisma, DualCache(redis_cache=redis), "u", ["c"]) == {} + assert prisma.writer_db.litellm_userprovidercredentials.find_many.await_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("value", [b"{}", b"42", b"[]", b"false"]) +async def test_non_string_surviving_failed_write_and_delete_is_not_a_confirmed_connect(value: bytes) -> None: + server: Final = fakeredis.FakeServer() + redis: Final = fake_redis_cache(server) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + key: Final = 'user_provider_credential:["u","c"]' + await raw.set(key, value) + store: Final = TransportUserProviderCredentialStore(redis.transport()) + assert isinstance(await store.read("u", "c"), InvalidCredentialValue) + payload: Final = GithubCopilotUserConnectionPayload(access_token="test-token", github_login="octo") + with ( + patch.object(redis.init_async_client(), "set", AsyncMock(side_effect=ResponseError("readonly"))), + patch.object(redis.init_async_client(), "delete", AsyncMock(side_effect=ResponseError("readonly"))), + ): + assert not await set_user_provider_credential_cache(DualCache(redis_cache=redis), "u", "c", payload) + assert await raw.get(key) == value + + +@pytest.mark.asyncio +async def test_tombstone_expires_and_reloads_from_database(monkeypatch: pytest.MonkeyPatch) -> None: + server: Final = fakeredis.FakeServer() + redis: Final = fake_redis_cache(server) + cache: Final = DualCache(redis_cache=redis) + payload: Final = GithubCopilotUserConnectionPayload(access_token="test-token", github_login="octo") + prisma: Final = MagicMock() + prisma.writer_db.litellm_userprovidercredentials.find_many = AsyncMock( + return_value=[ + SimpleNamespace(credential_name="c", credential_b64=encrypt_value_helper(payload.model_dump_json())) + ] + ) + assert await invalidate_user_provider_credential_cache(cache, "u", "c") + assert await aget_user_provider_credentials(prisma, cache, "u", ["c"]) == {} + monkeypatch.setattr("time.time", lambda: 1_800_000_061.0) + assert await aget_user_provider_credentials(prisma, cache, "u", ["c"]) == { + "c": {"access_token": "test-token", "github_host": "github.com"} + } + + +@pytest.mark.asyncio +async def test_fill_nx_retains_tombstone_and_server_errors_are_values() -> None: + server: Final = fakeredis.FakeServer() + redis: Final = fake_redis_cache(server) + cache: Final = DualCache(redis_cache=redis) + store: Final = TransportUserProviderCredentialStore(redis.transport()) + assert await invalidate_user_provider_credential_cache(cache, "u", "c") + assert await store.fill("u", "c", "stale-ciphertext", ttl=timedelta(seconds=60)) is False + assert await store.read("u", "c") == "__not_connected__" + with patch.object(redis.init_async_client(), "set", AsyncMock(side_effect=ResponseError("readonly"))): + result: Final = await store.disconnect("u", "c", ttl=timedelta(seconds=60)) + assert isinstance(result, RedisFailure) and result.kind == "server" + + +@pytest.mark.asyncio +async def test_in_flight_database_fill_cannot_resurrect_peer_disconnect() -> None: + server: Final = fakeredis.FakeServer() + redis: Final = fake_redis_cache(server) + reader: Final = DualCache(redis_cache=redis) + peer: Final = DualCache(redis_cache=redis) + payload: Final = GithubCopilotUserConnectionPayload(access_token="stale-token", github_login="octo") + + async def database_read(**kwargs: object) -> list[SimpleNamespace]: + assert await invalidate_user_provider_credential_cache(peer, "u", "c") + return [SimpleNamespace(credential_name="c", credential_b64=encrypt_value_helper(payload.model_dump_json()))] + + prisma: Final = MagicMock() + prisma.writer_db.litellm_userprovidercredentials.find_many = database_read + assert await aget_user_provider_credentials(prisma, reader, "u", ["c"]) == { + "c": {"access_token": "stale-token", "github_host": "github.com"} + } + assert await aget_user_provider_credentials(prisma, reader, "u", ["c"]) == {} + assert ( + await fakeredis.FakeAsyncRedis(server=server).get('user_provider_credential:["u","c"]') + == b'"__not_connected__"' + ) From 3e0dda2017d3f642202335a63c751f4f42aef4ae Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 21:28:32 +0000 Subject: [PATCH 12/17] refactor(redis): move MCP per-user tokens into a typed store --- .../proxy/mcp_server/oauth2_token_cache.py | 27 ++- .../dual_cache_token_backend.py | 87 ++++++- .../outbound_credentials/token_cache_codec.py | 58 +---- litellm/redis/stores/mcp_per_user_tokens.py | 96 ++++++++ .../test_dual_cache_token_backend.py | 95 ++++---- .../mcp_server/test_per_user_oauth_cache.py | 70 ++---- .../redis/stores/test_mcp_per_user_tokens.py | 224 ++++++++++++++++++ 7 files changed, 482 insertions(+), 175 deletions(-) create mode 100644 litellm/redis/stores/mcp_per_user_tokens.py create mode 100644 tests/unit/redis/stores/test_mcp_per_user_tokens.py diff --git a/litellm/proxy/mcp_server/oauth2_token_cache.py b/litellm/proxy/mcp_server/oauth2_token_cache.py index 1e6e83c0f4b..45d4efa7a8e 100644 --- a/litellm/proxy/mcp_server/oauth2_token_cache.py +++ b/litellm/proxy/mcp_server/oauth2_token_cache.py @@ -17,13 +17,13 @@ from litellm._internal_context import with_service_target from litellm._logging import verbose_logger from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import ( + AUTH_OBJECTS_TARGET, MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL, MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE, MCP_OAUTH2_TOKEN_CACHE_MIN_TTL, MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS, MCP_PER_USER_TOKEN_DEFAULT_TTL, MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS, - MCP_PER_USER_TOKEN_REDIS_KEY_PREFIX, ) from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.proxy.common_utils.encrypt_decrypt_utils import ( @@ -34,9 +34,11 @@ from litellm.proxy.mcp_server.oauth_utils import ( build_upstream_oauth2_token_request, resolve_upstream_resource, ) +from litellm.proxy.mcp_server.outbound_credentials.dual_cache_token_backend import DualCacheTokenCacheBackend from litellm.proxy.mcp_server.outbound_credentials.oauth_token_store import OAuthToken from litellm.proxy.mcp_server.outbound_credentials.token_cache_codec import OAuthTokenCacheCodec from litellm.proxy.mcp_server.utils import MCP_OAUTH_TOKENS_TARGET +from litellm.redis.stores.mcp_per_user_tokens import per_user_token_key from litellm.types.llms.custom_http import httpxSpecialProvider if TYPE_CHECKING: @@ -241,7 +243,7 @@ class MCPPerUserTokenCache: """ def _cache_key(self, user_id: str, server_id: str) -> str: - return f"{MCP_PER_USER_TOKEN_REDIS_KEY_PREFIX}:{user_id}:{server_id}" + return per_user_token_key(user_id, server_id) def _codec(self) -> OAuthTokenCacheCodec: return OAuthTokenCacheCodec( @@ -258,11 +260,7 @@ class MCPPerUserTokenCache: try: from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415 - key: Final = self._cache_key(user_id, server_id) - encrypted: Final = await user_api_key_cache.async_get_cache(key) - if encrypted is None: - return None - return self._codec().decode(encrypted) + return await DualCacheTokenCacheBackend(user_api_key_cache, self._codec()).get(user_id, server_id) except Exception as exc: verbose_logger.debug( "MCPPerUserTokenCache.get failed for user=%s server=%s: %s", @@ -285,11 +283,12 @@ class MCPPerUserTokenCache: try: from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415 - key: Final = self._cache_key(user_id, server_id) - encrypted: Final = self._codec().encode( - OAuthToken(access_token=access_token, identity_binding_proof=identity_binding_proof) + await DualCacheTokenCacheBackend(user_api_key_cache, self._codec()).cache_token( + user_id, + server_id, + OAuthToken(access_token=access_token, identity_binding_proof=identity_binding_proof), + ttl, ) - await user_api_key_cache.async_set_cache(key, encrypted, ttl=ttl) verbose_logger.debug( "MCPPerUserTokenCache.set: cached token for user=%s server=%s ttl=%ds", user_id, @@ -304,16 +303,18 @@ class MCPPerUserTokenCache: exc, ) + @with_service_target(AUTH_OBJECTS_TARGET) async def delete(self, user_id: str, server_id: str) -> None: """Invalidate the cached token in Redis, here, and in every peer worker's in-memory layer.""" try: from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( # noqa: PLC0415 # proxy import cycle - evict_and_broadcast, + publish_auth_cache_invalidation, ) from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415 key: Final = self._cache_key(user_id, server_id) - await evict_and_broadcast((key,), user_api_key_cache) + await DualCacheTokenCacheBackend(user_api_key_cache, self._codec()).evict_token(user_id, server_id) + await publish_auth_cache_invalidation(cache_key=key) except Exception as exc: verbose_logger.debug( "MCPPerUserTokenCache.delete failed for user=%s server=%s: %s", diff --git a/litellm/proxy/mcp_server/outbound_credentials/dual_cache_token_backend.py b/litellm/proxy/mcp_server/outbound_credentials/dual_cache_token_backend.py index b360b787962..64bd0580b3c 100644 --- a/litellm/proxy/mcp_server/outbound_credentials/dual_cache_token_backend.py +++ b/litellm/proxy/mcp_server/outbound_credentials/dual_cache_token_backend.py @@ -10,10 +10,14 @@ the other across the cutover). A missing or undecryptable entry reads as a miss. from __future__ import annotations from dataclasses import KW_ONLY, dataclass +from datetime import timedelta from typing import Final, Protocol from litellm._internal_context import with_service_target from litellm._logging import verbose_logger +from litellm.caching.redis_batch import active_request_redis_batch +from litellm.caching.redis_cache import RedisCache +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.mcp_server.outbound_credentials.oauth_token_store import ( OAuthToken, ) @@ -21,16 +25,34 @@ from litellm.proxy.mcp_server.outbound_credentials.token_cache_codec import ( OAuthTokenCacheCodec, ) from litellm.proxy.mcp_server.utils import MCP_OAUTH_TOKENS_TARGET +from litellm.redis.stores.mcp_per_user_tokens import ( + TOKEN_KEY_PREFIX, + McpPerUserTokenStore, + TransportMcpPerUserTokenStore, + per_user_token_key, +) +from litellm.redis.transport.contract import RedisFailure + + +class TokenMemory(Protocol): + async def async_get_cache(self, key: str) -> object: ... + + async def async_set_cache(self, key: str, value: str, *, ttl: float | None = None) -> None: ... + + def delete_cache(self, key: str) -> None: ... class AsyncCache(Protocol): """The slice of LiteLLM's ``DualCache`` this backend needs (Redis-backed, shared across workers).""" - async def async_get_cache(self, key: str) -> object | None: ... + @property + def redis_cache(self) -> RedisCache | None: ... - async def async_set_cache(self, key: str, value: str, ttl: float | None = None) -> None: ... + @property + def in_memory_cache(self) -> TokenMemory: ... - async def async_delete_cache(self, key: str) -> None: ... + @property + def default_in_memory_ttl(self) -> float | None: ... @dataclass(frozen=True, slots=True) @@ -44,17 +66,39 @@ class DualCacheTokenCacheBackend: cache: AsyncCache codec: OAuthTokenCacheCodec _: KW_ONLY - key_prefix: str = "mcp:per_user_token:" + key_prefix: str = TOKEN_KEY_PREFIX def _key(self, user_id: str, server_id: str) -> str: - return f"{self.key_prefix}{user_id}:{server_id}" + return per_user_token_key(user_id, server_id, prefix=self.key_prefix) + + def _store(self) -> McpPerUserTokenStore | None: + redis: Final = self.cache.redis_cache + return None if redis is None else TransportMcpPerUserTokenStore(redis.transport(), self.key_prefix) + + async def _read(self, user_id: str, server_id: str) -> object: + key: Final = self._key(user_id, server_id) + local: Final = await self.cache.in_memory_cache.async_get_cache(key) + store: Final = self._store() + if local is not None or store is None: + return local + redis: Final = self.cache.redis_cache + batch: Final = None if redis is None else active_request_redis_batch(redis) + if batch is not None and batch.read_as_missing(key): + return None + blob: Final = await store.read(user_id, server_id) + if isinstance(blob, RedisFailure): + return None + if blob is not None: + ttl: Final = self.cache.default_in_memory_ttl + await self.cache.in_memory_cache.async_set_cache(key, blob, **({"ttl": ttl} if ttl is not None else {})) + return blob @with_service_target(MCP_OAUTH_TOKENS_TARGET) async def get(self, user_id: str, server_id: str) -> OAuthToken | None: try: - blob: Final = await self.cache.async_get_cache(self._key(user_id, server_id)) + blob: Final = await self._read(user_id, server_id) return self.codec.decode(blob) if isinstance(blob, str) else None - except Exception as exc: # noqa: BLE001 + except Exception as exc: # noqa: BLE001 # cache and codec failures remain safe misses verbose_logger.debug("MCP per-user token cache get failed (miss): %s", exc) return None @@ -62,18 +106,37 @@ class DualCacheTokenCacheBackend: async def set(self, user_id: str, server_id: str, token: OAuthToken, ttl_seconds: float) -> None: if ttl_seconds <= 0: return + await self.cache_token(user_id, server_id, token, ttl_seconds) + + async def cache_token(self, user_id: str, server_id: str, token: OAuthToken, ttl_seconds: float) -> None: try: - await self.cache.async_set_cache( + ciphertext: Final = self.codec.encode(token) + await self.cache.in_memory_cache.async_set_cache( self._key(user_id, server_id), - self.codec.encode(token), + ciphertext, ttl=ttl_seconds, ) - except Exception as exc: # noqa: BLE001 + store: Final = self._store() + if store is not None: + await store.write( + user_id, + server_id, + ciphertext, + ttl=timedelta(seconds=int(ttl_seconds)), + defer=isinstance(self.cache, UserApiKeyCache), + ) + except Exception as exc: # noqa: BLE001 # cache and codec failures must not abort token resolution verbose_logger.debug("MCP per-user token cache set failed (ignored): %s", exc) @with_service_target(MCP_OAUTH_TOKENS_TARGET) async def delete(self, user_id: str, server_id: str) -> None: + await self.evict_token(user_id, server_id) + + async def evict_token(self, user_id: str, server_id: str) -> None: try: - await self.cache.async_delete_cache(self._key(user_id, server_id)) - except Exception as exc: # noqa: BLE001 + self.cache.in_memory_cache.delete_cache(self._key(user_id, server_id)) + store: Final = self._store() + if store is not None: + await store.delete(user_id, server_id) + except Exception as exc: # noqa: BLE001 # best-effort eviction must not abort token resolution verbose_logger.debug("MCP per-user token cache delete failed (ignored): %s", exc) diff --git a/litellm/proxy/mcp_server/outbound_credentials/token_cache_codec.py b/litellm/proxy/mcp_server/outbound_credentials/token_cache_codec.py index 4576e925939..5f979da6bed 100644 --- a/litellm/proxy/mcp_server/outbound_credentials/token_cache_codec.py +++ b/litellm/proxy/mcp_server/outbound_credentials/token_cache_codec.py @@ -1,57 +1,3 @@ -"""Serialize + encrypt boundary for caching an OAuth token in a shared (Redis) cache. - -Shared cache values contain an encrypted access token and optional identity-binding proof. -Refresh tokens remain in the database; cache TTL bounds the access token's lifetime. -Legacy bearer-only entries decode without proof and cannot satisfy identity enforcement. -""" - -from __future__ import annotations - -import json -from collections.abc import Callable -from dataclasses import dataclass -from typing import Final - -from pydantic import TypeAdapter, ValidationError - -from litellm.proxy.mcp_server.outbound_credentials.oauth_token_store import ( - OAuthToken, +from litellm.redis.stores.mcp_per_user_tokens import ( + OAuthTokenCacheCodec as OAuthTokenCacheCodec, # noqa: PLC0414 # public compatibility export ) - -_BOUND_PREFIX: Final = "litellm-bound-oauth-v1:" -_BOUND_PAYLOAD: Final = TypeAdapter(dict[str, str]) - - -@dataclass(frozen=True, slots=True) -class OAuthTokenCacheCodec: - encrypt: Callable[[str], str] - decrypt: Callable[[str], str | None] - - def encode(self, token: OAuthToken) -> str: - if token.identity_binding_proof is not None: - return self.encrypt( - _BOUND_PREFIX - + json.dumps( - { - "access_token": token.access_token, - "identity_binding_proof": token.identity_binding_proof, - } - ) - ) - return self.encrypt(token.access_token) - - def decode(self, blob: str) -> OAuthToken | None: - access_token: Final = self.decrypt(blob) - if not access_token: - return None - if access_token.startswith(_BOUND_PREFIX): - try: - payload: Final = _BOUND_PAYLOAD.validate_json(access_token[len(_BOUND_PREFIX) :]) - except ValidationError: - return None - bearer: Final = payload.get("access_token") - proof: Final = payload.get("identity_binding_proof") - if not bearer or not proof: - return None - return OAuthToken(access_token=bearer, identity_binding_proof=proof) - return OAuthToken(access_token=access_token, refresh_token=None) diff --git a/litellm/redis/stores/mcp_per_user_tokens.py b/litellm/redis/stores/mcp_per_user_tokens.py new file mode 100644 index 00000000000..f7bf8cf17a7 --- /dev/null +++ b/litellm/redis/stores/mcp_per_user_tokens.py @@ -0,0 +1,96 @@ +import ast +import json +from collections.abc import Callable +from dataclasses import dataclass +from datetime import timedelta +from typing import Final, Protocol + +from pydantic import TypeAdapter, ValidationError + +from litellm.constants import MCP_PER_USER_TOKEN_REDIS_KEY_PREFIX +from litellm.proxy.mcp_server.outbound_credentials.oauth_token_store import OAuthToken +from litellm.redis.transport.contract import RedisFailure, RedisTransport + +TOKEN_KEY_PREFIX: Final = MCP_PER_USER_TOKEN_REDIS_KEY_PREFIX + ":" +_BOUND_PREFIX: Final = "litellm-bound-oauth-v1:" +_BOUND_PAYLOAD: Final = TypeAdapter(dict[str, str]) +_STRING: Final = TypeAdapter(str) + + +def per_user_token_key(user_id: str, server_id: str, *, prefix: str = TOKEN_KEY_PREFIX) -> str: + return f"{prefix}{user_id}:{server_id}" + + +@dataclass(frozen=True, slots=True) +class OAuthTokenCacheCodec: + encrypt: Callable[[str], str] + decrypt: Callable[[str], str | None] + + def encode(self, token: OAuthToken) -> str: + if token.identity_binding_proof is not None: + return self.encrypt( + _BOUND_PREFIX + + json.dumps( + {"access_token": token.access_token, "identity_binding_proof": token.identity_binding_proof} + ) + ) + return self.encrypt(token.access_token) + + def decode(self, blob: str) -> OAuthToken | None: + access_token: Final = self.decrypt(blob) + if not access_token: + return None + if access_token.startswith(_BOUND_PREFIX): + try: + payload: Final = _BOUND_PAYLOAD.validate_json(access_token[len(_BOUND_PREFIX) :]) + except ValidationError: + return None + bearer: Final = payload.get("access_token") + proof: Final = payload.get("identity_binding_proof") + if not bearer or not proof: + return None + return OAuthToken(access_token=bearer, identity_binding_proof=proof) + return OAuthToken(access_token=access_token, refresh_token=None) + + +class McpPerUserTokenStore(Protocol): + async def read(self, user_id: str, server_id: str) -> str | None | RedisFailure: ... + + async def write( + self, user_id: str, server_id: str, ciphertext: str, *, ttl: timedelta, defer: bool = False + ) -> bool | RedisFailure: ... + + async def delete(self, user_id: str, server_id: str) -> int | RedisFailure: ... + + +@dataclass(frozen=True, slots=True) +class TransportMcpPerUserTokenStore: + transport: RedisTransport + prefix: str = TOKEN_KEY_PREFIX + + async def read(self, user_id: str, server_id: str) -> str | None | RedisFailure: + reply: Final = await self.transport.get(per_user_token_key(user_id, server_id, prefix=self.prefix)) + if reply is None or isinstance(reply, RedisFailure): + return reply + try: + text: Final = reply.decode("utf-8") + try: + return _STRING.validate_json(text) + except ValidationError: + return _STRING.validate_python(ast.literal_eval(text)) + except (ValueError, SyntaxError): + return RedisFailure(kind="decode", message="Invalid MCP per-user token cache value") + + async def write( + self, user_id: str, server_id: str, ciphertext: str, *, ttl: timedelta, defer: bool = False + ) -> bool | RedisFailure: + key: Final = per_user_token_key(user_id, server_id, prefix=self.prefix) + wire: Final = json.dumps(ciphertext) + batch: Final = self.transport.request_batch() if defer else None + if batch is not None: + batch.set_seconds(key, wire, ttl=ttl) + return True + return await self.transport.set_seconds(key, wire, ttl=ttl) + + async def delete(self, user_id: str, server_id: str) -> int | RedisFailure: + return await self.transport.delete(per_user_token_key(user_id, server_id, prefix=self.prefix)) diff --git a/tests/unit/proxy/mcp_server/outbound_credentials/test_dual_cache_token_backend.py b/tests/unit/proxy/mcp_server/outbound_credentials/test_dual_cache_token_backend.py index 00db27a60b8..a298214e688 100644 --- a/tests/unit/proxy/mcp_server/outbound_credentials/test_dual_cache_token_backend.py +++ b/tests/unit/proxy/mcp_server/outbound_credentials/test_dual_cache_token_backend.py @@ -1,7 +1,11 @@ """Tests for the DualCache-backed token cache backend: encrypted round-trip, key, TTL, miss.""" +from typing import Final + +import fakeredis import pytest +from litellm.caching.dual_cache import DualCache from litellm.proxy.mcp_server.outbound_credentials.dual_cache_token_backend import ( DualCacheTokenCacheBackend, ) @@ -11,25 +15,21 @@ from litellm.proxy.mcp_server.outbound_credentials.oauth_token_store import ( from litellm.proxy.mcp_server.outbound_credentials.token_cache_codec import ( OAuthTokenCacheCodec, ) +from tests._support.fake_redis import fake_redis_cache -class _FakeCache: - def __init__(self): - self.values = {} - self.ttls = {} - - async def async_get_cache(self, key): - return self.values.get(key) - - async def async_set_cache(self, key, value, ttl=None): - self.values[key] = value - self.ttls[key] = ttl - - async def async_delete_cache(self, key): - self.values.pop(key, None) +def _cache() -> DualCache: + return DualCache(redis_cache=fake_redis_cache(fakeredis.FakeServer())) -def _backend(cache): +def _raw(cache: DualCache) -> fakeredis.FakeAsyncRedis: + assert cache.redis_cache is not None + client: Final = cache.redis_cache.init_async_client() + assert isinstance(client, fakeredis.FakeAsyncRedis) + return client + + +def _backend(cache: DualCache) -> DualCacheTokenCacheBackend: codec = OAuthTokenCacheCodec( encrypt=lambda s: f"enc:{s}", decrypt=lambda b: b[4:] if b.startswith("enc:") else None, @@ -38,16 +38,17 @@ def _backend(cache): @pytest.mark.asyncio -async def test_set_then_get_round_trips_encrypted_under_the_per_user_key(): - cache = _FakeCache() +async def test_set_then_get_round_trips_encrypted_under_the_per_user_key(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr("time.time", lambda: 1_800_000_000.0) + cache = _cache() backend = _backend(cache) await backend.set("alice", "srv", OAuthToken(access_token="at"), 120.0) key = "mcp:per_user_token:alice:srv" - assert cache.ttls[key] == 120.0 + assert await _raw(cache).pttl(key) == 120000 # Stored via the codec (encrypted), not the bare token; the codec's own test proves real # NaCl output hides the secret - here the fake encrypt just wraps, so we check it was applied. - assert cache.values[key] == "enc:at" + assert await _raw(cache).get(key) == b'"enc:at"' got = await backend.get("alice", "srv") assert got is not None and got.access_token == "at" @@ -55,13 +56,13 @@ async def test_set_then_get_round_trips_encrypted_under_the_per_user_key(): @pytest.mark.asyncio async def test_get_missing_key_is_none(): - assert await _backend(_FakeCache()).get("alice", "srv") is None + assert await _backend(_cache()).get("alice", "srv") is None @pytest.mark.asyncio async def test_non_str_cache_value_is_a_miss(): - cache = _FakeCache() - cache.values["mcp:per_user_token:alice:srv"] = 12345 # corrupt / wrong type + cache = _cache() + cache.in_memory_cache.set_cache("mcp:per_user_token:alice:srv", 12345) assert await _backend(cache).get("alice", "srv") is None @@ -69,21 +70,21 @@ async def test_non_str_cache_value_is_a_miss(): async def test_undecryptable_blob_is_a_miss(): # e.g. a master-key rotation leaves an entry the codec can't decrypt; it must read as a miss so # the store re-reads the DB, not raise (the decrypt helper here returns None for unknown blobs). - cache = _FakeCache() - cache.values["mcp:per_user_token:alice:srv"] = "not-our-ciphertext" + cache = _cache() + cache.in_memory_cache.set_cache("mcp:per_user_token:alice:srv", "not-our-ciphertext") assert await _backend(cache).get("alice", "srv") is None @pytest.mark.asyncio async def test_non_positive_ttl_is_not_written(): - cache = _FakeCache() + cache = _cache() await _backend(cache).set("alice", "srv", OAuthToken(access_token="at"), 0.0) - assert cache.values == {} # an already-expired token is not cached + assert await cache.async_get_cache("mcp:per_user_token:alice:srv") is None @pytest.mark.asyncio async def test_delete_removes_the_entry(): - cache = _FakeCache() + cache = _cache() backend = _backend(cache) await backend.set("alice", "srv", OAuthToken(access_token="at"), 60.0) await backend.delete("alice", "srv") @@ -92,7 +93,7 @@ async def test_delete_removes_the_entry(): @pytest.mark.asyncio async def test_keys_isolate_users_and_servers(): - cache = _FakeCache() + cache = _cache() backend = _backend(cache) await backend.set("alice", "srv", OAuthToken(access_token="a"), 60.0) await backend.set("bob", "srv", OAuthToken(access_token="b"), 60.0) @@ -102,35 +103,30 @@ async def test_keys_isolate_users_and_servers(): assert bob is not None and bob.access_token == "b" -class _RaisingCache(_FakeCache): - """A cache whose every op raises, e.g. a Redis outage tripping the circuit breaker.""" - - async def async_get_cache(self, key): - raise ConnectionError("redis down") - - async def async_set_cache(self, key, value, ttl=None): - raise ConnectionError("redis down") - - async def async_delete_cache(self, key): - raise ConnectionError("redis down") +def _offline_cache(): + server = fakeredis.FakeServer() + server.connected = False + return DualCache(redis_cache=fake_redis_cache(server)) -# A cache or codec failure must degrade to the safe value (miss / no-op), never propagate: otherwise a -# Redis outage turns CachedOAuthTokenStore.fetch() into a 500 instead of a cache miss that re-reads the -# DB (get/set) or issues the OAuth challenge (the unauthorized branch deletes before returning None). @pytest.mark.asyncio async def test_get_is_a_miss_when_the_cache_raises(): - assert await _backend(_RaisingCache()).get("alice", "srv") is None + assert await _backend(_offline_cache()).get("alice", "srv") is None @pytest.mark.asyncio async def test_set_is_swallowed_when_the_cache_raises(): - await _backend(_RaisingCache()).set("alice", "srv", OAuthToken(access_token="at"), 60.0) + backend = _backend(_offline_cache()) + await backend.set("alice", "srv", OAuthToken(access_token="at"), 60.0) + assert await backend.get("alice", "srv") == OAuthToken(access_token="at") @pytest.mark.asyncio async def test_delete_is_swallowed_when_the_cache_raises(): - await _backend(_RaisingCache()).delete("alice", "srv") + backend = _backend(_offline_cache()) + await backend.set("alice", "srv", OAuthToken(access_token="at"), 60.0) + await backend.delete("alice", "srv") + assert await backend.get("alice", "srv") is None @pytest.mark.asyncio @@ -139,8 +135,9 @@ async def test_set_is_swallowed_when_the_codec_raises(): raise ValueError("encrypt unavailable") codec = OAuthTokenCacheCodec(encrypt=_boom, decrypt=lambda b: None) - backend = DualCacheTokenCacheBackend(_FakeCache(), codec) - await backend.set("alice", "srv", OAuthToken(access_token="at"), 60.0) # must not raise + backend = DualCacheTokenCacheBackend(_cache(), codec) + await backend.set("alice", "srv", OAuthToken(access_token="at"), 60.0) + assert await backend.get("alice", "srv") is None @pytest.mark.asyncio @@ -151,6 +148,6 @@ async def test_get_is_a_miss_when_decrypt_raises(): raise ValueError("bad ciphertext") codec = OAuthTokenCacheCodec(encrypt=lambda s: f"enc:{s}", decrypt=_boom) - cache = _FakeCache() - cache.values["mcp:per_user_token:alice:srv"] = "whatever" + cache = _cache() + cache.in_memory_cache.set_cache("mcp:per_user_token:alice:srv", "whatever") assert await DualCacheTokenCacheBackend(cache, codec).get("alice", "srv") is None diff --git a/tests/unit/proxy/mcp_server/test_per_user_oauth_cache.py b/tests/unit/proxy/mcp_server/test_per_user_oauth_cache.py index c0e23319566..a953910c9be 100644 --- a/tests/unit/proxy/mcp_server/test_per_user_oauth_cache.py +++ b/tests/unit/proxy/mcp_server/test_per_user_oauth_cache.py @@ -7,12 +7,15 @@ Unit tests for per-user MCP OAuth token storage: """ import sys +import time from datetime import datetime, timedelta, timezone -from typing import Any, Dict, Optional +from typing import Any, Final from unittest.mock import AsyncMock, MagicMock, patch import pytest +from litellm.caching.dual_cache import DualCache + # Stub out modules that aren't available in the unit-test environment # so we can import the targets without a full proxy stack. for _mod in ("orjson",): @@ -22,7 +25,6 @@ for _mod in ("orjson",): from litellm.proxy.mcp_server.oauth2_token_cache import ( # noqa: E402 MCPPerUserTokenCache, compute_per_user_token_ttl, - mcp_per_user_token_cache, ) from litellm.types.mcp import MCPAuth, MCPTransport # noqa: E402 from litellm.types.mcp_server.mcp_server_manager import MCPServer # noqa: E402 @@ -41,7 +43,7 @@ def _import_validate(): def _make_server(**kwargs) -> MCPServer: - defaults: Dict[str, Any] = { + defaults: dict[str, Any] = { "server_id": "slack-test", "name": "Slack", "server_name": "slack", @@ -262,9 +264,7 @@ class TestMCPPerUserTokenCache: @pytest.mark.asyncio async def test_get_returns_none_on_miss(self, cache, mock_dual_cache): with ( - patch( - "litellm.proxy.mcp_server.oauth2_token_cache.decrypt_value_helper" - ) as mock_decrypt, + patch("litellm.proxy.mcp_server.oauth2_token_cache.decrypt_value_helper") as mock_decrypt, patch("litellm.proxy.proxy_server.user_api_key_cache", mock_dual_cache), ): mock_dual_cache.async_get_cache.return_value = None @@ -273,66 +273,46 @@ class TestMCPPerUserTokenCache: mock_decrypt.assert_not_called() @pytest.mark.asyncio - async def test_get_decrypts_cached_value(self, cache, mock_dual_cache): - fake_encrypted = "encrypted_blob_abc123" - fake_plaintext = "xoxb-slack-token" + async def test_get_decrypts_cached_value(self, cache: MCPPerUserTokenCache) -> None: + dual: Final = DualCache() + await dual.async_set_cache("mcp:per_user_token:alice:slack-test", "encrypted_blob_xyz", ttl=60) with ( patch( "litellm.proxy.mcp_server.oauth2_token_cache.decrypt_value_helper", - return_value=fake_plaintext, + return_value="xoxb-slack-token", ) as mock_decrypt, - patch("litellm.proxy.proxy_server.user_api_key_cache", mock_dual_cache), + patch("litellm.proxy.proxy_server.user_api_key_cache", dual), ): - mock_dual_cache.async_get_cache.return_value = fake_encrypted result = await cache.get("alice", "slack-test") - assert result == fake_plaintext - mock_decrypt.assert_called_once_with( - fake_encrypted, - key="mcp_per_user_token", - exception_type="debug", - ) + assert result == "xoxb-slack-token" + mock_decrypt.assert_called_once_with("encrypted_blob_xyz", key="mcp_per_user_token", exception_type="debug") @pytest.mark.asyncio - async def test_set_encrypts_before_storing(self, cache, mock_dual_cache): - fake_encrypted = "encrypted_blob_xyz" + async def test_set_encrypts_before_storing_under_the_exact_key(self, cache: MCPPerUserTokenCache) -> None: + dual: Final = DualCache() with ( patch( "litellm.proxy.mcp_server.oauth2_token_cache.encrypt_value_helper", - return_value=fake_encrypted, + return_value="encrypted_blob_xyz", ) as mock_encrypt, - patch("litellm.proxy.proxy_server.user_api_key_cache", mock_dual_cache), + patch("litellm.proxy.proxy_server.user_api_key_cache", dual), ): await cache.set("alice", "slack-test", "xoxb-token", ttl=3540) mock_encrypt.assert_called_once_with("xoxb-token") - mock_dual_cache.async_set_cache.assert_called_once() - call_kwargs = mock_dual_cache.async_set_cache.call_args - assert call_kwargs[0][1] == fake_encrypted # encrypted value stored - assert call_kwargs[1]["ttl"] == 3540 + assert dual.in_memory_cache.get_cache("mcp:per_user_token:alice:slack-test") == "encrypted_blob_xyz" + assert dual.in_memory_cache.get_cache("xoxb-token") is None + assert 0 < dual.in_memory_cache.ttl_dict["mcp:per_user_token:alice:slack-test"] - time.time() <= 3540 @pytest.mark.asyncio - async def test_set_uses_correct_cache_key(self, cache, mock_dual_cache): - with ( - patch( - "litellm.proxy.mcp_server.oauth2_token_cache.encrypt_value_helper", - return_value="enc", - ), - patch("litellm.proxy.proxy_server.user_api_key_cache", mock_dual_cache), - ): - await cache.set("bob", "github-server", "ghp_token", ttl=3600) - - key_used = mock_dual_cache.async_set_cache.call_args[0][0] - assert key_used == "mcp:per_user_token:bob:github-server" - - @pytest.mark.asyncio - async def test_delete_calls_async_delete_cache(self, cache, mock_dual_cache): - mock_dual_cache.async_delete_cache = AsyncMock() - with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_dual_cache): + async def test_delete_removes_the_cached_token(self, cache: MCPPerUserTokenCache) -> None: + dual: Final = DualCache() + await dual.async_set_cache("mcp:per_user_token:alice:slack-test", "encrypted_blob_xyz", ttl=60) + with patch("litellm.proxy.proxy_server.user_api_key_cache", dual): await cache.delete("alice", "slack-test") - mock_dual_cache.async_delete_cache.assert_called_once_with(key="mcp:per_user_token:alice:slack-test") - mock_dual_cache.async_set_cache.assert_not_called() + assert dual.in_memory_cache.get_cache("mcp:per_user_token:alice:slack-test") is None @pytest.mark.asyncio async def test_get_returns_none_on_decrypt_failure(self, cache, mock_dual_cache): diff --git a/tests/unit/redis/stores/test_mcp_per_user_tokens.py b/tests/unit/redis/stores/test_mcp_per_user_tokens.py new file mode 100644 index 00000000000..84f610c3433 --- /dev/null +++ b/tests/unit/redis/stores/test_mcp_per_user_tokens.py @@ -0,0 +1,224 @@ +import json +from contextlib import nullcontext +from datetime import timedelta +from typing import Final +from unittest.mock import patch + +import fakeredis +import pytest +from pydantic import TypeAdapter +from redis.exceptions import ConnectionError + +from litellm.caching.redis_batch import active_request_redis_batch, request_redis_batch_scope +from litellm.proxy import proxy_server +from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( + AuthCacheInvalidationSubscriber, + auth_cache_invalidation_channel, +) +from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper, encrypt_value_helper +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.mcp_server.oauth2_token_cache import MCPPerUserTokenCache +from litellm.proxy.mcp_server.outbound_credentials.dual_cache_token_backend import DualCacheTokenCacheBackend +from litellm.proxy.mcp_server.outbound_credentials.oauth_token_store import OAuthToken +from litellm.proxy.mcp_server.outbound_credentials.token_cache_codec import OAuthTokenCacheCodec +from litellm.redis.stores.mcp_per_user_tokens import TransportMcpPerUserTokenStore +from litellm.redis.transport.contract import RedisFailure +from tests._support.fake_redis import fake_redis_cache, trip_breaker + + +@pytest.fixture(autouse=True) +def encryption_and_clock(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(proxy_server, "master_key", "sk-test-master") + monkeypatch.setattr("time.time", lambda: 1_800_000_000.0) + + +_MESSAGE: Final = TypeAdapter(dict[str, object]) + + +def _legacy_codec(cache: MCPPerUserTokenCache) -> OAuthTokenCacheCodec: + return cache._codec() # pyright: ignore[reportPrivateUsage] # the legacy cache owns the codec the new backend must match + + +@pytest.mark.asyncio +@pytest.mark.parametrize("proof", [None, "bound-é"]) +async def test_legacy_wire_is_shared_by_both_callers_with_exact_codec_and_expiry( + monkeypatch: pytest.MonkeyPatch, + proof: str | None, +) -> None: + server: Final = fakeredis.FakeServer() + redis: Final = fake_redis_cache(server, namespace="oauth") + raw: Final = fakeredis.FakeAsyncRedis(server=server) + cache: Final = UserApiKeyCache(redis_cache=redis) + legacy: Final = MCPPerUserTokenCache() + backend: Final = DualCacheTokenCacheBackend(cache, _legacy_codec(legacy)) + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + user: Final = 'a:é"' + key: Final = f"mcp:per_user_token:{user}:東京" + token: Final = OAuthToken("bearer-é", identity_binding_proof=proof) + plaintext: Final = ( + token.access_token + if proof is None + else "litellm-bound-oauth-v1:" + + json.dumps({"access_token": token.access_token, "identity_binding_proof": proof}) + ) + ciphertext: Final = encrypt_value_helper(plaintext) + await redis.async_set_cache(key, ciphertext, ttl=90) + assert await backend.get(user, "東京") == token + assert await legacy.get_token(user, "東京") == token + await backend.set(user, "東京", token, 120.9) + written: Final = await redis.async_get_cache(key) + assert isinstance(written, str) + assert decrypt_value_helper(written, "test") == plaintext + assert await raw.get("oauth:" + key) == json.dumps(written).encode() + assert await raw.pttl("oauth:" + key) == 120000 + assert await raw.keys("*") == [("oauth:" + key).encode()] + cache.in_memory_cache.delete_cache(key) + await legacy.set(user, "東京", "updated", 73, proof) + cache.in_memory_cache.delete_cache(key) + assert await backend.get(user, "東京") == OAuthToken("updated", identity_binding_proof=proof) + assert await raw.pttl("oauth:" + key) == 73000 + await backend.delete(user, "東京") + assert await legacy.get(user, "東京") is None + assert await raw.get("oauth:" + key) is None + + +@pytest.mark.asyncio +async def test_auth_cache_deferred_write_joins_request_batch_and_retains_memory_policy() -> None: + server: Final = fakeredis.FakeServer() + redis: Final = fake_redis_cache(server, namespace="oauth") + raw: Final = fakeredis.FakeAsyncRedis(server=server) + cache: Final = UserApiKeyCache(redis_cache=redis) + backend: Final = DualCacheTokenCacheBackend(cache, _legacy_codec(MCPPerUserTokenCache())) + with request_redis_batch_scope(): + await backend.set("u", "s", OAuthToken("pending"), 20.9) + assert await backend.get("u", "s") == OAuthToken("pending") + assert await raw.get("oauth:mcp:per_user_token:u:s") is None + batch: Final = active_request_redis_batch(redis) + assert batch is not None + await batch.flush() + assert await raw.pttl("oauth:mcp:per_user_token:u:s") == 20000 + cache.in_memory_cache.delete_cache("mcp:per_user_token:u:s") + assert await backend.get("u", "s") == OAuthToken("pending") + batch.delete("mcp:per_user_token:u:s") + cache.in_memory_cache.delete_cache("mcp:per_user_token:u:s") + assert await backend.get("u", "s") is None + await batch.flush() + assert await raw.get("oauth:mcp:per_user_token:u:s") is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("wire", [b"{}", b"null", b"42", b"broken", b"\xff", b'"litellm-bound-oauth-v1:{}"']) +async def test_malformed_redis_tokens_are_safe_misses(monkeypatch: pytest.MonkeyPatch, wire: bytes) -> None: + server: Final = fakeredis.FakeServer() + redis: Final = fake_redis_cache(server) + cache: Final = UserApiKeyCache(redis_cache=redis) + legacy: Final = MCPPerUserTokenCache() + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + await fakeredis.FakeAsyncRedis(server=server).set("mcp:per_user_token:u:s", wire) + assert await legacy.get("u", "s") is None + assert await DualCacheTokenCacheBackend(cache, _legacy_codec(legacy)).get("u", "s") is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("applied", [False, True]) +async def test_ambiguous_write_retains_local_value_without_retry_and_peer_sees_only_applied_value( + applied: bool, +) -> None: + server: Final = fakeredis.FakeServer() + redis: Final = fake_redis_cache(server, namespace="oauth") + raw: Final = fakeredis.FakeAsyncRedis(server=server) + local: Final = DualCacheTokenCacheBackend(UserApiKeyCache(redis_cache=redis), _legacy_codec(MCPPerUserTokenCache())) + peer: Final = DualCacheTokenCacheBackend(UserApiKeyCache(redis_cache=redis), _legacy_codec(MCPPerUserTokenCache())) + + async def lost_reply(*args: object, **kwargs: object) -> None: + if applied: + assert isinstance(args[1], (str, bytes)) + await raw.set("oauth:mcp:per_user_token:u:s", args[1], ex=20) + raise ConnectionError("lost reply") + + with patch.object(redis.init_async_client(), "set", lost_reply): + await local.set("u", "s", OAuthToken("test-token"), 20) + assert await local.get("u", "s") == OAuthToken("test-token") + assert await peer.get("u", "s") == (OAuthToken("test-token") if applied else None) + + +@pytest.mark.asyncio +async def test_expiry_and_outage_preserve_local_cache_and_safe_misses(monkeypatch: pytest.MonkeyPatch) -> None: + server: Final = fakeredis.FakeServer() + redis: Final = fake_redis_cache(server) + cache: Final = UserApiKeyCache(redis_cache=redis, default_in_memory_ttl=5) + backend: Final = DualCacheTokenCacheBackend(cache, _legacy_codec(MCPPerUserTokenCache())) + await backend.set("u", "s", OAuthToken("expires"), 2) + monkeypatch.setattr("time.time", lambda: 1_800_000_003.0) + assert await backend.get("u", "s") is None + server.connected = False + await backend.set("u", "s", OAuthToken("local"), 20) + assert await backend.get("u", "s") == OAuthToken("local") + await backend.delete("u", "s") + assert await backend.get("u", "s") is None + store: Final = TransportMcpPerUserTokenStore(redis.transport()) + assert isinstance(await store.read("u", "s"), RedisFailure) + assert isinstance(await store.write("u", "s", "encrypted", ttl=timedelta(seconds=2)), RedisFailure) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("deferred", [False, True]) +@pytest.mark.parametrize("ttl", [0.9, 0.99999999]) +async def test_subsecond_ttl_keeps_local_value_but_does_not_replace_shared_entry(deferred: bool, ttl: float) -> None: + server: Final = fakeredis.FakeServer() + redis: Final = fake_redis_cache(server) + cache: Final = UserApiKeyCache(redis_cache=redis) + backend: Final = DualCacheTokenCacheBackend(cache, _legacy_codec(MCPPerUserTokenCache())) + await backend.set("u", "s", OAuthToken("old"), 30) + with request_redis_batch_scope() if deferred else nullcontext(): + await backend.set("u", "s", OAuthToken("local"), ttl) + assert await backend.get("u", "s") == OAuthToken("local") + batch: Final = active_request_redis_batch(redis) + if batch is not None: + await batch.flush() + cache.in_memory_cache.delete_cache("mcp:per_user_token:u:s") + assert await backend.get("u", "s") == OAuthToken("old") + + +@pytest.mark.asyncio +async def test_open_breaker_preserves_existing_auth_memory_fallback() -> None: + redis: Final = fake_redis_cache(fakeredis.FakeServer()) + trip_breaker(redis) + store: Final = TransportMcpPerUserTokenStore(redis.transport()) + read: Final = await store.read("u", "s") + assert isinstance(read, RedisFailure) and read.kind == "breaker_open" + cache: Final = UserApiKeyCache(redis_cache=redis) + backend: Final = DualCacheTokenCacheBackend(cache, _legacy_codec(MCPPerUserTokenCache())) + await backend.set("u", "s", OAuthToken("local"), 60) + assert await backend.get("u", "s") == OAuthToken("local") + await backend.delete("u", "s") + assert await backend.get("u", "s") is None + + +@pytest.mark.asyncio +async def test_legacy_delete_evicts_redis_local_and_peer_through_auth_invalidation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + server: Final = fakeredis.FakeServer() + redis: Final = fake_redis_cache(server, namespace="oauth") + local: Final = UserApiKeyCache(redis_cache=redis) + peer: Final = UserApiKeyCache(redis_cache=redis) + legacy: Final = MCPPerUserTokenCache() + backend: Final = DualCacheTokenCacheBackend(peer, _legacy_codec(legacy)) + monkeypatch.setattr(proxy_server, "user_api_key_cache", local) + monkeypatch.setattr(proxy_server, "redis_usage_cache", redis) + await legacy.set("u", "s", "revoke-me", 60) + assert await backend.get("u", "s") == OAuthToken("revoke-me") + pubsub: Final = fakeredis.FakeAsyncRedis(server=server).pubsub() + await pubsub.subscribe(auth_cache_invalidation_channel(redis)) + await pubsub.get_message() + await legacy.delete("u", "s") + message: Final = _MESSAGE.validate_python(await pubsub.get_message(ignore_subscribe_messages=True, timeout=1)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # redis-py pubsub replies are untyped + data: Final = message["data"] + assert isinstance(data, (str, bytes)) + assert json.loads(data) == {"cache_key": "mcp:per_user_token:u:s"} + AuthCacheInvalidationSubscriber(redis, peer)._apply_message(message) # pyright: ignore[reportPrivateUsage] # drives the subscriber without a live pubsub loop + assert await backend.get("u", "s") is None + assert await legacy.get("u", "s") is None + assert await fakeredis.FakeAsyncRedis(server=server).get("oauth:mcp:per_user_token:u:s") is None + await pubsub.aclose() # pyright: ignore[reportAttributeAccessIssue] # redis-py stubs predate PubSub.aclose From 59f09d1d2826b3b52519f70d27c2a26651442d73 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 21:56:47 +0000 Subject: [PATCH 13/17] refactor(redis): migrate user deployment affinity into a typed store --- .../redis/stores/deployment_user_affinity.py | 85 ++++++++++ .../deployment_affinity_check.py | 58 +++++-- tests/_support/fake_redis.py | 25 ++- .../stores/test_deployment_user_affinity.py | 146 ++++++++++++++++++ .../test_deployment_affinity_check.py | 10 +- 5 files changed, 305 insertions(+), 19 deletions(-) create mode 100644 litellm/redis/stores/deployment_user_affinity.py create mode 100644 tests/unit/redis/stores/test_deployment_user_affinity.py diff --git a/litellm/redis/stores/deployment_user_affinity.py b/litellm/redis/stores/deployment_user_affinity.py new file mode 100644 index 00000000000..f465b4a4f44 --- /dev/null +++ b/litellm/redis/stores/deployment_user_affinity.py @@ -0,0 +1,85 @@ +import ast +import json +from dataclasses import dataclass +from typing import Final, Protocol + +from pydantic import JsonValue, TypeAdapter, ValidationError + +from litellm.redis.transport.contract import RedisFailure, RedisTransport + +_JSON: Final = TypeAdapter[JsonValue](JsonValue) +_OBJECT: Final = TypeAdapter(object) +_CLAIM: Final = """ +local current = redis.call('GET', KEYS[1]) +if current == false then + redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2]) + return ARGV[1] +end +if current == ARGV[1] then + redis.call('EXPIRE', KEYS[1], ARGV[2]) +end +return current +""" + + +def affinity_key(model_group: str, hashed_user_key: str) -> str: + return f"deployment_affinity:v1:{model_group}:{hashed_user_key}" + + +@dataclass(frozen=True, slots=True) +class UserPin: + value: object + + +@dataclass(frozen=True, slots=True) +class UnconfirmedUserPin: + pass + + +class DeploymentUserAffinityStore(Protocol): + async def read(self, model_group: str, hashed_user_key: str) -> UserPin | RedisFailure: ... + + async def claim( + self, model_group: str, hashed_user_key: str, model_id: str, *, ttl_seconds: int + ) -> UserPin | UnconfirmedUserPin | RedisFailure: ... + + +@dataclass(frozen=True, slots=True) +class TransportDeploymentUserAffinityStore: + transport: RedisTransport + + async def read(self, model_group: str, hashed_user_key: str) -> UserPin | RedisFailure: + raw: Final = await self.transport.get(affinity_key(model_group, hashed_user_key)) + if isinstance(raw, RedisFailure): + return raw + if raw is None: + return UserPin(None) + try: + decoded: Final = raw.decode("utf-8") + try: + return UserPin(_JSON.validate_json(decoded)) + except ValidationError: + return UserPin(_OBJECT.validate_python(ast.literal_eval(decoded))) + except (UnicodeError, ValueError, SyntaxError) as error: + return RedisFailure(kind="decode", message=str(error)) + + async def claim( + self, model_group: str, hashed_user_key: str, model_id: str, *, ttl_seconds: int + ) -> UserPin | UnconfirmedUserPin | RedisFailure: + raw: Final = await self.transport.eval( + _CLAIM, + (affinity_key(model_group, hashed_user_key),), + (json.dumps({"model_id": model_id}), int(ttl_seconds)), + ) + if isinstance(raw, RedisFailure): + return raw + if not isinstance(raw, bytes): + return UnconfirmedUserPin() + try: + decoded: Final = raw.decode("utf-8") + except UnicodeError as error: + return RedisFailure(kind="decode", message=str(error)) + try: + return UserPin(_JSON.validate_json(decoded)) + except ValidationError: + return UserPin(decoded) diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index daf7d6bb345..e8e12db5148 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -28,8 +28,16 @@ from litellm.caching.affinity_cache import ( set_local_affinity_pin, ) from litellm.caching.dual_cache import DualCache +from litellm.caching.redis_batch import active_request_redis_batch from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger, Span +from litellm.redis.stores.deployment_user_affinity import ( + DeploymentUserAffinityStore, + TransportDeploymentUserAffinityStore, + UnconfirmedUserPin, + UserPin, + affinity_key, +) from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.router_utils.common_utils import deployment_ids_failed_in_request from litellm.types.llms.openai import AllMessageValues @@ -263,7 +271,7 @@ class DeploymentAffinityCheck(CustomLogger): @classmethod def get_affinity_cache_key(cls, model_group: str, user_key: str) -> str: hashed_user_key: Final = cls._hash_user_key(user_key=user_key) - return f"{cls.CACHE_KEY_PREFIX}:{model_group}:{hashed_user_key}" + return affinity_key(model_group, hashed_user_key) @classmethod def get_session_affinity_cache_key(cls, model_group: str, session_id: str, user_key: str | None) -> str: @@ -346,7 +354,7 @@ class DeploymentAffinityCheck(CustomLogger): """Deployment id held by a stored pin, for both the dict shape this writes and the bare string older writers left behind. None when the value is neither.""" if isinstance(stored, dict): - model_id: Final = stored.get("model_id") + model_id: Final = TypeAdapter(Mapping[object, object]).validate_python(stored).get("model_id") return str(model_id) if model_id is not None else None if isinstance(stored, str): return stored @@ -355,6 +363,42 @@ class DeploymentAffinityCheck(CustomLogger): def _set_local_pin(self, cache_key: str, value: object, ttl_seconds: int) -> None: set_local_affinity_pin(self.cache, cache_key, value, ttl_seconds) + async def _read_user_pin(self, model_group: str, user_key: str) -> object: + cache_key: Final = self.get_affinity_cache_key(model_group, user_key) + local: Final = TypeAdapter(object).validate_python( + await self.cache.async_get_cache(key=cache_key, local_only=True) + ) + redis_cache: Final = self.cache.redis_cache + if local is not None or redis_cache is None: + return local + batch: Final = active_request_redis_batch(redis_cache) + if batch is not None and batch.read_as_missing(cache_key): + return None + store: Final[DeploymentUserAffinityStore] = TransportDeploymentUserAffinityStore(redis_cache.transport()) + result: Final = await store.read(model_group, self._hash_user_key(user_key)) + if not isinstance(result, UserPin): + return None + if result.value is not None: + await self.cache.async_set_cache( + cache_key, result.value, local_only=True, ttl=self.cache.default_in_memory_ttl + ) + return result.value + + async def _claim_user_pin(self, model_group: str, user_key: str, model_id: str) -> str | None: + cache_key: Final = self.get_affinity_cache_key(model_group, user_key) + redis_cache: Final = self.cache.redis_cache + if redis_cache is not None: + store: Final[DeploymentUserAffinityStore] = TransportDeploymentUserAffinityStore(redis_cache.transport()) + result: Final = await store.claim( + model_group, self._hash_user_key(user_key), model_id, ttl_seconds=self.ttl_seconds + ) + if isinstance(result, UnconfirmedUserPin): + return model_id + if isinstance(result, UserPin): + self._set_local_pin(cache_key, result.value, self.ttl_seconds) + return self._pinned_model_id(result.value) + return self._claim_pin_in_memory(cache_key, DeploymentAffinityCacheValue(model_id=model_id), self.ttl_seconds) + async def _claim_pin(self, cache_key: str, pin_value: DeploymentAffinityCacheValue, ttl_seconds: int) -> str | None: winner: Final = await claim_affinity_pin(self.cache, cache_key, pin_value, ttl_seconds) return self._pinned_model_id(winner) @@ -487,8 +531,7 @@ class DeploymentAffinityCheck(CustomLogger): if user_key is None: return typed_healthy_deployments - cache_key: Final = self.get_affinity_cache_key(model_group=stable_model_map_key, user_key=user_key) - cache_result: Final = await self.cache.async_get_cache(key=cache_key) + cache_result: Final = await self._read_user_pin(stable_model_map_key, user_key) model_id: str | None = None if isinstance(cache_result, dict): @@ -660,11 +703,8 @@ class DeploymentAffinityCheck(CustomLogger): if enable_user_key and user_key is not None: try: - cache_key: Final = self.get_affinity_cache_key(model_group=deployment_model_name, user_key=user_key) - claimed_user_pin: Final = await self._claim_pin( - cache_key=cache_key, - pin_value=pin_value, - ttl_seconds=self.ttl_seconds, + claimed_user_pin: Final = await self._claim_user_pin( + deployment_model_name, user_key, pin_value["model_id"] ) if claimed_user_pin == pin_value["model_id"]: verbose_router_logger.debug( diff --git a/tests/_support/fake_redis.py b/tests/_support/fake_redis.py index 79f7e271ec6..8b98f1f5f53 100644 --- a/tests/_support/fake_redis.py +++ b/tests/_support/fake_redis.py @@ -6,6 +6,7 @@ from pydantic import TypeAdapter import litellm from litellm._service_logger import ServiceLogging +from litellm.caching.dual_cache import DualCache from litellm.caching.redis_cache import RedisCache, RedisCircuitBreaker from litellm.redis.transport.events import RedisServiceEvents from litellm.redis.transport.redis_py import RedisPyTransport @@ -61,6 +62,7 @@ def fake_redis_transport( _FLOAT: Final = TypeAdapter(float) _INT: Final = TypeAdapter(int) _KEYS: Final = TypeAdapter(list[bytes]) +_OBJECT: Final = TypeAdapter(object) def set_reply_callback(cache: RedisCache, command: str, callback: Callable[..., object]) -> None: @@ -86,7 +88,7 @@ def trip_breaker(cache: RedisCache, failures: int | None = None) -> None: def async_client(cache: RedisCache) -> fakeredis.FakeAsyncRedis: - client: Final = cache.init_async_client() # pyright: ignore[reportUnknownMemberType] # legacy RedisCache API is untyped + client: Final = cache.init_async_client() # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # legacy RedisCache API is untyped assert isinstance(client, fakeredis.FakeAsyncRedis) return client @@ -95,7 +97,7 @@ def pin_async_client(cache: RedisCache) -> Callable[[], None]: client: Final = async_client(cache) def pin() -> None: - litellm.in_memory_llm_clients_cache.set_cache( + litellm.in_memory_llm_clients_cache.set_cache( # pyright: ignore[reportUnknownMemberType] # legacy cache API is untyped key=cache._get_async_client_cache_key(), # pyright: ignore[reportPrivateUsage] # re-seeds the per-loop client slot RedisCache reads value=client, ) @@ -105,3 +107,22 @@ def pin_async_client(cache: RedisCache) -> Callable[[], None]: async def legacy_set(cache: RedisCache, key: str, value: object, *, ttl: int | None = None) -> None: await cache.async_set_cache(key, value, ttl=ttl) # pyright: ignore[reportUnknownMemberType] # legacy RedisCache API is untyped + + +def seed_async_client(cache: RedisCache, client: object) -> None: + litellm.in_memory_llm_clients_cache.set_cache( # pyright: ignore[reportUnknownMemberType] # legacy cache API is untyped + key=cache._get_async_client_cache_key(), # pyright: ignore[reportPrivateUsage] # seeds the per-loop client slot RedisCache reads + value=client, + ) + + +async def legacy_get(cache: RedisCache, key: str) -> object: + return _OBJECT.validate_python(await cache.async_get_cache(key)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # legacy RedisCache API is untyped + + +def local_set(cache: DualCache, key: str, value: object, *, ttl: int | None = None) -> None: + cache.in_memory_cache.set_cache(key, value, ttl=ttl) # pyright: ignore[reportUnknownMemberType] # legacy InMemoryCache API is untyped + + +def local_get(cache: DualCache, key: str) -> object: + return _OBJECT.validate_python(cache.in_memory_cache.get_cache(key)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # legacy InMemoryCache API is untyped diff --git a/tests/unit/redis/stores/test_deployment_user_affinity.py b/tests/unit/redis/stores/test_deployment_user_affinity.py new file mode 100644 index 00000000000..123822dc6e9 --- /dev/null +++ b/tests/unit/redis/stores/test_deployment_user_affinity.py @@ -0,0 +1,146 @@ +import hashlib +import json +from collections.abc import Iterator +from typing import Final +from unittest.mock import patch + +import fakeredis +import pytest +from redis.exceptions import TimeoutError as RedisTimeoutError + +from litellm.caching.dual_cache import DualCache +from litellm.router_utils.pre_call_checks.deployment_affinity_check import DeploymentAffinityCheck +from tests._support.fake_redis import ( + fake_redis_cache, + keys_matching, + legacy_get, + local_get, + local_set, + seed_async_client, +) + + +async def _pre_call(checker: DeploymentAffinityCheck, request: dict[str, object], call_type: None = None) -> None: + await checker.async_pre_call_deployment_hook(request, call_type) # pyright: ignore[reportUnknownMemberType] # legacy cache API is untyped + + +async def _filter( + checker: DeploymentAffinityCheck, + model: str, + deployments: list[dict[str, object]], + messages: None, + request_kwargs: dict[str, object] | None, +) -> object: + return await checker.async_filter_deployments(model, deployments, messages, request_kwargs) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # legacy hook API is untyped + + +async def _claim_user(checker: DeploymentAffinityCheck, model_group: str, user_key: str, model_id: str) -> str | None: + return await checker._claim_user_pin(model_group, user_key, model_id) # pyright: ignore[reportPrivateUsage] # exercises the pin path without a full router request + + +@pytest.fixture(autouse=True) +def fixed_clock() -> Iterator[None]: + with patch("time.time", return_value=1700000000.0): + yield + + +class LostClaimReply(fakeredis.FakeAsyncRedis): + async def eval(self, script: str, numkeys: int, *keys_and_args: str | bytes | int) -> object: + await super().eval(script, numkeys, *keys_and_args) # pyright: ignore[reportUnknownMemberType] # fakeredis is untyped + raise RedisTimeoutError("reply lost after apply") + + +def _checker(server: fakeredis.FakeServer) -> DeploymentAffinityCheck: + return DeploymentAffinityCheck( + DualCache(redis_cache=fake_redis_cache(server, namespace="affinity")), + ttl_seconds=60, + enable_user_key_affinity=True, + enable_responses_api_affinity=False, + ) + + +def _request(model_id: str) -> dict[str, object]: + return { + "metadata": {"deployment_model_name": "group", "user_api_key_hash": "caller"}, + "model_info": {"id": model_id}, + } + + +def _key() -> str: + return "affinity:deployment_affinity:v1:group:" + hashlib.sha256(b"caller").hexdigest() + + +@pytest.mark.asyncio +async def test_user_claim_wire_winner_and_ttl() -> None: + server: Final = fakeredis.FakeServer() + raw: Final = fakeredis.FakeAsyncRedis(server=server) + checker: Final = _checker(server) + await _pre_call(checker, _request("部署"), None) + assert await raw.get(_key()) == b'{"model_id": "\\u90e8\\u7f72"}' + assert 59000 <= await raw.pttl(_key()) <= 60000 + redis_cache: Final = checker.cache.redis_cache + assert redis_cache is not None + assert await legacy_get(redis_cache, _key()) == {"model_id": "部署"} + await raw.pexpire(_key(), 20000) + await _pre_call(checker, _request("other"), None) + assert await raw.get(_key()) == b'{"model_id": "\\u90e8\\u7f72"}' + assert 19000 <= await raw.pttl(_key()) <= 20000 + await _pre_call(checker, _request("部署"), None) + assert 59000 <= await raw.pttl(_key()) <= 60000 + assert await keys_matching(raw, "affinity:affinity:*") == () + + +@pytest.mark.asyncio +@pytest.mark.parametrize("wire", [b'"chosen"', b'{"model_id": "chosen"}', b"{'model_id': 'chosen'}"]) +async def test_user_read_accepts_legacy_wire_without_refresh(wire: bytes) -> None: + server: Final = fakeredis.FakeServer() + raw: Final = fakeredis.FakeAsyncRedis(server=server) + checker: Final = _checker(server) + await raw.set(_key(), wire, ex=20) + deployments: Final[list[dict[str, object]]] = [ + {"model_name": "group", "litellm_params": {"model": "openai/group"}, "model_info": {"id": "chosen"}}, + {"model_name": "group", "litellm_params": {"model": "openai/group"}, "model_info": {"id": "other"}}, + ] + result: Final = await _filter(checker, "group", deployments, None, _request("chosen")) + assert result == deployments[:1] + assert 19000 <= await raw.pttl(_key()) <= 20000 + + +@pytest.mark.asyncio +async def test_user_late_attachment_and_outage_keep_local_winner() -> None: + server: Final = fakeredis.FakeServer() + checker: Final = DeploymentAffinityCheck(DualCache(), 60, True, False) + await _pre_call(checker, _request("local"), None) + checker.cache.redis_cache = fake_redis_cache(server, namespace="affinity") + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await raw.set(_key(), json.dumps({"model_id": "shared"}), ex=60) + await _pre_call(checker, _request("local"), None) + assert local_get(checker.cache, _key().removeprefix("affinity:")) == {"model_id": "shared"} + server.connected = False + await _pre_call(checker, _request("other"), None) + assert local_get(checker.cache, _key().removeprefix("affinity:")) == {"model_id": "shared"} + + +@pytest.mark.asyncio +async def test_user_unknown_write_does_not_replace_local_winner() -> None: + server: Final = fakeredis.FakeServer() + checker: Final = _checker(server) + redis_cache: Final = checker.cache.redis_cache + assert redis_cache is not None + seed_async_client(redis_cache, LostClaimReply(server=server)) + local_set(checker.cache, _key().removeprefix("affinity:"), {"model_id": "local"}, ttl=60) + await _pre_call(checker, _request("shared"), None) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + assert await raw.get(_key()) == b'{"model_id": "shared"}' + assert local_get(checker.cache, _key().removeprefix("affinity:")) == {"model_id": "local"} + + +@pytest.mark.asyncio +async def test_user_outage_preserves_legacy_mapping_with_non_string_keys() -> None: + server: Final = fakeredis.FakeServer() + server.connected = False + checker: Final = _checker(server) + winner: Final = {"model_id": "local", 1: "legacy"} + local_set(checker.cache, _key().removeprefix("affinity:"), winner, ttl=60) + assert await _claim_user(checker, "group", "caller", "shared") == "local" + assert local_get(checker.cache, _key().removeprefix("affinity:")) == winner diff --git a/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py b/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py index fa75f980cc0..9474f1fdf23 100644 --- a/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py @@ -522,8 +522,7 @@ async def test_async_filter_deployments_uses_stable_model_map_key_for_affinity_s user_key = "user-key-abc" stable_model_map_key = "claude-sonnet-4-5@20250929" - cache = AsyncMock() - cache.async_get_cache = AsyncMock() + cache = DualCache() callback = DeploymentAffinityCheck( cache=cache, @@ -550,12 +549,7 @@ async def test_async_filter_deployments_uses_stable_model_map_key_for_affinity_s user_key=user_key, ) - async def get_cache_side_effect(*, key: str): - if key == expected_cache_key: - return {"model_id": "deployment-2"} - return None - - cache.async_get_cache.side_effect = get_cache_side_effect + await cache.async_set_cache(expected_cache_key, {"model_id": "deployment-2"}, local_only=True) filtered = await callback.async_filter_deployments( model="some-router-model-group", From 5c682ccd1e44cde3d823610bd6155fd75ba30a0e Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 21:56:48 +0000 Subject: [PATCH 14/17] refactor(redis): migrate deployment session affinity into a typed store --- .../stores/deployment_session_affinity.py | 135 +++++++++++++ .../deployment_affinity_check.py | 136 +++++++++---- tests/_support/fake_redis.py | 7 + .../test_deployment_session_affinity.py | 181 ++++++++++++++++++ .../test_session_id_affinity.py | 80 ++++---- 5 files changed, 458 insertions(+), 81 deletions(-) create mode 100644 litellm/redis/stores/deployment_session_affinity.py create mode 100644 tests/unit/redis/stores/test_deployment_session_affinity.py diff --git a/litellm/redis/stores/deployment_session_affinity.py b/litellm/redis/stores/deployment_session_affinity.py new file mode 100644 index 00000000000..94ca36683e3 --- /dev/null +++ b/litellm/redis/stores/deployment_session_affinity.py @@ -0,0 +1,135 @@ +import ast +import json +from dataclasses import dataclass +from typing import Final, Protocol + +from pydantic import JsonValue, TypeAdapter, ValidationError + +from litellm.redis.transport.contract import RedisFailure, RedisTransport + +_JSON: Final = TypeAdapter[JsonValue](JsonValue) +_OBJECT: Final = TypeAdapter(object) +_CLAIM: Final = """ +local current = redis.call('GET', KEYS[1]) +if current == false then + redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2]) + return ARGV[1] +end +if ARGV[3] then + local decoded, stored = pcall(cjson.decode, current) + local legacy = current + if decoded and type(stored) == 'string' then legacy = stored end + for _, eligible in ipairs(cjson.decode(ARGV[3])) do + if eligible['model_id'] == legacy then + redis.call('EXPIRE', KEYS[1], ARGV[2]) + return current + end + end + if decoded and type(stored) == 'table' then + for _, eligible in ipairs(cjson.decode(ARGV[3])) do + local matches = true + for key, value in pairs(eligible) do + if stored[key] ~= value then matches = false; break end + end + for key, _ in pairs(stored) do + if eligible[key] == nil then matches = false; break end + end + if matches then + redis.call('EXPIRE', KEYS[1], ARGV[2]) + return current + end + end + end + redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2]) + return ARGV[1] +end +if current == ARGV[1] then + redis.call('EXPIRE', KEYS[1], ARGV[2]) +end +return current +""" + + +def affinity_key(model_group: str, caller_scope: str, session_id: str) -> str: + return f"deployment_affinity:v1:session:{model_group}:{caller_scope}:{session_id}" + + +@dataclass(frozen=True, slots=True) +class SessionPin: + value: object + + +@dataclass(frozen=True, slots=True) +class UnconfirmedSessionPin: + pass + + +class DeploymentSessionAffinityStore(Protocol): + async def read(self, model_group: str, caller_scope: str, session_id: str) -> SessionPin | RedisFailure: ... + + async def claim( + self, + model_group: str, + caller_scope: str, + session_id: str, + model_id: str, + *, + ttl_seconds: int, + eligible_model_ids: frozenset[str] | None = None, + ) -> SessionPin | UnconfirmedSessionPin | RedisFailure: ... + + +@dataclass(frozen=True, slots=True) +class TransportDeploymentSessionAffinityStore: + transport: RedisTransport + + async def read(self, model_group: str, caller_scope: str, session_id: str) -> SessionPin | RedisFailure: + raw: Final = await self.transport.get(affinity_key(model_group, caller_scope, session_id)) + if isinstance(raw, RedisFailure): + return raw + if raw is None: + return SessionPin(None) + try: + decoded: Final = raw.decode("utf-8") + try: + return SessionPin(_JSON.validate_json(decoded)) + except ValidationError: + return SessionPin(_OBJECT.validate_python(ast.literal_eval(decoded))) + except (UnicodeError, ValueError, SyntaxError) as error: + return RedisFailure(kind="decode", message=str(error)) + + async def claim( + self, + model_group: str, + caller_scope: str, + session_id: str, + model_id: str, + *, + ttl_seconds: int, + eligible_model_ids: frozenset[str] | None = None, + ) -> SessionPin | UnconfirmedSessionPin | RedisFailure: + raw: Final = await self.transport.eval( + _CLAIM, + (affinity_key(model_group, caller_scope, session_id),), + ( + json.dumps({"model_id": model_id}), + int(ttl_seconds), + *( + (json.dumps(tuple({"model_id": value} for value in eligible_model_ids)),) + if eligible_model_ids is not None + else () + ), + ), + ) + if isinstance(raw, RedisFailure): + return raw + if not isinstance(raw, bytes): + return UnconfirmedSessionPin() + try: + decoded: Final = raw.decode("utf-8") + except UnicodeError as error: + return RedisFailure(kind="decode", message=str(error)) + try: + return SessionPin(_JSON.validate_json(decoded)) + except ValidationError: + return SessionPin(decoded) diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index e8e12db5148..6a9e6a984a7 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -23,7 +23,6 @@ from litellm._internal_context import with_service_target from litellm._logging import verbose_router_logger from litellm.caching.affinity_cache import ( ROUTER_SESSION_PINS_TARGET, - claim_affinity_pin, claim_affinity_pin_in_memory, set_local_affinity_pin, ) @@ -31,6 +30,15 @@ from litellm.caching.dual_cache import DualCache from litellm.caching.redis_batch import active_request_redis_batch from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger, Span +from litellm.redis.stores.deployment_session_affinity import ( + DeploymentSessionAffinityStore, + SessionPin, + TransportDeploymentSessionAffinityStore, + UnconfirmedSessionPin, +) +from litellm.redis.stores.deployment_session_affinity import ( + affinity_key as session_affinity_key, +) from litellm.redis.stores.deployment_user_affinity import ( DeploymentUserAffinityStore, TransportDeploymentUserAffinityStore, @@ -280,7 +288,7 @@ class DeploymentAffinityCheck(CustomLogger): `"unscoped"` covers direct Router usage with no authenticated caller, matching the complexity router's own session pin key.""" hashed_user_key: Final = cls._hash_user_key(user_key) if user_key is not None else "unscoped" - return f"{cls.CACHE_KEY_PREFIX}:session:{model_group}:{hashed_user_key}:{session_id}" + return session_affinity_key(model_group, hashed_user_key, session_id) @staticmethod def _get_session_id_from_metadata_dict(metadata: Mapping[object, object]) -> str | None: @@ -399,8 +407,45 @@ class DeploymentAffinityCheck(CustomLogger): return self._pinned_model_id(result.value) return self._claim_pin_in_memory(cache_key, DeploymentAffinityCacheValue(model_id=model_id), self.ttl_seconds) - async def _claim_pin(self, cache_key: str, pin_value: DeploymentAffinityCacheValue, ttl_seconds: int) -> str | None: - winner: Final = await claim_affinity_pin(self.cache, cache_key, pin_value, ttl_seconds) + async def _claim_session_pin( + self, + model_group: str, + session_id: str, + user_key: str | None, + model_id: str, + ttl_seconds: int, + *, + eligible_model_ids: frozenset[str] | None = None, + ) -> str | None: + cache_key: Final = self.get_session_affinity_cache_key(model_group, session_id, user_key) + redis_cache: Final = self.cache.redis_cache + if redis_cache is not None: + store: Final[DeploymentSessionAffinityStore] = TransportDeploymentSessionAffinityStore( + redis_cache.transport() + ) + result: Final = await store.claim( + model_group, + self._hash_user_key(user_key) if user_key is not None else "unscoped", + session_id, + model_id, + ttl_seconds=ttl_seconds, + eligible_model_ids=eligible_model_ids, + ) + if isinstance(result, UnconfirmedSessionPin): + return model_id + if isinstance(result, SessionPin): + self._set_local_pin(cache_key, result.value, ttl_seconds) + return self._pinned_model_id(result.value) + winner: Final = claim_affinity_pin_in_memory( + self.cache, + cache_key, + DeploymentAffinityCacheValue(model_id=model_id), + ttl_seconds, + eligible_values=tuple(DeploymentAffinityCacheValue(model_id=value) for value in eligible_model_ids) + if eligible_model_ids is not None + else None, + preserve_legacy_model_ids=eligible_model_ids is not None, + ) return self._pinned_model_id(winner) def _claim_pin_in_memory( @@ -487,17 +532,12 @@ class DeploymentAffinityCheck(CustomLogger): if session_affinity_active: session_id: Final = self._get_session_id_from_request_kwargs(request_kwargs=request_kwargs) if session_id is not None: - session_cache_key: Final = self.get_session_affinity_cache_key( - model_group=stable_model_map_key, session_id=session_id, user_key=user_key - ) - session_cache_result: Final = ( - await self._read_shared_session_pin( - session_cache_key, - self._get_marker_session_affinity_ttl(request_kwargs) or self.ttl_seconds, - parent_otel_span, - ) - if self.session_affinity_repin_on_failover - else await self.cache.async_get_cache(key=session_cache_key) + session_cache_result: Final = await self._read_session_pin( + stable_model_map_key, + session_id, + user_key, + self._get_marker_session_affinity_ttl(request_kwargs) or self.ttl_seconds, + parent_otel_span, ) session_model_id: str | None = None @@ -561,18 +601,36 @@ class DeploymentAffinityCheck(CustomLogger): ) return [deployment] - async def _read_shared_session_pin(self, cache_key: str, ttl_seconds: int, parent_otel_span: Span | None) -> object: + async def _read_session_pin( + self, + model_group: str, + session_id: str, + user_key: str | None, + ttl_seconds: int, + parent_otel_span: Span | None, + ) -> object: + cache_key: Final = self.get_session_affinity_cache_key(model_group, session_id, user_key) redis_cache: Final = self.cache.redis_cache - if redis_cache is None: - return TypeAdapter(object).validate_python(await self.cache.async_get_cache(key=cache_key)) + local: Final = TypeAdapter(object).validate_python( + await self.cache.async_get_cache(key=cache_key, local_only=True) + ) + if redis_cache is None or (not self.session_affinity_repin_on_failover and local is not None): + return local local_before: Final = TypeAdapter(object).validate_python(self.cache.in_memory_cache.get_cache(cache_key)) - try: - shared: Final[object] = await redis_cache.async_get_cache( - key=cache_key, parent_otel_span=parent_otel_span, raise_on_error=True - ) - except Exception: # noqa: BLE001 # Optional affinity falls back locally on Redis admission or read failure - verbose_router_logger.debug("Session pin read failed, using pod-local affinity") + batch: Final = active_request_redis_batch(redis_cache) + if not self.session_affinity_repin_on_failover and batch is not None and batch.read_as_missing(cache_key): + return None + store: Final[DeploymentSessionAffinityStore] = TransportDeploymentSessionAffinityStore(redis_cache.transport()) + result: Final = await store.read( + model_group, self._hash_user_key(user_key) if user_key is not None else "unscoped", session_id + ) + if not isinstance(result, SessionPin): return TypeAdapter(object).validate_python(await self.cache.async_get_cache(key=cache_key, local_only=True)) + shared: Final = result.value + if not self.session_affinity_repin_on_failover: + if shared is not None: + await self.cache.async_set_cache(key=cache_key, value=shared, local_only=True) + return shared if self.cache.in_memory_cache.get_cache(cache_key) is local_before: if shared is None: self.cache.in_memory_cache.delete_cache(cache_key) @@ -614,16 +672,13 @@ class DeploymentAffinityCheck(CustomLogger): model_id: Final = model_info.get("id") if model_id is None or str(model_id) not in eligible_deployment_ids: return - cache_key: Final = self.get_session_affinity_cache_key( - model_group, session_id, self.get_user_key_from_request_kwargs(request_kwargs) - ) - await claim_affinity_pin( - self.cache, - cache_key, - DeploymentAffinityCacheValue(model_id=str(model_id)), + await self._claim_session_pin( + model_group, + session_id, + self.get_user_key_from_request_kwargs(request_kwargs), + str(model_id), marker_ttl if marker_ttl is not None else self.ttl_seconds, - eligible_values=tuple(DeploymentAffinityCacheValue(model_id=value) for value in eligible_deployment_ids), - preserve_legacy_model_ids=True, + eligible_model_ids=eligible_deployment_ids, ) async def async_pre_call_deployment_hook( @@ -733,19 +788,18 @@ class DeploymentAffinityCheck(CustomLogger): if session_id is not None: try: session_affinity_ttl: Final = marker_session_ttl if marker_session_ttl is not None else self.ttl_seconds - session_cache_key: Final = self.get_session_affinity_cache_key( - model_group=deployment_model_name, session_id=session_id, user_key=user_key - ) - claimed_session_pin: Final = await self._claim_pin( - cache_key=session_cache_key, - pin_value=pin_value, + claimed_session_pin: Final = await self._claim_session_pin( + model_group=deployment_model_name, + session_id=session_id, + user_key=user_key, + model_id=str(model_id), ttl_seconds=session_affinity_ttl, ) if claimed_session_pin == pin_value["model_id"]: verbose_router_logger.debug( "DeploymentAffinityCheck: set session affinity mapping model_map_key=%s deployment=%s ttl=%s session_id=%s", deployment_model_name, - model_id, + pin_value["model_id"], session_affinity_ttl, session_id, ) @@ -754,7 +808,7 @@ class DeploymentAffinityCheck(CustomLogger): "DeploymentAffinityCheck: session pin already claimed model_map_key=%s existing=%s ours=%s session_id=%s", deployment_model_name, claimed_session_pin, - model_id, + pin_value["model_id"], session_id, ) except Exception as e: diff --git a/tests/_support/fake_redis.py b/tests/_support/fake_redis.py index 8b98f1f5f53..1b0328ef99a 100644 --- a/tests/_support/fake_redis.py +++ b/tests/_support/fake_redis.py @@ -3,6 +3,7 @@ from typing import Final, TypeVar import fakeredis from pydantic import TypeAdapter +from redis.asyncio import Redis import litellm from litellm._service_logger import ServiceLogging @@ -93,6 +94,12 @@ def async_client(cache: RedisCache) -> fakeredis.FakeAsyncRedis: return client +def raw_client(cache: RedisCache) -> "Redis[bytes]": + client: Final = cache.init_async_client() # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # legacy RedisCache API is untyped + assert isinstance(client, Redis) + return client # pyright: ignore[reportUnknownVariableType] # redis-py erases the reply type at runtime + + def pin_async_client(cache: RedisCache) -> Callable[[], None]: client: Final = async_client(cache) diff --git a/tests/unit/redis/stores/test_deployment_session_affinity.py b/tests/unit/redis/stores/test_deployment_session_affinity.py new file mode 100644 index 00000000000..ed4a5b71a70 --- /dev/null +++ b/tests/unit/redis/stores/test_deployment_session_affinity.py @@ -0,0 +1,181 @@ +import hashlib +import json +from collections.abc import Callable, Iterator +from typing import Final +from unittest.mock import patch + +import fakeredis +import pytest +from redis.exceptions import TimeoutError as RedisTimeoutError + +from litellm.caching.dual_cache import DualCache +from litellm.router_utils.pre_call_checks.deployment_affinity_check import DeploymentAffinityCheck +from tests._support.fake_redis import ( + fake_redis_cache, + keys_matching, + local_get, + local_set, + seed_async_client, +) + + +async def _pre_call(checker: DeploymentAffinityCheck, request: dict[str, object], call_type: None = None) -> None: + await checker.async_pre_call_deployment_hook(request, call_type) # pyright: ignore[reportUnknownMemberType] # legacy cache API is untyped + + +async def _read_session( + checker: DeploymentAffinityCheck, + model_group: str, + session_id: str, + user_key: str | None, + ttl_seconds: int, + parent_otel_span: None, +) -> object: + return await checker._read_session_pin(model_group, session_id, user_key, ttl_seconds, parent_otel_span) # pyright: ignore[reportPrivateUsage] # exercises the pin path without a full router request + + +@pytest.fixture(autouse=True) +def fixed_clock() -> Iterator[None]: + with patch("time.time", return_value=1700000000.0): + yield + + +def _checker(server: fakeredis.FakeServer, *, repin: bool = True) -> DeploymentAffinityCheck: + return DeploymentAffinityCheck( + DualCache(redis_cache=fake_redis_cache(server, namespace="session")), + 60, + False, + False, + enable_session_id_affinity=True, + session_affinity_repin_on_failover=repin, + ) + + +def _key(caller: str = "caller") -> str: + return "session:deployment_affinity:v1:session:group:" + hashlib.sha256(caller.encode()).hexdigest() + ":same" + + +def _request(model_id: str, caller: str = "caller") -> dict[str, object]: + return { + "model_info": {"id": model_id}, + "metadata": { + "deployment_model_name": "group", + "model_group": "group", + "session_id": "same", + "user_api_key_hash": caller, + }, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("repin", [False, True]) +@pytest.mark.parametrize("wire", [b'"chosen"', b'{"model_id": "chosen"}', b"{'model_id': 'chosen'}"]) +async def test_session_reads_legacy_values_without_refresh(wire: bytes, repin: bool) -> None: + server: Final = fakeredis.FakeServer() + checker: Final = _checker(server, repin=repin) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await raw.set(_key(), wire, ex=20) + assert await _read_session(checker, "group", "same", "caller", 60, None) in ("chosen", {"model_id": "chosen"}) + assert await raw.get(_key()) == wire + assert await raw.pttl(_key()) == 20000 + + +@pytest.mark.asyncio +async def test_session_claim_bytes_scope_and_claim_time_ttl() -> None: + server: Final = fakeredis.FakeServer() + checker: Final = _checker(server) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await _pre_call(checker, _request("部署"), None) + assert await raw.get(_key()) == b'{"model_id": "\\u90e8\\u7f72"}' + assert await raw.pttl(_key()) == 60000 + await raw.pexpire(_key(), 20000) + await _pre_call(checker, _request("other"), None) + assert await raw.pttl(_key()) == 20000 + await _pre_call(checker, _request("部署"), None) + assert await raw.pttl(_key()) == 60000 + await _pre_call(checker, _request("other", "second"), None) + assert await raw.get(_key("second")) == b'{"model_id": "other"}' + assert await keys_matching(raw, "session:session:*") == () + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "wire,expected", + [ + (b'{"model_id": "chosen"}', b'{"model_id": "chosen"}'), + (b'"chosen"', b'"chosen"'), + (b"chosen", b"chosen"), + (b'{"model_id": "removed"}', b'{"model_id": "replacement"}'), + (b'{"model_id": "chosen", "extra": true}', b'{"model_id": "replacement"}'), + (b"invalid{", b'{"model_id": "replacement"}'), + ], +) +async def test_success_claim_preserves_only_eligible_complete_values(wire: bytes, expected: bytes) -> None: + server: Final = fakeredis.FakeServer() + checker: Final = _checker(server) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await raw.set(_key(), wire, ex=20) + await checker.async_repin_session_after_success( + _request("replacement"), "group", frozenset(("chosen", "replacement")) + ) + assert await raw.get(_key()) == expected + assert await raw.pttl(_key()) == 60000 + + +@pytest.mark.asyncio +async def test_shared_miss_clears_local_while_outage_retains_it() -> None: + server: Final = fakeredis.FakeServer() + checker: Final = _checker(server) + local_key: Final = _key().removeprefix("session:") + local_set(checker.cache, local_key, {"model_id": "stale"}, ttl=60) + assert await _read_session(checker, "group", "same", "caller", 60, None) is None + assert local_get(checker.cache, local_key) is None + local_set(checker.cache, local_key, {"model_id": "local"}, ttl=60) + server.connected = False + assert await _read_session(checker, "group", "same", "caller", 60, None) == {"model_id": "local"} + + +class ConcurrentRead(fakeredis.FakeAsyncRedis): + after_read: Callable[[], None] + + async def get(self, name: str) -> object: # pyright: ignore[reportIncompatibleMethodOverride] # narrows fakeredis for the test + result = await super().get(name) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # fakeredis is untyped + self.after_read() + return result # pyright: ignore[reportUnknownVariableType] # fakeredis is untyped + + +@pytest.mark.asyncio +async def test_shared_read_does_not_overwrite_a_newer_local_claim() -> None: + server: Final = fakeredis.FakeServer() + checker: Final = _checker(server) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await raw.set(_key(), json.dumps({"model_id": "shared"}), ex=20) + local_key: Final = _key().removeprefix("session:") + client: Final = ConcurrentRead(server=server) + client.after_read = lambda: local_set(checker.cache, local_key, {"model_id": "newer"}, ttl=60) + redis_cache: Final = checker.cache.redis_cache + assert redis_cache is not None + seed_async_client(redis_cache, client) + assert await _read_session(checker, "group", "same", "caller", 60, None) == {"model_id": "shared"} + assert local_get(checker.cache, local_key) == {"model_id": "newer"} + + +class LostClaimReply(fakeredis.FakeAsyncRedis): + async def eval(self, script: str, numkeys: int, *keys_and_args: str | bytes | int) -> object: + await super().eval(script, numkeys, *keys_and_args) # pyright: ignore[reportUnknownMemberType] # fakeredis is untyped + raise RedisTimeoutError("reply lost after apply") + + +@pytest.mark.asyncio +async def test_session_unknown_write_retains_local_continuity() -> None: + server: Final = fakeredis.FakeServer() + checker: Final = _checker(server) + redis_cache: Final = checker.cache.redis_cache + assert redis_cache is not None + seed_async_client(redis_cache, LostClaimReply(server=server)) + local_key: Final = _key().removeprefix("session:") + local_set(checker.cache, local_key, {"model_id": "local"}, ttl=60) + await _pre_call(checker, _request("shared"), None) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + assert await raw.get(_key()) == b'{"model_id": "shared"}' + assert local_get(checker.cache, local_key) == {"model_id": "local"} diff --git a/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py b/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py index 6b4ea74894a..d317f0d5741 100644 --- a/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py +++ b/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py @@ -1,8 +1,8 @@ import asyncio import json -from typing import Final from collections.abc import AsyncIterator, Mapping from datetime import datetime, timezone +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -19,6 +19,7 @@ from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESS from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( DeploymentAffinityCheck, ) +from tests._support.fake_redis import async_client, fake_redis_cache, raw_client, seed_async_client def _repinning_router(enabled: bool | None = True) -> litellm.Router: @@ -121,6 +122,17 @@ async def test_unsuccessful_responses_do_not_repin(status: str, wrapped: bool) - router.discard() +async def _claim_session( + callback: DeploymentAffinityCheck, + model_group: str, + session_id: str, + user_key: str | None, + model_id: str, + ttl_seconds: int, +) -> object: + return await callback._claim_session_pin(model_group, session_id, user_key, model_id, ttl_seconds) # pyright: ignore[reportPrivateUsage] # exercises the pin path without a full router request + + async def test_concurrent_successful_replacements_keep_first_winner() -> None: router: Final = _repinning_router() try: @@ -288,8 +300,8 @@ async def test_delayed_shared_pin_read_preserves_a_newer_local_write( await resume.wait() return value - reader: Final = DelayedReader(connection_pool=redis.init_async_client().connection_pool) - litellm.in_memory_llm_clients_cache.set_cache(key=redis._get_async_client_cache_key(), value=reader) + reader: Final = DelayedReader(connection_pool=raw_client(redis).connection_pool) + seed_async_client(redis, reader) pending: Final = asyncio.create_task(_selected_session_ids(router)) try: await observed.wait() @@ -817,12 +829,12 @@ async def test_in_memory_reclaim_slides_idle_window_only_for_the_stored_deployme cache.in_memory_cache.set_cache(pin_key, stored_pin, ttl=10) first_expiry = cache.in_memory_cache.ttl_dict[pin_key] - reclaimed = await callback._claim_pin(cache_key=pin_key, pin_value={"model_id": "deployment-1"}, ttl_seconds=777) + reclaimed = await _claim_session(callback, "target-group", "slide-session", "key-1", "deployment-1", 777) assert reclaimed == "deployment-1" assert cache.in_memory_cache.ttl_dict[pin_key] == pytest.approx(time_module.time() + 777, abs=5) assert cache.in_memory_cache.ttl_dict[pin_key] > first_expiry - lost = await callback._claim_pin(cache_key=pin_key, pin_value={"model_id": "deployment-2"}, ttl_seconds=10) + lost = await _claim_session(callback, "target-group", "slide-session", "key-1", "deployment-2", 10) assert lost == "deployment-1" assert cache.in_memory_cache.ttl_dict[pin_key] == pytest.approx(time_module.time() + 777, abs=5) @@ -840,32 +852,21 @@ async def test_claim_pin_uses_redis_attached_after_construction(): enable_responses_api_affinity=False, ) - captured = {} - - async def fake_runner(keys, args, client=None): - captured["keys"] = keys - captured["args"] = args - return b'{"model_id": "other-pod-winner"}' - - late_redis = MagicMock() - late_redis.async_register_script = MagicMock(return_value=fake_runner) + server = FakeServer() + late_redis = fake_redis_cache(server, namespace="late") cache.redis_cache = late_redis import time as time_module pin_key = _session_pin_key("late-redis-session", "key-1") + await late_redis.async_set_cache(pin_key, {"model_id": "other-pod-winner"}, ttl=10) cache.in_memory_cache.set_cache(pin_key, {"model_id": "other-pod-winner"}, ttl=10) - claimed = await callback._claim_pin( - cache_key=pin_key, - pin_value={"model_id": "our-deployment"}, - ttl_seconds=777, - ) + claimed = await _claim_session(callback, "target-group", "late-redis-session", "key-1", "our-deployment", 777) assert claimed == "other-pod-winner" assert cache.in_memory_cache.ttl_dict[pin_key] == pytest.approx(time_module.time() + 777, abs=5) - assert captured["keys"] == (pin_key,) - assert captured["args"] == ('{"model_id": "our-deployment"}', 777) + assert await late_redis.async_get_cache(pin_key) == {"model_id": "other-pod-winner"} assert cache.in_memory_cache.get_cache(_session_pin_key("late-redis-session", "key-1")) == { "model_id": "other-pod-winner" } @@ -885,20 +886,18 @@ async def test_claim_pin_falls_back_to_pod_local_when_redis_is_down(): enable_responses_api_affinity=False, ) - async def exploding_runner(keys, args, client=None): - raise ConnectionError("redis is down") - - down_redis = MagicMock() - down_redis.async_register_script = MagicMock(return_value=exploding_runner) + server = FakeServer() + server.connected = False + down_redis = fake_redis_cache(server, namespace="down") cache.redis_cache = down_redis key = _session_pin_key("outage-session", "key-1") - claimed = await callback._claim_pin(cache_key=key, pin_value={"model_id": "our-deployment"}, ttl_seconds=777) + claimed = await _claim_session(callback, "target-group", "outage-session", "key-1", "our-deployment", 777) assert claimed == "our-deployment" assert cache.in_memory_cache.get_cache(key) == {"model_id": "our-deployment"} - second = await callback._claim_pin(cache_key=key, pin_value={"model_id": "another-deployment"}, ttl_seconds=777) + second = await _claim_session(callback, "target-group", "outage-session", "key-1", "another-deployment", 777) assert second == "our-deployment" @@ -979,14 +978,15 @@ async def test_legacy_deployment_claim_retains_decoder_and_keepalive( enable_user_key_affinity=False, enable_responses_api_affinity=False, ) - cache.in_memory_cache.set_cache("deployment-pin", stored, ttl=10) + pin_key: Final = callback.get_session_affinity_cache_key("target-group", "decoder", None) + cache.in_memory_cache.set_cache(pin_key, stored, ttl=10) clock.return_value = 105.0 - winner: Final = await callback._claim_pin("deployment-pin", {"model_id": "7"}, 30) + winner: Final = await _claim_session(callback, "target-group", "decoder", None, "7", 30) assert winner == expected - assert cache.in_memory_cache.ttl_dict["deployment-pin"] == (135.0 if refresh else 110.0) - assert cache.in_memory_cache.get_cache("deployment-pin") == ({"model_id": "7"} if refresh else stored) + assert cache.in_memory_cache.ttl_dict[pin_key] == (135.0 if refresh else 110.0) + assert cache.in_memory_cache.get_cache(pin_key) == ({"model_id": "7"} if refresh else stored) @pytest.mark.asyncio @@ -998,16 +998,13 @@ async def test_legacy_deployment_claim_retains_decoder_and_keepalive( ("winner", "winner", "winner"), (b"winner", "winner", "winner"), ('{"model": "winner"}', None, {"model": "winner"}), - (None, "candidate", None), - (123, "candidate", None), - ({"model_id": "winner"}, "candidate", None), + (None, "candidate", {"model_id": "candidate"}), ], ) async def test_redis_deployment_claim_preserves_legacy_result_decoding( - raw: object, expected: str | None, stored: object + raw: str | bytes | None, expected: str | None, stored: object ) -> None: - redis: Final = MagicMock() - redis.async_register_script.return_value = AsyncMock(return_value=raw) + redis: Final = fake_redis_cache(FakeServer(), namespace="decode") cache: Final = DualCache(redis_cache=redis) callback: Final = DeploymentAffinityCheck( cache=cache, @@ -1016,10 +1013,13 @@ async def test_redis_deployment_claim_preserves_legacy_result_decoding( enable_responses_api_affinity=False, ) - winner: Final = await callback._claim_pin("deployment-pin", {"model_id": "candidate"}, 30) + pin_key: Final = callback.get_session_affinity_cache_key("target-group", "decoder", None) + if raw is not None: + await async_client(redis).set(redis.check_and_fix_namespace(pin_key), raw, ex=10) + winner: Final = await _claim_session(callback, "target-group", "decoder", None, "candidate", 30) assert winner == expected - assert cache.in_memory_cache.get_cache("deployment-pin") == stored + assert cache.in_memory_cache.get_cache(pin_key) == stored @pytest.mark.asyncio From 24964fdea057376768b206d9d039d126f01903fd Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 21:56:49 +0000 Subject: [PATCH 15/17] refactor(redis): migrate complexity session affinity into a typed store --- .../stores/complexity_session_affinity.py | 82 +++++++ .../complexity_router/complexity_router.py | 109 +++++++-- .../test_complexity_session_affinity.py | 230 ++++++++++++++++++ 3 files changed, 397 insertions(+), 24 deletions(-) create mode 100644 litellm/redis/stores/complexity_session_affinity.py create mode 100644 tests/unit/redis/stores/test_complexity_session_affinity.py diff --git a/litellm/redis/stores/complexity_session_affinity.py b/litellm/redis/stores/complexity_session_affinity.py new file mode 100644 index 00000000000..703b6b27c96 --- /dev/null +++ b/litellm/redis/stores/complexity_session_affinity.py @@ -0,0 +1,82 @@ +import json +from collections.abc import Mapping +from dataclasses import dataclass +from datetime import timedelta +from typing import Final, Protocol + +from pydantic import JsonValue, TypeAdapter, ValidationError + +from litellm.redis.transport.contract import RedisFailure, RedisTransport + +_JSON: Final = TypeAdapter[JsonValue](JsonValue) +_READ: Final = """ +local current = redis.call('GET', KEYS[1]) +if current ~= false then + redis.call('EXPIRE', KEYS[1], ARGV[1]) +end +return current +""" + + +def affinity_key(model_name: str, caller_scope: str, session_id: str) -> str: + return f"complexity_router_session_affinity:v1:{model_name}:{caller_scope}:{session_id}" + + +def pin_value(model: str, tier: str | None) -> Mapping[str, str | None]: + return {"model": model, "tier": tier} + + +@dataclass(frozen=True, slots=True) +class ComplexitySessionPin: + value: object + + +class ComplexitySessionAffinityStore(Protocol): + async def read_and_refresh( + self, model_name: str, caller_scope: str, session_id: str, *, ttl_seconds: int + ) -> ComplexitySessionPin | RedisFailure: ... + + async def overwrite( + self, model_name: str, caller_scope: str, session_id: str, model: str, tier: str | None, *, ttl_seconds: int + ) -> None | RedisFailure: ... + + +@dataclass(frozen=True, slots=True) +class TransportComplexitySessionAffinityStore: + transport: RedisTransport + coalesce_reads: bool = False + + async def read_and_refresh( + self, model_name: str, caller_scope: str, session_id: str, *, ttl_seconds: int + ) -> ComplexitySessionPin | RedisFailure: + key: Final = affinity_key(model_name, caller_scope, session_id) + batch: Final = self.transport.request_batch() if self.coalesce_reads else None + raw: Final = ( + await batch.eval(_READ, (key,), (int(ttl_seconds),)) + if batch is not None + else await self.transport.eval(_READ, (key,), (int(ttl_seconds),)) + ) + if isinstance(raw, RedisFailure): + return raw + if not isinstance(raw, bytes): + return ComplexitySessionPin(None) + try: + decoded: Final = raw.decode("utf-8") + except UnicodeError as error: + return RedisFailure(kind="decode", message=str(error)) + try: + return ComplexitySessionPin(_JSON.validate_json(decoded)) + except ValidationError: + return ComplexitySessionPin(decoded) + + async def overwrite( + self, model_name: str, caller_scope: str, session_id: str, model: str, tier: str | None, *, ttl_seconds: int + ) -> None | RedisFailure: + result: Final = await self.transport.set_seconds( + affinity_key(model_name, caller_scope, session_id), + json.dumps(pin_value(model, tier)), + ttl=timedelta(seconds=int(ttl_seconds)), + ) + if isinstance(result, RedisFailure): + return result + return None if result else RedisFailure(kind="decode", message="Unexpected session overwrite reply") diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 30f1021251f..a6adc14dfd9 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -36,7 +36,8 @@ from typing_extensions import ReadOnly, TypedDict from litellm._internal_context import with_service_target from litellm._logging import verbose_router_logger -from litellm.caching.affinity_cache import ROUTER_SESSION_PINS_TARGET, claim_affinity_pin, read_affinity_pin +from litellm.caching.affinity_cache import ROUTER_SESSION_PINS_TARGET, claim_affinity_pin, set_local_affinity_pin +from litellm.caching.redis_batch import active_request_redis_batch from litellm.constants import ( COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY, EMPTY_MAPPING, @@ -64,6 +65,17 @@ from litellm.llms.anthropic.common_utils import is_claude_code_user_agent from litellm.llms.base_llm.base_utils import type_to_response_format_param from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.llms.databricks.decisions.transformation import DATABRICKS_DECISIONS_CONFIG +from litellm.redis.stores.complexity_session_affinity import ( + ComplexitySessionAffinityStore, + ComplexitySessionPin, + TransportComplexitySessionAffinityStore, +) +from litellm.redis.stores.complexity_session_affinity import ( + affinity_key as session_affinity_key, +) +from litellm.redis.stores.complexity_session_affinity import ( + pin_value as session_pin_value, +) from litellm.router_strategy.adaptive_router.classifier import classify_prompt from litellm.router_strategy.complexity_router.context_compaction import compaction_pending from litellm.router_strategy.complexity_router.tier_predictor import ( @@ -1388,7 +1400,7 @@ def _parse_session_affinity_pin(value: object, active_tiers: tuple[str, ...]) -> def _session_affinity_cache_value(model: str, tier: ComplexityTier | str | None) -> Mapping[str, str | None]: tier_value: Final = _tier_name(tier) if tier is not None else None - return {"model": model, "tier": tier_value} + return session_pin_value(model, tier_value) @dataclass(frozen=True, slots=True) @@ -4412,7 +4424,7 @@ class ComplexityRouter(CustomLogger): return self._get_session_id_from_request_kwargs(request_kwargs) @staticmethod - def _get_user_api_key_hash_from_request_kwargs(request_kwargs: dict) -> str | None: + def _get_user_api_key_hash_from_request_kwargs(request_kwargs: Mapping[str, object]) -> str | None: """Resolve the proxy-derived API key hash, the same trust boundary DeploymentAffinityCheck uses for its own key-based affinity (not the client-supplied OpenAI `user` param, which isn't authenticated).""" @@ -4422,13 +4434,69 @@ class ComplexityRouter(CustomLogger): return str(user_key) return None - def _get_session_affinity_cache_key(self, session_id: str, request_kwargs: dict) -> str: + def _get_session_affinity_cache_key(self, session_id: str, request_kwargs: Mapping[str, object]) -> str: # Namespace by the caller's API key hash so two different callers reusing the # same client-supplied session_id can't poison each other's routing pin. Falls # back to "unscoped" only when there's no authenticated caller to scope by # (e.g. direct Router usage without the proxy layer). caller_scope: Final = self._get_user_api_key_hash_from_request_kwargs(request_kwargs) or "unscoped" - return f"complexity_router_session_affinity:v1:{self.model_name}:{caller_scope}:{session_id}" + return session_affinity_key(self.model_name, caller_scope, session_id) + + async def _read_session_affinity_pin(self, session_id: str, request_kwargs: Mapping[str, object]) -> object: + cache: Final = self.litellm_router_instance.cache + cache_key: Final = self._get_session_affinity_cache_key(session_id, request_kwargs) + ttl: Final = self.config.session_affinity_ttl_seconds + redis_cache: Final = cache.redis_cache + if redis_cache is not None: + batch: Final = active_request_redis_batch(redis_cache) + store: Final[ComplexitySessionAffinityStore] = TransportComplexitySessionAffinityStore( + redis_cache.transport(), coalesce_reads=batch is not None and batch.pending > 0 + ) + result: Final = await store.read_and_refresh( + self.model_name, + self._get_user_api_key_hash_from_request_kwargs(request_kwargs) or "unscoped", + session_id, + ttl_seconds=ttl, + ) + if isinstance(result, ComplexitySessionPin): + if result.value is None: + cache.in_memory_cache.delete_cache(cache_key) + else: + set_local_affinity_pin(cache, cache_key, result.value, ttl) + return result.value + local: Final = TypeAdapter(object).validate_python(cache.in_memory_cache.get_cache(cache_key)) + if local is not None: + set_local_affinity_pin(cache, cache_key, local, ttl) + return local + + async def _store_session_affinity_pin( + self, + session_id: str, + request_kwargs: Mapping[str, object], + model: str, + tier: ComplexityTier | str | None, + ) -> None: + cache: Final = self.litellm_router_instance.cache + ttl: Final = self.config.session_affinity_ttl_seconds + await cache.async_set_cache( + key=self._get_session_affinity_cache_key(session_id, request_kwargs), + value=_session_affinity_cache_value(model, tier), + ttl=ttl, + local_only=True, + ) + redis_cache: Final = cache.redis_cache + if redis_cache is not None: + store: Final[ComplexitySessionAffinityStore] = TransportComplexitySessionAffinityStore( + redis_cache.transport() + ) + await store.overwrite( + self.model_name, + self._get_user_api_key_hash_from_request_kwargs(request_kwargs) or "unscoped", + session_id, + model, + _tier_name(tier) if tier is not None else None, + ttl_seconds=ttl, + ) @property def _uses_tier_pin(self) -> bool: @@ -4456,7 +4524,7 @@ class ComplexityRouter(CustomLogger): async def async_pre_routing_hook( self, model: str, - request_kwargs: dict, + request_kwargs: dict[str, object], messages: list[dict[str, object]] | None = None, input: str | list | None = None, specific_deployment: bool | None = False, @@ -4496,7 +4564,6 @@ class ComplexityRouter(CustomLogger): use_session_affinity: Final = self._uses_tier_pin session_id: Final = self._get_tier_pin_session_id(request_kwargs) if use_session_affinity else None - cache_key = self._get_session_affinity_cache_key(session_id, request_kwargs) if session_id is not None else None # In 'user_turn' mode a held pin is replayed only on continuation turns; a new human # ask falls through and re-classifies. session_affinity restores pin-first for asks too. @@ -4507,10 +4574,8 @@ class ComplexityRouter(CustomLogger): or codex_compaction ) - if cache_key is not None and pin_replay_allowed: - pinned_value: Final = await read_affinity_pin( - self.litellm_router_instance.cache, cache_key, self.config.session_affinity_ttl_seconds - ) + if session_id is not None and pin_replay_allowed: + pinned_value: Final = await self._read_session_affinity_pin(session_id, request_kwargs) pinned_pin: Final = _parse_session_affinity_pin(pinned_value, self.config.tier_names()) if pinned_pin is not None: user_message: Final = ( @@ -4608,10 +4673,8 @@ class ComplexityRouter(CustomLogger): else floor_model ) if escalated: - await self.litellm_router_instance.cache.async_set_cache( - key=cache_key, - value=_session_affinity_cache_value(session_model, resolved_pin_tier), - ttl=self.config.session_affinity_ttl_seconds, + await self._store_session_affinity_pin( + session_id, request_kwargs, session_model, resolved_pin_tier ) if self.config.adaptive: from litellm.router_strategy.adaptive_router.config import ( @@ -4705,19 +4768,17 @@ class ComplexityRouter(CustomLogger): # the hard floor constrained its pick, so pinning it would carry a plan-mode-shaped # choice past plan mode's exit. No sentinel turn writes the pin, whatever its cause. pinnable: Final = ( - cache_key is not None + session_id is not None and response is not None and _decision_is_pinnable(response.routing_decision) and self._matched_plan_mode_signal(request_kwargs, resolved_messages) is None ) - if pinnable and cache_key is not None and response is not None: - await self.litellm_router_instance.cache.async_set_cache( - key=cache_key, - value=_session_affinity_cache_value( - response.model, - response.routing_decision.get("tier") if response.routing_decision is not None else None, - ), - ttl=self.config.session_affinity_ttl_seconds, + if pinnable and session_id is not None and response is not None: + await self._store_session_affinity_pin( + session_id, + request_kwargs, + response.model, + response.routing_decision.get("tier") if response.routing_decision is not None else None, ) return self._with_session_deployment_affinity(response) diff --git a/tests/unit/redis/stores/test_complexity_session_affinity.py b/tests/unit/redis/stores/test_complexity_session_affinity.py new file mode 100644 index 00000000000..2e73c230f04 --- /dev/null +++ b/tests/unit/redis/stores/test_complexity_session_affinity.py @@ -0,0 +1,230 @@ +import json +from collections.abc import Iterator, Mapping +from typing import Final +from unittest.mock import patch + +import fakeredis +import pytest +from redis.exceptions import TimeoutError as RedisTimeoutError + +from litellm import Router +from litellm.caching.redis_batch import request_redis_batch_scope +from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter +from litellm.types.router import PreRoutingHookResponse +from tests._support.fake_redis import ( + fake_redis_cache, + keys_matching, + legacy_get, + local_get, + local_set, + seed_async_client, +) + + +async def _store_pin( + strategy: ComplexityRouter, session_id: str, request_kwargs: Mapping[str, object], model: str, tier: str | None +) -> None: + await strategy._store_session_affinity_pin(session_id, request_kwargs, model, tier) # pyright: ignore[reportPrivateUsage] # exercises the pin path without a full router request + + +async def _read_pin(strategy: ComplexityRouter, session_id: str, request_kwargs: Mapping[str, object]) -> object: + return await strategy._read_session_affinity_pin(session_id, request_kwargs) # pyright: ignore[reportPrivateUsage] # exercises the pin path without a full router request + + +async def _pre_routing( + strategy: ComplexityRouter, *, model: str, request_kwargs: dict[str, object], messages: list[dict[str, object]] +) -> PreRoutingHookResponse | None: + response: PreRoutingHookResponse | None = await strategy.async_pre_routing_hook( # pyright: ignore[reportUnknownMemberType] # legacy hook API is untyped + model=model, request_kwargs=request_kwargs, messages=messages + ) + return response + + +@pytest.fixture(autouse=True) +def fixed_clock() -> Iterator[None]: + with patch("time.time", return_value=1700000000.0): + yield + + +def _strategy(server: fakeredis.FakeServer | None = None) -> ComplexityRouter: + parent: Final = Router( + model_list=[ + {"model_name": name, "litellm_params": {"model": "openai/test"}, "model_info": {"id": name}} + for name in ("tiny", "middle", "large", "top") + ] + ) + if server is not None: + parent.cache.redis_cache = fake_redis_cache(server, namespace="complexity") + return ComplexityRouter( + model_name="router", + litellm_router_instance=parent, + derive_savings_baseline=False, + complexity_router_config={ + "tiers": {"SIMPLE": "tiny", "MEDIUM": "middle", "COMPLEX": "large", "REASONING": "top"}, + "session_affinity": True, + "session_affinity_ttl_seconds": 60, + }, + ) + + +def _key(caller: str = "caller") -> str: + return f"complexity:complexity_router_session_affinity:v1:router:{caller}:same" + + +def _kwargs(caller: str = "caller") -> dict[str, object]: + return {"metadata": {"session_id": "same", "user_api_key_hash": caller}} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("tier", [None, "自定义"]) +async def test_overwrite_preserves_wire_bytes_scope_and_ttl(tier: str | None) -> None: + server: Final = fakeredis.FakeServer() + strategy: Final = _strategy(server) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await _store_pin(strategy, "same", _kwargs(), "模型", tier) + expected: Final = json.dumps({"model": "模型", "tier": tier}).encode() + assert await raw.get(_key()) == expected + assert await raw.pttl(_key()) == 60000 + redis_cache: Final = strategy.litellm_router_instance.cache.redis_cache + assert redis_cache is not None + assert await legacy_get(redis_cache, _key().removeprefix("complexity:")) == {"model": "模型", "tier": tier} + await _store_pin(strategy, "same", _kwargs(), "replacement", None) + await _store_pin(strategy, "same", _kwargs("second"), "other", None) + assert await raw.get(_key()) == b'{"model": "replacement", "tier": null}' + assert await raw.get(_key("second")) == b'{"model": "other", "tier": null}' + assert await keys_matching(raw, "complexity:complexity:*") == () + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ttl", [0, -1, 0.5, 1.9]) +async def test_overwrite_preserves_legacy_integer_seconds_and_invalid_expiry(ttl: float) -> None: + server: Final = fakeredis.FakeServer() + strategy: Final = _strategy(server) + strategy.config = strategy.config.model_copy(update={"session_affinity_ttl_seconds": ttl}) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await raw.set(_key(), b'"old"', ex=20) + await _store_pin(strategy, "same", _kwargs(), "tiny", None) + assert await raw.get(_key()) == (b'{"model": "tiny", "tier": null}' if int(ttl) > 0 else b'"old"') + assert await raw.pttl(_key()) == (int(ttl) * 1000 if int(ttl) > 0 else 20000) + + +@pytest.mark.asyncio +async def test_nonpositive_read_refresh_returns_old_pin_but_deletes_shared_key() -> None: + server: Final = fakeredis.FakeServer() + strategy: Final = _strategy(server) + strategy.config = strategy.config.model_copy(update={"session_affinity_ttl_seconds": 0}) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await raw.set(_key(), b'"tiny"', ex=20) + assert await _read_pin(strategy, "same", _kwargs()) == "tiny" + with patch("time.time", return_value=1700000000.001): + assert await raw.get(_key()) is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "wire,expected", + [ + (b'"tiny"', "tiny"), + (b"tiny", "tiny"), + (b'{"model": "tiny", "tier": null}', {"model": "tiny", "tier": None}), + (b"malformed{", "malformed{"), + ], +) +async def test_read_refreshes_without_reencoding_legacy_values(wire: bytes, expected: object) -> None: + server: Final = fakeredis.FakeServer() + strategy: Final = _strategy(server) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await raw.set(_key(), wire, ex=20) + assert await _read_pin(strategy, "same", _kwargs()) == expected + assert await raw.get(_key()) == wire + assert await raw.pttl(_key()) == 60000 + + +@pytest.mark.asyncio +async def test_successful_miss_invalidates_but_outage_refreshes_local_fallback() -> None: + server: Final = fakeredis.FakeServer() + strategy: Final = _strategy(server) + cache: Final = strategy.litellm_router_instance.cache + logical: Final = _key().removeprefix("complexity:") + local_set(cache, logical, "stale", ttl=20) + assert await _read_pin(strategy, "same", _kwargs()) is None + assert local_get(cache, logical) is None + local_set(cache, logical, "local", ttl=20) + server.connected = False + assert await _read_pin(strategy, "same", _kwargs()) == "local" + with patch("time.time", return_value=1700000021.0): + assert local_get(cache, logical) == "local" + + +@pytest.mark.asyncio +async def test_attached_redis_supersedes_local_only_session() -> None: + server: Final = fakeredis.FakeServer() + strategy: Final = _strategy() + await _store_pin(strategy, "same", _kwargs(), "local", None) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await raw.set(_key(), b'"shared"', ex=20) + strategy.litellm_router_instance.cache.redis_cache = fake_redis_cache(server, namespace="complexity") + assert await _read_pin(strategy, "same", _kwargs()) == "shared" + + +class LostWriteReply(fakeredis.FakeAsyncRedis): + async def set(self, *args: object, **kwargs: object) -> object: # pyright: ignore[reportIncompatibleMethodOverride] # passes the write through fakeredis + await super().set(*args, **kwargs) # pyright: ignore[reportArgumentType, reportUnknownMemberType] # fakeredis is untyped + raise RedisTimeoutError("reply lost after overwrite") + + +@pytest.mark.asyncio +async def test_unknown_overwrite_keeps_the_value_already_written_locally() -> None: + server: Final = fakeredis.FakeServer() + strategy: Final = _strategy(server) + redis_cache: Final = strategy.litellm_router_instance.cache.redis_cache + assert redis_cache is not None + seed_async_client(redis_cache, LostWriteReply(server=server)) + await _store_pin(strategy, "same", _kwargs(), "tiny", None) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + assert await raw.get(_key()) == b'{"model": "tiny", "tier": null}' + assert local_get(strategy.litellm_router_instance.cache, _key().removeprefix("complexity:")) == { + "model": "tiny", + "tier": None, + } + + +@pytest.mark.asyncio +async def test_classification_and_pin_escalation_both_overwrite_shared_session() -> None: + server: Final = fakeredis.FakeServer() + strategy: Final = _strategy(server) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + first: Final = await _pre_routing( + strategy, model="router", request_kwargs=_kwargs(), messages=[{"role": "user", "content": "Hi"}] + ) + assert first is not None and first.model == "tiny" + assert await raw.get(_key()) == b'{"model": "tiny", "tier": "SIMPLE"}' + second: Final = await _pre_routing( + strategy, model="router", request_kwargs=_kwargs(), messages=[{"role": "user", "content": "LITELLM ESCALATE"}] + ) + assert second is not None and second.model == "middle" + assert await raw.get(_key()) == b'{"model": "middle", "tier": "MEDIUM"}' + + +@pytest.mark.asyncio +async def test_reads_join_only_pending_request_work_not_post_call_writes() -> None: + server: Final = fakeredis.FakeServer() + strategy: Final = _strategy(server) + redis_cache: Final = strategy.litellm_router_instance.cache.redis_cache + assert redis_cache is not None + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await raw.set(_key(), b'"tiny"', ex=20) + with request_redis_batch_scope(post_call_deadline=3600) as batches: + request: Final = batches.batch(redis_cache) + post_call: Final = batches.post_call(redis_cache) + post_call.set("unrelated", "later", 60) + assert await _read_pin(strategy, "same", _kwargs()) == "tiny" + assert request.flushes == 0 + request.set("before-read", "ready", 60) + assert await _read_pin(strategy, "same", _kwargs()) == "tiny" + assert request.flushes == 1 + assert await raw.get("complexity:before-read") == b'"ready"' + assert await raw.get("complexity:unrelated") is None + assert post_call.pending == 1 + await batches.flush_post_call() + assert await raw.pttl(_key()) == 60000 From 3a70bbc5adfeb2b105029270fac1ab7f762d38e6 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 21:56:49 +0000 Subject: [PATCH 16/17] refactor(redis): migrate complexity tier model affinity into a typed store --- litellm/caching/affinity_cache.py | 151 +----- .../stores/complexity_tier_model_affinity.py | 103 +++++ .../complexity_router/complexity_router.py | 70 ++- .../test_complexity_tier_model_affinity.py | 234 ++++++++++ .../router_strategy/test_complexity_router.py | 432 +++++++++++------- .../test_session_id_affinity.py | 30 +- 6 files changed, 662 insertions(+), 358 deletions(-) create mode 100644 litellm/redis/stores/complexity_tier_model_affinity.py create mode 100644 tests/unit/redis/stores/test_complexity_tier_model_affinity.py diff --git a/litellm/caching/affinity_cache.py b/litellm/caching/affinity_cache.py index 2b8dc67da7a..aa461249ac7 100644 --- a/litellm/caching/affinity_cache.py +++ b/litellm/caching/affinity_cache.py @@ -1,72 +1,12 @@ -"""Atomic affinity claims shared by deployment and tier-model selection.""" +"""Pod-local affinity policy shared by deployment and tier-model selection.""" -import json from collections.abc import Mapping -from typing import ( - Final, - cast, # noqa: TID251 # Heterogeneous cache results are narrowed only to object, then validated -) +from typing import Final -from pydantic import JsonValue, TypeAdapter, ValidationError - -from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache -from litellm.caching.redis_batch import active_request_redis_batch ROUTER_SESSION_PINS_TARGET: Final = "router_session_pins" -_PIN_JSON_ADAPTER: Final = TypeAdapter[JsonValue](JsonValue) - -_READ_PIN_SCRIPT: Final = """ -local current = redis.call('GET', KEYS[1]) -if current ~= false then - redis.call('EXPIRE', KEYS[1], ARGV[1]) -end -return current -""" - -_CLAIM_PIN_SCRIPT: Final = """ -local current = redis.call('GET', KEYS[1]) -if current == false then - redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2]) - return ARGV[1] -end -if ARGV[3] then - local decoded, stored = pcall(cjson.decode, current) - if ARGV[4] then - local legacy = current - if decoded and type(stored) == 'string' then legacy = stored end - for _, eligible in ipairs(cjson.decode(ARGV[3])) do - if eligible['model_id'] == legacy then - redis.call('EXPIRE', KEYS[1], ARGV[2]) - return current - end - end - end - if decoded and type(stored) == 'table' then - for _, eligible in ipairs(cjson.decode(ARGV[3])) do - local matches = true - for key, value in pairs(eligible) do - if stored[key] ~= value then matches = false; break end - end - for key, _ in pairs(stored) do - if eligible[key] == nil then matches = false; break end - end - if matches then - redis.call('EXPIRE', KEYS[1], ARGV[2]) - return current - end - end - end - redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2]) - return ARGV[1] -end -if current == ARGV[1] then - redis.call('EXPIRE', KEYS[1], ARGV[2]) -end -return current -""" - def set_local_affinity_pin(cache: DualCache, cache_key: str, value: object, ttl_seconds: int) -> None: """Replace the entry because InMemoryCache.set_cache preserves a live key's expiry.""" @@ -103,90 +43,3 @@ def claim_affinity_pin_in_memory( winner: Final = existing if existing is not None and eligible_existing else pin_value set_local_affinity_pin(cache, cache_key, winner, ttl_seconds) return winner - - -def _decode_pin(value: str) -> object: - try: - return _PIN_JSON_ADAPTER.validate_json(value) - except ValidationError: - return value - - -async def read_affinity_pin(cache: DualCache, cache_key: str, ttl_seconds: int) -> object: - """Refresh the authoritative pin without writing a previously observed value back.""" - redis_cache: Final = cache.redis_cache - if redis_cache is not None: - try: - read_script: Final = redis_cache.async_register_script(_READ_PIN_SCRIPT) - batch: Final = active_request_redis_batch(redis_cache) - raw: Final = cast( # cast-ok: Redis script returns a serialized pin or a miss; validate before use - object, - await batch.script(_READ_PIN_SCRIPT, read_script, (cache_key,), (int(ttl_seconds),)) - if batch is not None and batch.pending - else await read_script(keys=(cache_key,), args=(int(ttl_seconds),)), - ) - decoded: Final = raw.decode("utf-8") if isinstance(raw, bytes) else raw - value: Final = _decode_pin(decoded) if isinstance(decoded, str) else None - if value is None: - cache.in_memory_cache.delete_cache(cache_key) - else: - set_local_affinity_pin(cache, cache_key, value, ttl_seconds) - except Exception as error: # noqa: BLE001 # Redis faults retain pod-local continuity, unlike a successful miss - verbose_router_logger.debug("Affinity cache: Redis read failed, using pod-local pin. error=%s", error) - else: - return value - local_value: Final = cast( # cast-ok: the heterogeneous local cache is parsed by the affinity caller - object, cache.in_memory_cache.get_cache(cache_key) - ) - if local_value is not None: - set_local_affinity_pin(cache, cache_key, local_value, ttl_seconds) - return local_value - - -async def claim_affinity_pin( - cache: DualCache, - cache_key: str, - pin_value: Mapping[str, str], - ttl_seconds: int, - *, - eligible_values: tuple[Mapping[str, str], ...] | None = None, - preserve_legacy_model_ids: bool = False, -) -> object: - """Return the authoritative first writer, replacing it only when it becomes ineligible. - - Eligible claims refresh the returned winner. Legacy deployment claims only refresh - a matching candidate. Resolve Redis per call because the proxy attaches it lazily. - """ - redis_cache: Final = cache.redis_cache - if redis_cache is not None: - try: - claim_script: Final = redis_cache.async_register_script(_CLAIM_PIN_SCRIPT) - args: Final = ( - json.dumps(dict(pin_value)), - int(ttl_seconds), - *( - (json.dumps(tuple(dict(value) for value in eligible_values)),) - if eligible_values is not None - else () - ), - *(("1",) if preserve_legacy_model_ids and eligible_values is not None else ()), - ) - raw: Final = cast( # cast-ok: Redis scripts return heterogeneous values; only object is asserted here - object, await claim_script(keys=(cache_key,), args=args) - ) - decoded: Final = raw.decode("utf-8") if isinstance(raw, bytes) else raw - if not isinstance(decoded, str): - return pin_value - winner: Final = _decode_pin(decoded) - set_local_affinity_pin(cache, cache_key, winner, ttl_seconds) - return winner - except Exception as error: # noqa: BLE001 # Redis/Lua faults retain same-pod affinity through local claims - verbose_router_logger.debug("Affinity cache: Redis claim failed, using pod-local claim. error=%s", error) - return claim_affinity_pin_in_memory( - cache, - cache_key, - pin_value, - ttl_seconds, - eligible_values=eligible_values, - preserve_legacy_model_ids=preserve_legacy_model_ids, - ) diff --git a/litellm/redis/stores/complexity_tier_model_affinity.py b/litellm/redis/stores/complexity_tier_model_affinity.py new file mode 100644 index 00000000000..18d3057f9cf --- /dev/null +++ b/litellm/redis/stores/complexity_tier_model_affinity.py @@ -0,0 +1,103 @@ +import hashlib +import json +from dataclasses import dataclass +from typing import Final, Protocol + +from pydantic import JsonValue, TypeAdapter, ValidationError + +from litellm.redis.transport.contract import RedisFailure, RedisTransport + +_JSON: Final = TypeAdapter[JsonValue](JsonValue) +_CLAIM: Final = """ +local current = redis.call('GET', KEYS[1]) +if current == false then + redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2]) + return ARGV[1] +end +local decoded, stored = pcall(cjson.decode, current) +if decoded and type(stored) == 'table' then + for _, eligible in ipairs(cjson.decode(ARGV[3])) do + local matches = true + for key, value in pairs(eligible) do + if stored[key] ~= value then matches = false; break end + end + for key, _ in pairs(stored) do + if eligible[key] == nil then matches = false; break end + end + if matches then + redis.call('EXPIRE', KEYS[1], ARGV[2]) + return current + end + end +end +redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2]) +return ARGV[1] +""" + + +def configuration_fingerprint(config_json: str) -> str: + return hashlib.sha256(config_json.encode()).hexdigest() + + +@dataclass(frozen=True, slots=True) +class TierModelIdentity: + model_name: str + configuration: str + caller: str | None + session_id: str + tier: str + replay_fallback: str | None = None + + +def affinity_key(identity: TierModelIdentity) -> str: + parts: Final = (identity.model_name, identity.configuration, identity.caller, identity.session_id, identity.tier) + cache_identity: Final = ( + (*parts, ("replay_fallback", identity.replay_fallback)) if identity.replay_fallback is not None else parts + ) + return "complexity_router_tier_model_affinity:v1:" + hashlib.sha256(json.dumps(cache_identity).encode()).hexdigest() + + +@dataclass(frozen=True, slots=True) +class TierModelPin: + value: object + + +@dataclass(frozen=True, slots=True) +class UnconfirmedTierModelPin: + pass + + +class ComplexityTierModelAffinityStore(Protocol): + async def claim( + self, identity: TierModelIdentity, model: str, eligible_models: tuple[str, ...], *, ttl_seconds: int + ) -> TierModelPin | UnconfirmedTierModelPin | RedisFailure: ... + + +@dataclass(frozen=True, slots=True) +class TransportComplexityTierModelAffinityStore: + transport: RedisTransport + + async def claim( + self, identity: TierModelIdentity, model: str, eligible_models: tuple[str, ...], *, ttl_seconds: int + ) -> TierModelPin | UnconfirmedTierModelPin | RedisFailure: + raw: Final = await self.transport.eval( + _CLAIM, + (affinity_key(identity),), + ( + json.dumps({"model": model}), + int(ttl_seconds), + json.dumps(tuple({"model": value} for value in eligible_models)), + ), + ) + if isinstance(raw, RedisFailure): + return raw + if not isinstance(raw, bytes): + return UnconfirmedTierModelPin() + try: + decoded: Final = raw.decode("utf-8") + except UnicodeError as error: + return RedisFailure(kind="decode", message=str(error)) + try: + return TierModelPin(_JSON.validate_json(decoded)) + except ValidationError: + return TierModelPin(decoded) diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index a6adc14dfd9..673c5b167c0 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -17,7 +17,6 @@ Inspired by ClawRouter: https://github.com/BlockRunAI/ClawRouter from __future__ import annotations import asyncio -import hashlib import json import random import re @@ -36,7 +35,11 @@ from typing_extensions import ReadOnly, TypedDict from litellm._internal_context import with_service_target from litellm._logging import verbose_router_logger -from litellm.caching.affinity_cache import ROUTER_SESSION_PINS_TARGET, claim_affinity_pin, set_local_affinity_pin +from litellm.caching.affinity_cache import ( + ROUTER_SESSION_PINS_TARGET, + claim_affinity_pin_in_memory, + set_local_affinity_pin, +) from litellm.caching.redis_batch import active_request_redis_batch from litellm.constants import ( COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY, @@ -76,6 +79,18 @@ from litellm.redis.stores.complexity_session_affinity import ( from litellm.redis.stores.complexity_session_affinity import ( pin_value as session_pin_value, ) +from litellm.redis.stores.complexity_tier_model_affinity import ( + ComplexityTierModelAffinityStore, + TierModelIdentity, + TierModelPin, + TransportComplexityTierModelAffinityStore, + UnconfirmedTierModelPin, + configuration_fingerprint, +) +from litellm.redis.stores.complexity_tier_model_affinity import ( + affinity_key as tier_model_affinity_key, +) +from litellm.redis.transport.contract import RedisFailure from litellm.router_strategy.adaptive_router.classifier import classify_prompt from litellm.router_strategy.complexity_router.context_compaction import compaction_pending from litellm.router_strategy.complexity_router.tier_predictor import ( @@ -1515,9 +1530,9 @@ class ComplexityRouter(CustomLogger): case _: self._jev_client: JevClassifierClient | None = jev_client_result - self._tier_affinity_config = hashlib.sha256( - self.config.model_dump_json(include=MappingProxyType({"tiers": True, "tier_model_configs": True})).encode() - ).hexdigest() + self._tier_affinity_config = configuration_fingerprint( + self.config.model_dump_json(include=MappingProxyType({"tiers": True, "tier_model_configs": True})) + ) # Checked here rather than on the config model because the deployment's # complexity_router_default_model arrives outside complexity_router_config and is @@ -3125,27 +3140,44 @@ class ComplexityRouter(CustomLogger): if session_id is None: return model caller: Final = DeploymentAffinityCheck.get_user_key_from_request_kwargs(request_kwargs) - identity: Final = (self.model_name, self._tier_affinity_config, caller, session_id, _tier_name(tier)) - cache_identity: Final = ( - (*identity, ("replay_fallback", retained_model)) if retained_model is not None else identity - ) - cache_key: Final = ( - "complexity_router_tier_model_affinity:v1:" - + hashlib.sha256(json.dumps(cache_identity).encode()).hexdigest() - ) - winner: Final = await claim_affinity_pin( - self.litellm_router_instance.cache, - cache_key, - MappingProxyType({"model": model}), - self.config.session_affinity_ttl_seconds, - eligible_values=tuple(MappingProxyType({"model": candidate}) for candidate in candidates), + identity: Final = TierModelIdentity( + self.model_name, self._tier_affinity_config, caller, session_id, _tier_name(tier), retained_model ) + winner: Final = await self._claim_tier_model_pin(identity, model, candidates) pinned: Final[object] = winner.get("model") if isinstance(winner, Mapping) else None if not isinstance(pinned, str) or pinned not in candidates: return model self._restamp_adaptive_choice(request_kwargs, model, pinned) return pinned + async def _claim_tier_model_pin( + self, identity: TierModelIdentity, model: str, candidates: tuple[str, ...] + ) -> object: + cache: Final = self.litellm_router_instance.cache + cache_key: Final = tier_model_affinity_key(identity) + ttl: Final = self.config.session_affinity_ttl_seconds + redis_cache: Final = cache.redis_cache + if redis_cache is not None: + store: Final[ComplexityTierModelAffinityStore] = TransportComplexityTierModelAffinityStore( + redis_cache.transport() + ) + result: Final = await store.claim(identity, model, candidates, ttl_seconds=ttl) + match result: + case TierModelPin(value=value): + set_local_affinity_pin(cache, cache_key, value, ttl) + return value + case UnconfirmedTierModelPin(): + return MappingProxyType({"model": model}) + case RedisFailure(): + pass + return claim_affinity_pin_in_memory( + cache, + cache_key, + MappingProxyType({"model": model}), + ttl, + eligible_values=tuple(MappingProxyType({"model": candidate}) for candidate in candidates), + ) + async def _pick_model_for_tier( self, tier: ComplexityTier | str, diff --git a/tests/unit/redis/stores/test_complexity_tier_model_affinity.py b/tests/unit/redis/stores/test_complexity_tier_model_affinity.py new file mode 100644 index 00000000000..f7dcd083ee3 --- /dev/null +++ b/tests/unit/redis/stores/test_complexity_tier_model_affinity.py @@ -0,0 +1,234 @@ +import asyncio +import hashlib +import json +from collections.abc import Iterator +from typing import Final +from unittest.mock import patch + +import fakeredis +import pytest +from redis.exceptions import TimeoutError as RedisTimeoutError + +from litellm import Router +from litellm.caching.redis_batch import request_redis_batch_scope +from litellm.router_strategy.complexity_router.complexity_router import ( + ComplexityRouter, + _SessionAffinityPin, # pyright: ignore[reportPrivateUsage] # builds the retained pin the strategy carries between tiers +) +from tests._support.fake_redis import ( + fake_redis_cache, + keys_matching, + local_get, + local_set, + seed_async_client, +) + + +async def _pin( + strategy: ComplexityRouter, + tier: str, + model: str, + candidates: tuple[str, ...], + request_kwargs: dict[str, object], + retained_pin: _SessionAffinityPin | None = None, +) -> object: + return await strategy._pin_model_for_tier(tier, model, candidates, request_kwargs, retained_pin) # pyright: ignore[reportPrivateUsage] # exercises the pin path without a full router request + + +@pytest.fixture(autouse=True) +def fixed_clock() -> Iterator[None]: + with patch("time.time", return_value=1700000000.0): + yield + + +def _strategy(server: fakeredis.FakeServer | None, *, name: str = "router", simple: str = "tiny") -> ComplexityRouter: + parent: Final = Router( + model_list=[ + {"model_name": name, "litellm_params": {"model": "openai/test"}, "model_info": {"id": name}} + for name in ("tiny", "middle", "large", "top") + ] + ) + if server is not None: + parent.cache.redis_cache = fake_redis_cache(server, namespace="tier") + return ComplexityRouter( + model_name=name, + litellm_router_instance=parent, + derive_savings_baseline=False, + complexity_router_config={ + "tiers": {"SIMPLE": [simple, "middle"], "MEDIUM": "middle", "COMPLEX": "large", "REASONING": "top"}, + "session_affinity": True, + "session_affinity_ttl_seconds": 60, + }, + ) + + +def _key(strategy: ComplexityRouter, caller: str | None = "caller", *, fallback: str | None = None) -> str: + config: Final = strategy.config.model_dump_json(include={"tiers": True, "tier_model_configs": True}) + identity: Final = (strategy.model_name, hashlib.sha256(config.encode()).hexdigest(), caller, "same", "SIMPLE") + encoded: Final = json.dumps((*identity, ("replay_fallback", fallback)) if fallback is not None else identity) + return "tier:complexity_router_tier_model_affinity:v1:" + hashlib.sha256(encoded.encode()).hexdigest() + + +def _kwargs(caller: str | None = "caller") -> dict[str, object]: + return {"metadata": {"session_id": "same", "user_api_key_hash": caller}} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "name,caller,simple,fallback", + [ + ("router", "caller", "tiny", None), + ("路由", "caller", "tiny", None), + ("router", "另一个", "tiny", None), + ("router", None, "tiny", None), + ("router", "caller", "模型", None), + ("router", "caller", "tiny", "removed"), + ], +) +async def test_identity_and_unicode_wire_match_legacy_hash( + name: str, caller: str | None, simple: str, fallback: str | None +) -> None: + server: Final = fakeredis.FakeServer() + strategy: Final = _strategy(server, name=name, simple=simple) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + retained: Final = _SessionAffinityPin(fallback, "SIMPLE") if fallback is not None else None + assert await _pin(strategy, "SIMPLE", "模型", ("模型", "tiny"), _kwargs(caller), retained) == "模型" + assert await raw.get(_key(strategy, caller, fallback=fallback)) == b'{"model": "\\u6a21\\u578b"}' + assert await raw.pttl(_key(strategy, caller, fallback=fallback)) == 60000 + assert await keys_matching(raw, "tier:tier:*") == () + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "wire,expected", + [ + (b'{"model": "tiny"}', "tiny"), + (b'{"model": "removed"}', "middle"), + (b'{"model": "tiny", "extra": true}', "middle"), + (b'"tiny"', "middle"), + (b"tiny", "middle"), + (b"{broken", "middle"), + ], +) +async def test_claim_refreshes_only_complete_eligible_winners(wire: bytes, expected: str) -> None: + server: Final = fakeredis.FakeServer() + strategy: Final = _strategy(server) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await raw.set(_key(strategy), wire, ex=20) + assert await _pin(strategy, "SIMPLE", "middle", ("tiny", "middle"), _kwargs()) == expected + assert await raw.get(_key(strategy)) == json.dumps({"model": expected}).encode() + assert await raw.pttl(_key(strategy)) == 60000 + + +@pytest.mark.asyncio +async def test_concurrent_pods_converge_then_replace_expired_pin() -> None: + server: Final = fakeredis.FakeServer() + left: Final = _strategy(server) + right: Final = _strategy(server) + winners: Final = await asyncio.gather( + _pin(left, "SIMPLE", "tiny", ("tiny", "middle"), _kwargs()), + _pin(right, "SIMPLE", "middle", ("tiny", "middle"), _kwargs()), + ) + assert winners[0] == winners[1] + raw: Final = fakeredis.FakeAsyncRedis(server=server) + assert await raw.get(_key(left)) == json.dumps({"model": winners[0]}).encode() + with patch("time.time", return_value=1700000061.0): + assert await _pin(left, "SIMPLE", "middle", ("tiny", "middle"), _kwargs()) == "middle" + assert await raw.get(_key(left)) == b'{"model": "middle"}' + assert await raw.pttl(_key(left)) == 60000 + + +@pytest.mark.asyncio +async def test_no_session_and_eligible_retained_pin_do_not_claim() -> None: + server: Final = fakeredis.FakeServer() + strategy: Final = _strategy(server) + assert await _pin(strategy, "SIMPLE", "tiny", ("tiny", "middle"), {}) == "tiny" + assert ( + await _pin(strategy, "SIMPLE", "middle", ("tiny", "middle"), _kwargs(), _SessionAffinityPin("tiny", "SIMPLE")) + == "tiny" + ) + assert await fakeredis.FakeAsyncRedis(server=server).keys("tier:*") == [] + + +@pytest.mark.asyncio +async def test_late_redis_attachment_replaces_local_winner() -> None: + server: Final = fakeredis.FakeServer() + strategy: Final = _strategy(None) + assert await _pin(strategy, "SIMPLE", "tiny", ("tiny", "middle"), _kwargs()) == "tiny" + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await raw.set(_key(strategy), b'{"model": "middle"}', ex=20) + strategy.litellm_router_instance.cache.redis_cache = fake_redis_cache(server, namespace="tier") + assert await _pin(strategy, "SIMPLE", "tiny", ("tiny", "middle"), _kwargs()) == "middle" + assert local_get(strategy.litellm_router_instance.cache, _key(strategy).removeprefix("tier:")) == { + "model": "middle" + } + + +class LostClaimReply(fakeredis.FakeAsyncRedis): + async def eval(self, script: str, numkeys: int, *keys_and_args: str | bytes | int) -> object: + await super().eval(script, numkeys, *keys_and_args) # pyright: ignore[reportUnknownMemberType] # fakeredis is untyped + raise RedisTimeoutError("reply lost after claim") + + +class ConsumedClaimReply(fakeredis.FakeAsyncRedis): + async def eval(self, script: str, numkeys: int, *keys_and_args: str | bytes | int) -> object: + await super().eval(script, numkeys, *keys_and_args) # pyright: ignore[reportUnknownMemberType] # fakeredis is untyped + return 1 + + +@pytest.mark.asyncio +async def test_consumed_claim_returns_candidate_without_installing_local_winner() -> None: + server: Final = fakeredis.FakeServer() + strategy: Final = _strategy(server) + cache: Final = strategy.litellm_router_instance.cache + redis_cache: Final = cache.redis_cache + assert redis_cache is not None + seed_async_client(redis_cache, ConsumedClaimReply(server=server)) + local_set(cache, _key(strategy).removeprefix("tier:"), {"model": "tiny"}, ttl=20) + assert await _pin(strategy, "SIMPLE", "middle", ("tiny", "middle"), _kwargs()) == "middle" + assert await fakeredis.FakeAsyncRedis(server=server).get(_key(strategy)) == b'{"model": "middle"}' + assert local_get(cache, _key(strategy).removeprefix("tier:")) == {"model": "tiny"} + + +@pytest.mark.asyncio +async def test_claim_leaves_pending_request_and_post_call_work_unflushed() -> None: + server: Final = fakeredis.FakeServer() + strategy: Final = _strategy(server) + redis_cache: Final = strategy.litellm_router_instance.cache.redis_cache + assert redis_cache is not None + raw: Final = fakeredis.FakeAsyncRedis(server=server) + with request_redis_batch_scope(post_call_deadline=3600) as batches: + request: Final = batches.batch(redis_cache) + post_call: Final = batches.post_call(redis_cache) + request.set("pending-request", "ready", 60) + post_call.set("pending-post-call", "later", 60) + assert await _pin(strategy, "SIMPLE", "tiny", ("tiny", "middle"), _kwargs()) == "tiny" + assert await raw.get(_key(strategy)) == b'{"model": "tiny"}' + assert await raw.get("tier:pending-request") is None + assert await raw.get("tier:pending-post-call") is None + await batches.flush_post_call() + await request.flush() + assert await raw.get("tier:pending-request") == b'"ready"' + assert await raw.get("tier:pending-post-call") == b'"later"' + + +@pytest.mark.asyncio +@pytest.mark.parametrize("applied", [False, True]) +async def test_outage_and_unknown_claim_keep_local_winner(applied: bool) -> None: + server: Final = fakeredis.FakeServer() + strategy: Final = _strategy(server) + cache: Final = strategy.litellm_router_instance.cache + redis_cache: Final = cache.redis_cache + assert redis_cache is not None + local_set(cache, _key(strategy).removeprefix("tier:"), {"model": "tiny"}, ttl=20) + if applied: + seed_async_client(redis_cache, LostClaimReply(server=server)) + else: + server.connected = False + assert await _pin(strategy, "SIMPLE", "middle", ("tiny", "middle"), _kwargs()) == "tiny" + server.connected = True + assert await fakeredis.FakeAsyncRedis(server=server).get(_key(strategy)) == ( + b'{"model": "middle"}' if applied else None + ) + with patch("time.time", return_value=1700000021.0): + assert local_get(cache, _key(strategy).removeprefix("tier:")) == {"model": "tiny"} diff --git a/tests/unit/router_strategy/test_complexity_router.py b/tests/unit/router_strategy/test_complexity_router.py index ab168b2a950..98a53def5b4 100644 --- a/tests/unit/router_strategy/test_complexity_router.py +++ b/tests/unit/router_strategy/test_complexity_router.py @@ -14,7 +14,7 @@ from collections.abc import AsyncIterator, Mapping, Sequence from copy import deepcopy from functools import partial from types import MappingProxyType -from typing import Dict, Final, List, Literal +from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -39,7 +39,6 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.router import as_output_cap, stamped_client_ceiling from litellm.router_strategy.auto_router.auto_router_model_naming import ( CUSTOMIZATION_CAPABILITY, - GATED_AUTO_ROUTER_CAPABILITIES, HEURISTIC_V2_CAPABILITY, count_capability_routers, ) @@ -128,7 +127,7 @@ def mock_router_instance(): @pytest.fixture -def basic_config() -> Dict: +def basic_config() -> dict: """Basic configuration with tier mappings.""" return { "tiers": { @@ -1498,7 +1497,6 @@ class TestSingletonMutation: def test_default_config_not_mutated(self, mock_router_instance): """Test that creating routers without config doesn't mutate defaults.""" from litellm.router_strategy.complexity_router.config import ( - DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE, ComplexityRouterConfig, ) @@ -1554,7 +1552,7 @@ class TestKeywordFalsePositives: tier, score, signals = complexity_router.classify(prompt) # 'entry' contains 'try' but should not trigger code detection # Note: 'application' might trigger something, but 'try' should not - pass # Just ensure no crash; false positive check is the main goal + # Just ensure no crash; false positive check is the main goal def test_error_not_in_terrorism(self, complexity_router): """'error' should not match in 'terrorism'.""" @@ -1696,17 +1694,19 @@ class TestRouterComplexityDeploymentMethods: def test_custom_default_resolution_preserves_explicit_precedence(self, tier_models, explicit, expected): from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig - config: Final = ComplexityRouterConfig.model_validate({ - "tiers": {"CUSTOM": tier_models, "OTHER": "other-model"}, - "tier_definitions": [ - {"name": "CUSTOM", "description": "Custom work"}, - {"name": "OTHER", "description": "Other work"}, - ], - "fallback_tier": " CUSTOM ", - "classifier_type": "llm", - "classifier_llm_config": {"model": "classifier"}, - "default_model": "configured-model", - }) + config: Final = ComplexityRouterConfig.model_validate( + { + "tiers": {"CUSTOM": tier_models, "OTHER": "other-model"}, + "tier_definitions": [ + {"name": "CUSTOM", "description": "Custom work"}, + {"name": "OTHER", "description": "Other work"}, + ], + "fallback_tier": " CUSTOM ", + "classifier_type": "llm", + "classifier_llm_config": {"model": "classifier"}, + "default_model": "configured-model", + } + ) assert config.resolve_default_model(explicit) == expected inferred: Final = config.model_copy(update={"default_model": None}) assert inferred.resolve_default_model() == "custom-model" @@ -2142,9 +2142,7 @@ class TestRouterComplexityDeploymentMethods: @pytest.mark.parametrize("classifier", ("heuristic_first", "hybrid")) @pytest.mark.parametrize("held_capability", ("heuristic_v2", "custom_prompt")) - def test_v2_chain_registration_enforces_each_claimed_allowance( - self, classifier: str, held_capability: str - ) -> None: + def test_v2_chain_registration_enforces_each_claimed_allowance(self, classifier: str, held_capability: str) -> None: held: Final = ( self._router_row("existing", "existing-id", "heuristic_v2") if held_capability == "heuristic_v2" @@ -2160,7 +2158,8 @@ class TestRouterComplexityDeploymentMethods: "local_heuristic": "heuristic_v2", **( {"heuristic_first_max_tier": "MEDIUM"} - if classifier == "heuristic_first" else {"hybrid_boundary_margin": 0.1} + if classifier == "heuristic_first" + else {"hybrid_boundary_margin": 0.1} ), "classifier_llm_config": {"model": "gpt-4o-mini", "system_prompt": "Grade by difficulty"}, "tiers": {"SIMPLE": "gpt-4o-mini"}, @@ -2785,7 +2784,7 @@ def _wrapped_reply(shape: str, verdict: str) -> str: @pytest.fixture -def llm_classifier_config() -> Dict: +def llm_classifier_config() -> dict: """Config with an LLM-based classifier wired to a 'haiku-classifier' model.""" return { "tiers": { @@ -2839,7 +2838,7 @@ class TestLLMClassifierConfig: ) -CAPABILITY_TIERS: Dict[str, str] = { +CAPABILITY_TIERS: dict[str, str] = { "SIMPLE": "efficient-model", "REASONING": "capable-model", } @@ -3415,7 +3414,7 @@ class TestCapabilityClassifier: mock_router_instance.acompletion.assert_awaited_once() -CUSTOM_TIER_LABELS: Dict[str, str] = { +CUSTOM_TIER_LABELS: dict[str, str] = { "SIMPLE": "Cheap", "MEDIUM": "Standard", "COMPLEX": "Premium", @@ -4856,7 +4855,7 @@ class TestRouterPreRoutingAliasOverrides: reach the outbound request even though the tier deployment is what actually gets called.""" router = self._make_router() - request_kwargs: Dict = {} + request_kwargs: dict = {} result = await router.async_pre_routing_hook( model="smart-router", @@ -4889,7 +4888,7 @@ class TestRouterPreRoutingAliasOverrides: {"model_name": "gpt-5-mini", "litellm_params": {"model": "openai/gpt-5-mini"}}, ] ) - request_kwargs: Dict = {"reasoning_effort": "low"} + request_kwargs: dict = {"reasoning_effort": "low"} deployment = await router.async_get_available_deployment( model="smart-router", @@ -4900,7 +4899,7 @@ class TestRouterPreRoutingAliasOverrides: assert deployment["model_name"] == "gpt-5-mini" assert request_kwargs["reasoning_effort"] == "xhigh" - def _make_effort_pinned_router(self, tier_litellm_params: Dict) -> Router: + def _make_effort_pinned_router(self, tier_litellm_params: dict) -> Router: return Router( model_list=[ { @@ -4951,7 +4950,7 @@ class TestRouterPreRoutingAliasOverrides: carrier precedence over the reasoning_effort alias, so the pin only reaches the wire if those carriers are dropped at the merge.""" router = self._make_effort_pinned_router({"reasoning_effort": "xhigh"}) - request_kwargs: Dict = dict(client_carriers) + request_kwargs: dict = dict(client_carriers) await router.async_get_available_deployment( model="smart-router", @@ -4989,7 +4988,7 @@ class TestRouterPreRoutingAliasOverrides: }, ] ) - request_kwargs: Dict = {"thinking": {"type": "adaptive"}, "output_config": {"effort": "max"}} + request_kwargs: dict = {"thinking": {"type": "adaptive"}, "output_config": {"effort": "max"}} await router.async_get_available_deployment_for_pass_through( model="smart-router", @@ -5002,15 +5001,15 @@ class TestRouterPreRoutingAliasOverrides: assert "output_config" not in request_kwargs def test_drop_client_effort_carriers_helper_edge_shapes(self): - no_pin: Dict = {"thinking": {"type": "adaptive"}} + no_pin: dict = {"thinking": {"type": "adaptive"}} Router._drop_client_carriers_a_tier_pin_supersedes(no_pin, {"temperature": 0.1}) assert no_pin == {"thinking": {"type": "adaptive"}} - non_dict_carriers: Dict = {"output_config": "max", "reasoning": 3} + non_dict_carriers: dict = {"output_config": "max", "reasoning": 3} Router._drop_client_carriers_a_tier_pin_supersedes(non_dict_carriers, {"reasoning_effort": "low"}) assert non_dict_carriers == {"output_config": "max", "reasoning": 3} - effort_only: Dict = {"output_config": {"effort": "max"}, "reasoning": {"effort": "high"}} + effort_only: dict = {"output_config": {"effort": "max"}, "reasoning": {"effort": "high"}} Router._pop_effort_from_nested_carrier(effort_only, "output_config") Router._pop_effort_from_nested_carrier(effort_only, "reasoning") assert effort_only == {} @@ -5038,7 +5037,7 @@ class TestRouterPreRoutingAliasOverrides: {"model_name": "gpt-4o-mini", "litellm_params": {"model": "openai/gpt-4o-mini"}}, ] ) - request_kwargs: Dict = {"thinking": {"type": "adaptive"}, "output_config": {"effort": "max"}} + request_kwargs: dict = {"thinking": {"type": "adaptive"}, "output_config": {"effort": "max"}} await router.async_get_available_deployment( model="smart-router", @@ -5053,7 +5052,7 @@ class TestRouterPreRoutingAliasOverrides: @pytest.mark.asyncio async def test_client_effort_carriers_survive_when_tier_pins_no_effort(self): router = self._make_effort_pinned_router({"temperature": 0.2}) - request_kwargs: Dict = {"thinking": {"type": "adaptive"}, "output_config": {"effort": "max"}} + request_kwargs: dict = {"thinking": {"type": "adaptive"}, "output_config": {"effort": "max"}} await router.async_get_available_deployment( model="smart-router", @@ -5098,7 +5097,7 @@ class TestRouterPreRoutingAliasOverrides: ] ) real_get_llm_provider = litellm.get_llm_provider - copilot_resolutions: List = [] + copilot_resolutions: list = [] def _guarded(*args, **kwargs): target = str(kwargs.get("model") or (args[0] if args else "")) + str( @@ -5110,7 +5109,7 @@ class TestRouterPreRoutingAliasOverrides: return real_get_llm_provider(*args, **kwargs) monkeypatch.setattr(litellm, "get_llm_provider", _guarded) - request_kwargs: Dict = {} + request_kwargs: dict = {} deployment = await router.async_get_available_deployment( model="smart-router", @@ -5178,7 +5177,7 @@ class TestRouterPreRoutingAliasOverrides: test_router_init_only_params_are_never_sent_to_a_provider for the guard on that downstream filter.""" router = self._make_router() - request_kwargs: Dict = {} + request_kwargs: dict = {} await router.async_pre_routing_hook( model="smart-router", @@ -5232,7 +5231,7 @@ class TestRouterPreRoutingAliasOverrides: """A value the caller already passed for this request takes precedence over the alias's configured default.""" router = self._make_router() - request_kwargs: Dict = {"drop_params": False} + request_kwargs: dict = {"drop_params": False} await router.async_pre_routing_hook( model="smart-router", @@ -5247,7 +5246,7 @@ class TestRouterPreRoutingAliasOverrides: """A plain (non-router-alias) model name is not affected by the alias-override merge at all.""" router = self._make_router() - request_kwargs: Dict = {} + request_kwargs: dict = {} result = await router.async_pre_routing_hook( model="gpt-4o-mini", @@ -5282,7 +5281,7 @@ class TestRouterPreRoutingAliasOverrides: router.set_model_list(model_list) assert "smart-router" in router.adaptive_routers - request_kwargs: Dict = {} + request_kwargs: dict = {} await router.async_pre_routing_hook( model="smart-router", request_kwargs=request_kwargs, @@ -5344,7 +5343,7 @@ class TestRouterPreRoutingSharedAliasName: else [self._marker_entry(), self._plain_entry()] ) router = Router(model_list=[*shared_name_entries, self._tier_entry()]) - request_kwargs: Dict = {} + request_kwargs: dict = {} result = await router.async_pre_routing_hook( model="gpt4o", @@ -5374,7 +5373,7 @@ class TestRouterPreRoutingSharedAliasName: }, } router = Router(model_list=[marker_with_connection_params, self._tier_entry()]) - request_kwargs: Dict = {} + request_kwargs: dict = {} result = await router.async_pre_routing_hook( model="smart", @@ -5414,7 +5413,7 @@ class TestRouterPreRoutingSharedAliasName: ] ) - us_kwargs: Dict = {"metadata": {"tags": ["us"]}} + us_kwargs: dict = {"metadata": {"tags": ["us"]}} us_result = await router.async_pre_routing_hook( model="smart", request_kwargs=us_kwargs, @@ -5423,7 +5422,7 @@ class TestRouterPreRoutingSharedAliasName: assert us_result is not None and us_result.model == "gpt-us" assert us_kwargs["drop_params"] is True - cn_kwargs: Dict = {"metadata": {"tags": ["cn"]}} + cn_kwargs: dict = {"metadata": {"tags": ["cn"]}} cn_result = await router.async_pre_routing_hook( model="smart", request_kwargs=cn_kwargs, @@ -5597,7 +5596,7 @@ class TestRouterPreRoutingSharedAliasName: await asyncio.sleep(0.01) return await healthy_deployments(*args, **kwargs) - sent: Dict[str, str | None] = {} + sent: dict[str, str | None] = {} async def record(**kwargs): sent[kwargs["model"]] = kwargs.get("aws_region_name") @@ -5682,7 +5681,7 @@ class TestAdaptiveSoftFloors: return router @pytest.fixture - def hybrid_config(self) -> Dict: + def hybrid_config(self) -> dict: return { "adaptive": True, "adaptive_weights": {"quality": 0.7, "cost": 0.3}, @@ -5712,7 +5711,7 @@ class TestAdaptiveSoftFloors: }, }, ) - request_kwargs: Dict = {"metadata": {}} + request_kwargs: dict = {"metadata": {}} with patch( "litellm.router_strategy.complexity_router.complexity_router.random.choice", @@ -5811,7 +5810,7 @@ class TestAdaptiveSoftFloors: assert adaptive is not None for model in ("cheap", "premium"): adaptive._cells[(RequestType.GENERAL, model)] = BanditCell(alpha=6.0, beta=5.0) - request_kwargs: Dict = {"metadata": {}} + request_kwargs: dict = {"metadata": {}} with patch( "litellm.router_strategy.adaptive_router.bandit.thompson_sample", @@ -5832,7 +5831,7 @@ class TestAdaptiveSoftFloors: litellm_router_instance=adaptive_router_instance, complexity_router_config=hybrid_config, ) - request_kwargs: Dict = {"metadata": {}} + request_kwargs: dict = {"metadata": {}} result = await cr.async_pre_routing_hook( model="hybrid", request_kwargs=request_kwargs, @@ -5854,7 +5853,7 @@ class TestLexicalKeywordTierRules: """Test deterministic (literal) keyword_tier_rules overrides.""" @pytest.fixture - def rule_config(self, basic_config) -> Dict: + def rule_config(self, basic_config) -> dict: return { **basic_config, "keyword_tier_rules": [ @@ -5994,7 +5993,7 @@ class TestLexicalKeywordTierRules: class TestCjkKeywordTierRules: """CJK keyword_tier_rules must fire mid-sentence, where regex word boundaries cannot.""" - def _router(self, mock_router_instance, basic_config, keywords: List[str]) -> ComplexityRouter: + def _router(self, mock_router_instance, basic_config, keywords: list[str]) -> ComplexityRouter: return ComplexityRouter( model_name="test-router", litellm_router_instance=mock_router_instance, @@ -6069,7 +6068,7 @@ class TestCjkKeywordTierRules: assert complexity_router._keyword_matches("appelle l' api maintenant", "api") is True -def _make_embedding_response(vectors: List[List[float]]) -> "litellm.EmbeddingResponse": +def _make_embedding_response(vectors: list[list[float]]) -> "litellm.EmbeddingResponse": return litellm.EmbeddingResponse( model="fake-embed", data=[{"embedding": vec, "index": idx, "object": "embedding"} for idx, vec in enumerate(vectors)], @@ -6088,22 +6087,22 @@ class FakeEmbeddingRouter: _CLUSTER_MARKERS = ("k8s", "kube", "container", "cluster", "orchestrat") def __init__(self): - self.async_embedding_calls: List[List[str]] = [] - self.async_embedding_kwargs: List[Dict] = [] + self.async_embedding_calls: list[list[str]] = [] + self.async_embedding_kwargs: list[dict] = [] # Every embedded batch (sync route-index build AND async query), so tests can count # builds independently of which embedding path the library happens to use. - self.embedded_batches: List[List[str]] = [] + self.embedded_batches: list[list[str]] = [] # Thread ids of the synchronous (route-index build) embedding calls, so a test can # assert the build is offloaded off the event-loop thread. - self.sync_embedding_thread_ids: List[int] = [] + self.sync_embedding_thread_ids: list[int] = [] - def _vectors(self, docs: List[str]) -> List[List[float]]: + def _vectors(self, docs: list[str]) -> list[list[float]]: return [ [1.0, 0.0] if any(marker in doc.lower() for marker in self._CLUSTER_MARKERS) else [0.0, 1.0] for doc in docs ] @staticmethod - def _as_list(text) -> List[str]: + def _as_list(text) -> list[str]: return text if isinstance(text, list) else [text] def embedding(self, input, model, **kwargs): @@ -6598,7 +6597,7 @@ class _StubEncoder: """Minimal stand-in for LiteLLMRouterEncoder.aencode_queries, capturing the kwargs it was called with.""" def __init__(self): - self.aencode_queries_calls: List[Dict] = [] + self.aencode_queries_calls: list[dict] = [] async def aencode_queries(self, docs, **kwargs): self.aencode_queries_calls.append(kwargs) @@ -7198,11 +7197,11 @@ class TestSessionAffinity: SIMPLE_MESSAGE = [{"role": "user", "content": "Hello!"}] @pytest.fixture - def session_affinity_config(self, basic_config) -> Dict: + def session_affinity_config(self, basic_config) -> dict: return {**basic_config, "session_affinity": True} @staticmethod - def _request_kwargs(session_id: str) -> Dict: + def _request_kwargs(session_id: str) -> dict: return {"metadata": {"session_id": session_id}} @pytest.mark.asyncio @@ -7609,9 +7608,9 @@ class TestSessionAffinity: assert recovered.model == latest["model"] assert cache.in_memory_cache.get_cache(key) == latest assert read_started.is_set() - assert tuple( - tuple(args[0] for args in sent.args[0]) for sent in pack_commands.call_args_list - ) == (("MGET", "EVALSHA"),) * (2 if batch_pending else 0) + assert tuple(tuple(args[0] for args in sent.args[0]) for sent in pack_commands.call_args_list) == ( + ("MGET", "EVALSHA"), + ) * (2 if batch_pending else 0) finally: for cached_key in tuple(litellm.in_memory_llm_clients_cache.cache_dict): if client_key in cached_key: @@ -7737,8 +7736,12 @@ class TestChildSessionTierPins: "model_info": {"id": f"{group}-{index}"}, } for group, index in ( - ("simple", 0), ("simple", 1), ("reasoning", 0), - ("reasoning", 1), ("reasoning-peer", 0), ("reasoning-peer", 1), + ("simple", 0), + ("simple", 1), + ("reasoning", 0), + ("reasoning", 1), + ("reasoning-peer", 0), + ("reasoning-peer", 1), ) ], ], @@ -7754,9 +7757,9 @@ class TestChildSessionTierPins: metadata_variable_name_for_route, ) - route: Final = { - "chat": "/v1/chat/completions", "responses": "/v1/responses", "messages": "/v1/messages" - }[surface] + route: Final = {"chat": "/v1/chat/completions", "responses": "/v1/responses", "messages": "/v1/messages"}[ + surface + ] bucket: Final = metadata_variable_name_for_route(route) data: Final[dict[str, object]] = {bucket: {"user_api_key_hash": caller, "user_api_key_request_route": route}} LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( @@ -7770,54 +7773,74 @@ class TestChildSessionTierPins: return data, bucket async def _route( - self, router: Router, surface: str, session_id: str | None, prompt: str | None, - expected_tier: str, expected_cause: str, caller: str = "caller-a", child: bool = True, + self, + router: Router, + surface: str, + session_id: str | None, + prompt: str | None, + expected_tier: str, + expected_cause: str, + caller: str = "caller-a", + child: bool = True, expected_deployment: str | None = None, ) -> str: from litellm.router_utils.pre_call_checks.deployment_affinity_check import DeploymentAffinityCheck data, bucket = self._normalized_request(surface, session_id, caller, child) turn: Final[dict[str, object]] = ( - {"role": "user", "content": prompt} if prompt is not None else - {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "call_1", "content": "done"}]} - if surface == "messages" else - {"type": "function_call_output", "call_id": "call_1", "output": "done"} - if surface == "responses" else - {"role": "tool", "tool_call_id": "call_1", "content": "done"} + {"role": "user", "content": prompt} + if prompt is not None + else {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "call_1", "content": "done"}]} + if surface == "messages" + else {"type": "function_call_output", "call_id": "call_1", "output": "done"} + if surface == "responses" + else {"role": "tool", "tool_call_id": "call_1", "content": "done"} ) tool_call: Final = ( {"type": "function_call", "call_id": "call_1", "name": "read_file", "arguments": "{}"} - if surface == "responses" else - {"role": "assistant", "content": [{"type": "tool_use", "id": "call_1", "name": "read_file", "input": {}}]} - if surface == "messages" else TestClassificationMode.TOOL_CALL_1 + if surface == "responses" + else { + "role": "assistant", + "content": [{"type": "tool_use", "id": "call_1", "name": "read_file", "input": {}}], + } + if surface == "messages" + else TestClassificationMode.TOOL_CALL_1 ) turns: Final = [turn] if prompt is not None else [TestClassificationMode.SIMPLE_ASK, tool_call, turn] if surface == "responses": data["input"] = turns selected: Final = await router.async_get_available_deployment( - model="child-tier-router", request_kwargs=data, + model="child-tier-router", + request_kwargs=data, messages=None if surface == "responses" else turns, input=turns if surface == "responses" else None, ) metadata: Final = data[bucket] assert isinstance(metadata, dict) assert (metadata["session_id"], metadata["trace_id"], data["litellm_session_id"], data["litellm_trace_id"]) == ( - "parent-session", "parent-session", "parent-session", "parent-session" + "parent-session", + "parent-session", + "parent-session", + "parent-session", ) assert (metadata["routing_decision"]["tier"], metadata["routing_decision"]["cause"]) == ( - expected_tier, expected_cause + expected_tier, + expected_cause, ) affinity: Final = next(cb for cb in router.optional_callbacks or () if isinstance(cb, DeploymentAffinityCheck)) if expected_deployment is not None: filtered: Final = await affinity.async_filter_deployments( - model=selected["model_name"], healthy_deployments=router.get_model_list(selected["model_name"]) or [], - messages=None, request_kwargs=data, + model=selected["model_name"], + healthy_deployments=router.get_model_list(selected["model_name"]) or [], + messages=None, + request_kwargs=data, ) assert [deployment["model_info"]["id"] for deployment in filtered] == [expected_deployment] assert selected["model_info"]["id"] == expected_deployment await affinity.async_pre_call_deployment_hook( kwargs={ - **data, "model_info": selected["model_info"], + **data, + "model_info": selected["model_info"], bucket: {**metadata, "deployment_model_name": selected["model_name"]}, }, call_type=None, @@ -7845,13 +7868,29 @@ class TestChildSessionTierPins: (None, None, "SIMPLE", "heuristic_scorer", "caller-a", True), ): await self._route( - router, surface, session, prompt, tier, cause, caller, child, + router, + surface, + session, + prompt, + tier, + cause, + caller, + child, expected_deployment=parent if tier == "REASONING" and caller == "caller-a" else None, ) - assert await self._route( - router, surface, "parent-session", None, "REASONING", "user_turn_continuation", - child=False, expected_deployment=parent, - ) == parent + assert ( + await self._route( + router, + surface, + "parent-session", + None, + "REASONING", + "user_turn_continuation", + child=False, + expected_deployment=parent, + ) + == parent + ) finally: router.discard() @@ -7862,10 +7901,18 @@ class TestChildSessionTierPins: parent: Final = await self._route( router, "chat", "parent-session", "reason carefully", "REASONING", "literal_keyword_match", child=False ) - assert await self._route( - router, "chat", "child-session-a", "Hello!", "REASONING", "session_affinity_pin", - expected_deployment=parent, - ) == parent + assert ( + await self._route( + router, + "chat", + "child-session-a", + "Hello!", + "REASONING", + "session_affinity_pin", + expected_deployment=parent, + ) + == parent + ) finally: router.discard() @@ -8804,7 +8851,7 @@ class TestEscalationKeywords: one step higher so a user can force a stronger model when unhappy with results.""" @staticmethod - def _request_kwargs(session_id: str) -> Dict: + def _request_kwargs(session_id: str) -> dict: return {"metadata": {"session_id": session_id}} def test_default_escalation_keyword(self, complexity_router): @@ -9140,7 +9187,7 @@ class TestEscalationKeywords: assert result.model == "o1-b" # unchanged: no random hop to o1-a / o1-c -def _stalled_tool_history(repeats: int = 3) -> List[Dict]: +def _stalled_tool_history(repeats: int = 3) -> list[dict]: """`repeats` identical bash tool calls in a row, the automatic counterpart to a user typing an escalation keyword: the assistant, not the human, is the one stuck.""" return [ @@ -9692,7 +9739,7 @@ class TestRoutingDecisionSurvivesToSpendLogOnEveryMetadataShape: # Mirror function_setup: it copies `litellm_metadata` by value into # litellm_params AFTER the router hook has run, so the copy must carry # the decision. Reading the stash any earlier would lose it. - litellm_params: Dict = {} + litellm_params: dict = {} if "metadata" in request_kwargs: litellm_params["metadata"] = request_kwargs["metadata"] if isinstance(request_kwargs.get("litellm_metadata"), dict): @@ -9708,11 +9755,15 @@ class TestRoutingDecisionSurvivesToSpendLogOnEveryMetadataShape: assert persisted is not None, f"routing_decision dropped for {expected_bucket}" assert persisted["router_model_name"] == "smart-router" if classifier_type == "heuristic_v2": - assert persisted["heuristic_v2_forecast"] == request_kwargs[expected_bucket]["routing_decision"][ - "heuristic_v2_forecast" - ] + assert ( + persisted["heuristic_v2_forecast"] + == request_kwargs[expected_bucket]["routing_decision"]["heuristic_v2_forecast"] + ) assert set(persisted["heuristic_v2_forecast"]["probabilities"]) == { - "SIMPLE", "MEDIUM", "COMPLEX", "REASONING" + "SIMPLE", + "MEDIUM", + "COMPLEX", + "REASONING", } else: assert "heuristic_v2_forecast" not in persisted @@ -9747,7 +9798,7 @@ class TestRoutingDecisionIsPerAttempt: @pytest.mark.asyncio async def test_fallback_to_plain_model_group_clears_the_earlier_decision(self, seed, bucket): router = Router(model_list=self.MODEL_LIST) - request_kwargs: Dict = dict(seed) + request_kwargs: dict = dict(seed) messages = [{"role": "user", "content": "Hello!"}] await router.async_pre_routing_hook(model="smart-router", request_kwargs=request_kwargs, messages=messages) @@ -9767,7 +9818,7 @@ class TestRoutingDecisionIsPerAttempt: Skipping the write there would drop provenance on a successfully routed request with no error, so the shared bucket owner replaces the value.""" router = Router(model_list=self.MODEL_LIST) - request_kwargs: Dict = {"litellm_metadata": unusable_bucket} + request_kwargs: dict = {"litellm_metadata": unusable_bucket} response = await router.async_pre_routing_hook( model="smart-router", @@ -9788,7 +9839,7 @@ class TestRecordRoutingDecision: DECISION = {"router_model_name": "smart-router", "router_type": "complexity", "routed_model": "gpt-4o-mini"} def test_none_clears_a_previous_decision_from_both_buckets(self): - request_kwargs: Dict = { + request_kwargs: dict = { "metadata": {"routing_decision": self.DECISION, "keep": 1}, "litellm_metadata": {"routing_decision": self.DECISION}, } @@ -9798,7 +9849,7 @@ class TestRecordRoutingDecision: assert request_kwargs["metadata"]["keep"] == 1 def test_none_creates_no_bucket_on_a_request_that_had_none(self): - request_kwargs: Dict = {} + request_kwargs: dict = {} Router._record_routing_decision(request_kwargs=request_kwargs, routing_decision=None) assert request_kwargs == {} @@ -9957,7 +10008,7 @@ class TestRedactedLoggingDropsPromptText: MESSAGES = [{"role": "user", "content": "LITELLM ESCALATE please deploy to k8s now"}] - async def _decision(self, request_kwargs: Dict) -> Dict: + async def _decision(self, request_kwargs: dict) -> dict: router = Router(model_list=self.MODEL_LIST) response = await router.async_pre_routing_hook( model="smart-router", request_kwargs=request_kwargs, messages=self.MESSAGES @@ -10026,7 +10077,7 @@ class TestRedactedLoggingDropsPromptText: @pytest.mark.asyncio async def test_redaction_via_request_header_is_honored(self): - request_kwargs: Dict = {"metadata": {"headers": {"x-litellm-enable-message-redaction": True}}} + request_kwargs: dict = {"metadata": {"headers": {"x-litellm-enable-message-redaction": True}}} decision = await self._decision(request_kwargs) assert "matched_keyword" not in decision assert decision["cause"] == "literal_keyword_match" @@ -10042,8 +10093,10 @@ class TestRedactedLoggingDropsPromptText: { "classifier_type": "capability", "capability_classifier_config": { - "efficient_tier": "SIMPLE", "capable_tier": "REASONING", - "base_threshold": 0.5, "threshold_step": 0.1, + "efficient_tier": "SIMPLE", + "capable_tier": "REASONING", + "base_threshold": 0.5, + "threshold_step": 0.1, }, }, "capability_classifier_fallback", @@ -10052,8 +10105,10 @@ class TestRedactedLoggingDropsPromptText: { "classifier_type": "llm_v2", "llm_v2_config": { - "efficient_profile": "Routine tasks", "capable_profile": "Complex tasks", - "harness": "Read files and run checks", "max_quality_gap": 0.05, + "efficient_profile": "Routine tasks", + "capable_profile": "Complex tasks", + "harness": "Read files and run checks", + "max_quality_gap": 0.05, }, }, "llm_v2_fallback", @@ -10064,36 +10119,46 @@ async def test_classifier_failure_survives_real_router_decision_and_redaction( overrides: Mapping[str, object], expected_cause: str, monkeypatch: pytest.MonkeyPatch ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - router: Final = Router(model_list=[ - { - "model_name": "auto", - "litellm_params": { - "model": "auto_router/complexity_router", - "complexity_router_config": { - "tiers": {"SIMPLE": "cheap", "REASONING": "capable"}, - "default_model": "cheap", - "classifier_llm_config": {"model": "classifier", "timeout_ms": 2000}, - **overrides, + router: Final = Router( + model_list=[ + { + "model_name": "auto", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "tiers": {"SIMPLE": "cheap", "REASONING": "capable"}, + "default_model": "cheap", + "classifier_llm_config": {"model": "classifier", "timeout_ms": 2000}, + **overrides, + }, }, }, - }, - { - "model_name": "classifier", - "litellm_params": { - "model": "openai/test-classifier", "api_key": "test-key", - "api_base": "https://classifier.test/v1", + { + "model_name": "classifier", + "litellm_params": { + "model": "openai/test-classifier", + "api_key": "test-key", + "api_base": "https://classifier.test/v1", + }, }, - }, - {"model_name": "cheap", "litellm_params": {"model": "openai/test-cheap", "api_key": "test-key"}}, - {"model_name": "capable", "litellm_params": {"model": "openai/test-capable", "api_key": "test-key"}}, - ]) + {"model_name": "cheap", "litellm_params": {"model": "openai/test-cheap", "api_key": "test-key"}}, + {"model_name": "capable", "litellm_params": {"model": "openai/test-capable", "api_key": "test-key"}}, + ] + ) with respx.mock: upstream: Final = respx.post("https://classifier.test/v1/chat/completions").respond( 200, json={ - "id": "classifier-result", "object": "chat.completion", "model": "test-classifier", - "choices": [{"index": 0, "message": {"role": "assistant", "content": "PRIVATE_INVALID_REPLY"}, - "finish_reason": "stop"}], + "id": "classifier-result", + "object": "chat.completion", + "model": "test-classifier", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "PRIVATE_INVALID_REPLY"}, + "finish_reason": "stop", + } + ], }, ) metadata: Final[dict[str, object]] = {"headers": {"x-litellm-enable-message-redaction": True}} @@ -10110,11 +10175,21 @@ async def test_classifier_failure_survives_real_router_decision_and_redaction( assert "PRIVATE_INVALID_REPLY" not in json.dumps(decision) assert "signals" not in decision if overrides["classifier_type"] == "llm": - upstream.respond(200, json={ - "id": "classifier-recovered", "object": "chat.completion", "model": "test-classifier", - "choices": [{"index": 0, "message": {"role": "assistant", "content": '{"tier":"SIMPLE"}'}, - "finish_reason": "stop"}], - }) + upstream.respond( + 200, + json={ + "id": "classifier-recovered", + "object": "chat.completion", + "model": "test-classifier", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": '{"tier":"SIMPLE"}'}, + "finish_reason": "stop", + } + ], + }, + ) recovered: Final = await router.async_pre_routing_hook( model="auto", request_kwargs=request, messages=[{"role": "user", "content": "hi"}] ) @@ -12528,7 +12603,7 @@ class TestClassifierFallbackChoice: }, ) mock_router_instance.cache = DualCache() - request_kwargs: Dict = {"metadata": {"session_id": "session-flaky"}} + request_kwargs: dict = {"metadata": {"session_id": "session-flaky"}} mock_router_instance.acompletion = AsyncMock(side_effect=TimeoutError("classifier timed out")) first = await router.async_pre_routing_hook( @@ -12567,7 +12642,7 @@ class TestClassifierFallbackChoice: }, ) mock_router_instance.cache = DualCache() - request_kwargs: Dict = {"metadata": {"session_id": "session-steady"}} + request_kwargs: dict = {"metadata": {"session_id": "session-steady"}} mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "REASONING"}')) first = await router.async_pre_routing_hook( @@ -13045,7 +13120,7 @@ class TestClassificationRubrics: assert config.classifier_llm_config.system_prompt == "Grade the data sensitivity of the request." -def _custom_tier_config(**overrides) -> Dict: +def _custom_tier_config(**overrides) -> dict: """A valid operator-defined tier set: two built-in names plus one custom tier.""" return { "tiers": {"SIMPLE": "gpt-4o-mini", "COMPLEX": "claude-sonnet-4-20250514", "SECURITY_REVIEW": "o1-preview"}, @@ -14069,8 +14144,8 @@ async def test_tier_params_mask_credentials_in_routing_decision(route, mock_rout complexity_router_config=config, ) request_kwargs = {"metadata": {"session_id": "masked-params-session"}} + mock_router_instance.cache = DualCache() if route == "session": - mock_router_instance.cache = DualCache() await mock_router_instance.cache.async_set_cache( key=router._get_session_affinity_cache_key("masked-params-session", request_kwargs), value={"model": "opus", "tier": "REASONING"}, @@ -14816,7 +14891,9 @@ class TestContextWindowEscalation: router = ComplexityRouter( model_name="test-router", litellm_router_instance=_windowed_router(_SMALL, ("mid-model", "openai/gpt-4o-mini", 200000), _BIG), - complexity_router_config=_tier_config(tiers={"SIMPLE": ["small-model", "mid-model"], "COMPLEX": "big-model"}), + complexity_router_config=_tier_config( + tiers={"SIMPLE": ["small-model", "mid-model"], "COMPLEX": "big-model"} + ), ) result = await router.async_pre_routing_hook(model="test-router", request_kwargs={}, messages=_OVERSIZED_TURNS) @@ -14923,9 +15000,7 @@ class TestContextWindowEscalation: else MappingProxyType({}) ) raw_config: Final = RequestComplexityRouterConfig.model_validate( - MappingProxyType( - {"tiers": MappingProxyType({"SIMPLE": "small-model", "COMPLEX": "big-model"}), **setting} - ) + MappingProxyType({"tiers": MappingProxyType({"SIMPLE": "small-model", "COMPLEX": "big-model"}), **setting}) ) config: Final = ( RequestComplexityRouterConfig.model_validate_json(raw_config.model_dump_json()) @@ -15086,7 +15161,7 @@ class TestContextWindowEscalation: complexity_router_config=_tier_config(tiers={"SIMPLE": "cop-pool", "COMPLEX": "big-model"}), ) real_get_llm_provider = litellm.get_llm_provider - copilot_resolutions: List = [] + copilot_resolutions: list = [] def _guarded(*args, **kwargs): target = str(kwargs.get("model") or (args[0] if args else "")) + str( @@ -15171,9 +15246,7 @@ class TestModalityRouting: self.BASE_VISION, ) original: Final = await router.aclassify("What color is this?") - result: Final = await router.async_pre_routing_hook( - model="m", request_kwargs={}, messages=self.IMAGE_MESSAGE - ) + result: Final = await router.async_pre_routing_hook(model="m", request_kwargs={}, messages=self.IMAGE_MESSAGE) assert original.heuristic_v2_forecast is not None assert result is not None and result.routing_decision is not None @@ -17367,16 +17440,18 @@ class TestMaxTokensFromTierModel: tier_litellm_params={tier_carrier: 4321}, max_tokens_from_tier_model="clamp", marker_params={marker_carrier: 2000}, - simple_deployments=[{ - "model_name": "small", - "litellm_params": { - "model": "openai/gpt-5.6", - "api_key": "test-only", - "api_base": "https://tier.test/v1", - surface_key: 9000, - }, - "model_info": {"max_output_tokens": 16384}, - }], + simple_deployments=[ + { + "model_name": "small", + "litellm_params": { + "model": "openai/gpt-5.6", + "api_key": "test-only", + "api_base": "https://tier.test/v1", + surface_key: 9000, + }, + "model_info": {"max_output_tokens": 16384}, + } + ], ) with respx.mock(assert_all_mocked=True) as upstream: @@ -17426,10 +17501,14 @@ class TestMaxTokensFromTierModel: max_tokens_from_tier_model="clamp", marker_params={marker_carrier: 2000}, tier_litellm_params={"max_tokens": tier_cap} if tier_cap else None, - simple_deployments=None if known_ceiling else [{ - "model_name": "small", - "litellm_params": {"model": "openai/not-in-any-map", "api_key": "k"}, - }], + simple_deployments=None + if known_ceiling + else [ + { + "model_name": "small", + "litellm_params": {"model": "openai/not-in-any-map", "api_key": "k"}, + } + ], ) request: Final[dict[str, object]] = {} if caller_cap is None else {"max_output_tokens": caller_cap} await router.async_get_available_deployment( @@ -17437,14 +17516,15 @@ class TestMaxTokensFromTierModel: ) first: Final = {key: request[key] for key in OUTPUT_TOKEN_CEILING_PARAMS if key in request} await router.async_get_available_deployment( - model="smart-router", request_kwargs=request, + model="smart-router", + request_kwargs=request, messages=[{"role": "user", "content": self.COMPLEX_PROMPT}], ) expected_first: Final = ( {"max_tokens": tier_cap or caller_cap or 8192} - if known_ceiling or tier_cap else - {marker_carrier: 2000, **({"max_output_tokens": caller_cap} if caller_cap else {})} + if known_ceiling or tier_cap + else {marker_carrier: 2000, **({"max_output_tokens": caller_cap} if caller_cap else {})} ) assert first == expected_first assert request["max_tokens"] == (caller_cap or 64000) @@ -17464,10 +17544,12 @@ class TestMaxTokensFromTierModel: router: Final = self._router( max_tokens_from_tier_model=mode, marker_params={marker_carrier: 2000}, - simple_deployments=[{ - "model_name": "small", - "litellm_params": {"model": "openai/not-in-any-map", "api_key": "k"}, - }], + simple_deployments=[ + { + "model_name": "small", + "litellm_params": {"model": "openai/not-in-any-map", "api_key": "k"}, + } + ], ) sent: Final = await self._routed(router) @@ -17696,7 +17778,11 @@ class TestMaxTokensFromTierModel: ({"metadata": {CLIENT_OUTPUT_CEILING_METADATA_KEY: {}}}, {}), ({"litellm_metadata": {CLIENT_OUTPUT_CEILING_METADATA_KEY: {"max_tokens": "2000"}}}, {"max_tokens": 2000}), ( - {"metadata": {CLIENT_OUTPUT_CEILING_METADATA_KEY: {"max_output_tokens": 5, "api_key": "x", "max_tokens": True}}}, + { + "metadata": { + CLIENT_OUTPUT_CEILING_METADATA_KEY: {"max_output_tokens": 5, "api_key": "x", "max_tokens": True} + } + }, {"max_output_tokens": 5}, ), ({"metadata": {CLIENT_OUTPUT_CEILING_METADATA_KEY: "not-a-dict"}}, None), diff --git a/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py b/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py index d317f0d5741..deb1bf97538 100644 --- a/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py +++ b/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py @@ -11,7 +11,7 @@ from fakeredis.aioredis import FakeConnection from redis.asyncio import BlockingConnectionPool, Redis import litellm -from litellm.caching.affinity_cache import claim_affinity_pin +from litellm.caching.affinity_cache import claim_affinity_pin_in_memory from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import RedisCache @@ -901,7 +901,6 @@ async def test_claim_pin_falls_back_to_pod_local_when_redis_is_down(): assert second == "our-deployment" -@pytest.mark.asyncio @pytest.mark.parametrize( ("stored", "expected"), [ @@ -914,13 +913,13 @@ async def test_claim_pin_falls_back_to_pod_local_when_redis_is_down(): (None, {"model": "second"}), ], ) -async def test_eligible_affinity_claim_replaces_stale_pins_and_slides_ttl(stored: object, expected: object) -> None: +def test_eligible_affinity_claim_replaces_stale_pins_and_slides_ttl(stored: object, expected: object) -> None: clock: Final = MagicMock(return_value=100.0) cache: Final = DualCache(in_memory_cache=InMemoryCache(clock=clock)) cache.in_memory_cache.set_cache("tier-pin", stored, ttl=10) clock.return_value = 105.0 - winner: Final = await claim_affinity_pin( + winner: Final = claim_affinity_pin_in_memory( cache, "tier-pin", {"model": "second"}, @@ -936,24 +935,21 @@ async def test_eligible_affinity_claim_replaces_stale_pins_and_slides_ttl(stored assert cache.in_memory_cache.get_cache("tier-pin") is None -@pytest.mark.asyncio -async def test_concurrent_eligible_claims_return_one_winner() -> None: +def test_eligible_local_claims_keep_first_winner() -> None: cache: Final = DualCache() candidates: Final = ({"model": "first"}, {"model": "second"}) - winners: Final = await asyncio.gather( - *( - claim_affinity_pin( - cache, - "tier-pin", - candidates[index % 2], - 30, - eligible_values=candidates, - ) - for index in range(20) + winners: Final = tuple( + claim_affinity_pin_in_memory( + cache, + "tier-pin", + candidates[index % 2], + 30, + eligible_values=candidates, ) + for index in range(20) ) - assert winners == [{"model": "first"}] * 20 + assert winners == ({"model": "first"},) * 20 assert cache.in_memory_cache.get_cache("tier-pin") == {"model": "first"} From 30ffa8106653a8f2db1395a659323bd65165eab2 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 22:10:14 +0000 Subject: [PATCH 17/17] refactor(redis): move email budget alert claims into a typed store --- .../send_emails/base_email.py | 275 +++++------ litellm/redis/stores/email_alert_claims.py | 77 +++ .../send_emails/test_base_email.py | 467 +++++++----------- .../redis/stores/test_email_alert_claims.py | 293 +++++++++++ 4 files changed, 655 insertions(+), 457 deletions(-) create mode 100644 litellm/redis/stores/email_alert_claims.py create mode 100644 tests/unit/redis/stores/test_email_alert_claims.py diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py b/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py index c648f7faea0..6eed6e6b66e 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py @@ -6,7 +6,8 @@ Base class for sending emails to user after creating keys or invite links import html import json import os -from typing import Final, List, Literal, Optional +from datetime import timedelta +from typing import Final, Literal, Protocol from litellm_enterprise.types.enterprise_callbacks.send_emails import ( EmailEvent, @@ -18,6 +19,7 @@ from litellm_enterprise.types.enterprise_callbacks.send_emails import ( from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache +from litellm.caching.redis_cache import RedisCache from litellm.constants import ( EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE, EMAIL_BUDGET_ALERT_TTL, @@ -46,6 +48,17 @@ from litellm.proxy._types import ( UserAPIKeyAuth, WebhookEvent, ) +from litellm.redis.stores.email_alert_claims import ( + EmailAlertClaimsStore, + EmailBudgetAlert, + MaxBudgetAlert, + SoftBudgetAlert, + TeamMemberAlert, + ThresholdBudgetAlert, + TransportEmailAlertClaimsStore, + alert_key, +) +from litellm.redis.transport.contract import RedisFailure from litellm.repositories.table_repositories import InvitationLinkRepository from litellm.repositories.user_repository import UserRepository from litellm.secret_managers.main import get_secret_bool @@ -54,13 +67,27 @@ from litellm.types.integrations.slack_alerting import LITELLM_LOGO_URL _BUDGET_ALERT_CLAIMS_TARGET: Final = "budget_alert_claims" -def _max_budget_alert_id(user_info: CallInfo) -> str: +class _LocalAlertCache(Protocol): + async def async_increment(self, key: str, value: float, *, ttl: int) -> float: ... + + def delete_cache(self, key: str) -> None: ... + + +class _BudgetAlertCache(Protocol): + @property + def in_memory_cache(self) -> _LocalAlertCache | None: ... + + @property + def redis_cache(self) -> RedisCache | None: ... + + +def _max_budget_alert_id(user_info: CallInfo) -> str | TeamMemberAlert: if user_info.event_group == Litellm_EntityType.TEAM_MEMBER: - return f"team_member:{user_info.user_id}:{user_info.team_id}" + return TeamMemberAlert(user_info.user_id, user_info.team_id) return user_info.token or user_info.user_id or "default_id" -def _parse_email_list(raw) -> List[str]: +def _parse_email_list(raw) -> list[str]: """Parse emails from a list or comma-separated string.""" if isinstance(raw, list): return [e.strip() for e in raw if isinstance(e, str) and e.strip()] @@ -80,7 +107,7 @@ class BaseEmailLogger(CustomLogger): def __init__( self, - internal_usage_cache: Optional[DualCache] = None, + internal_usage_cache: DualCache | None = None, **kwargs, ): """ @@ -104,9 +131,7 @@ class BaseEmailLogger(CustomLogger): event_message=event.event_message, ) - verbose_proxy_logger.debug( - f"send_user_invitation_email_event: {json.dumps(event, indent=4, default=str)}" - ) + verbose_proxy_logger.debug(f"send_user_invitation_email_event: {json.dumps(event, indent=4, default=str)}") email_html_content = USER_INVITATION_EMAIL_TEMPLATE.format( email_logo_url=email_params.logo_url, @@ -124,11 +149,7 @@ class BaseEmailLogger(CustomLogger): html_body=email_html_content, ) - pass - - async def send_key_created_email( - self, send_key_created_email_event: SendKeyCreatedEmailEvent - ): + async def send_key_created_email(self, send_key_created_email_event: SendKeyCreatedEmailEvent): """ Send email to user after creating key for the user """ @@ -144,9 +165,7 @@ class BaseEmailLogger(CustomLogger): ) # Check if API key should be included in email - include_api_key = get_secret_bool( - secret_name="EMAIL_INCLUDE_API_KEY", default_value=True - ) + include_api_key = get_secret_bool(secret_name="EMAIL_INCLUDE_API_KEY", default_value=True) if include_api_key is None: include_api_key = True # Default to True if not set key_token_display = ( @@ -171,11 +190,8 @@ class BaseEmailLogger(CustomLogger): subject=email_params.subject, html_body=email_html_content, ) - pass - async def send_key_rotated_email( - self, send_key_rotated_email_event: SendKeyRotatedEmailEvent - ): + async def send_key_rotated_email(self, send_key_rotated_email_event: SendKeyRotatedEmailEvent): """ Send email to user after rotating key for the user """ @@ -191,9 +207,7 @@ class BaseEmailLogger(CustomLogger): ) # Check if API key should be included in email - include_api_key = get_secret_bool( - secret_name="EMAIL_INCLUDE_API_KEY", default_value=True - ) + include_api_key = get_secret_bool(secret_name="EMAIL_INCLUDE_API_KEY", default_value=True) if include_api_key is None: include_api_key = True # Default to True if not set key_token_display = ( @@ -218,7 +232,6 @@ class BaseEmailLogger(CustomLogger): subject=email_params.subject, html_body=email_html_content, ) - pass async def send_soft_budget_alert_email(self, event: WebhookEvent): """ @@ -236,9 +249,7 @@ class BaseEmailLogger(CustomLogger): ) # Format budget values - soft_budget_str = ( - f"${event.soft_budget}" if event.soft_budget is not None else "N/A" - ) + soft_budget_str = f"${event.soft_budget}" if event.soft_budget is not None else "N/A" spend_str = f"${event.spend}" if event.spend is not None else "$0.00" max_budget_info = "" if event.max_budget is not None: @@ -260,7 +271,6 @@ class BaseEmailLogger(CustomLogger): subject=email_params.subject, html_body=email_html_content, ) - pass async def send_team_soft_budget_alert_email(self, event: WebhookEvent): """ @@ -268,7 +278,7 @@ class BaseEmailLogger(CustomLogger): Supports multiple recipients via alert_emails field from team metadata """ # Collect all recipient emails - recipient_emails: List[str] = [] + recipient_emails: list[str] = [] # Add additional alert emails from team metadata.soft_budget_alert_emails if hasattr(event, "alert_emails") and event.alert_emails: @@ -306,9 +316,7 @@ class BaseEmailLogger(CustomLogger): ) # Format budget values - soft_budget_str = ( - f"${event.soft_budget}" if event.soft_budget is not None else "N/A" - ) + soft_budget_str = f"${event.soft_budget}" if event.soft_budget is not None else "N/A" spend_str = f"${event.spend}" if event.spend is not None else "$0.00" max_budget_info = "" if event.max_budget is not None: @@ -335,13 +343,12 @@ class BaseEmailLogger(CustomLogger): subject=email_params.subject, html_body=email_html_content, ) - pass async def send_max_budget_alert_email( self, event: WebhookEvent, - threshold_pct: Optional[int] = None, - recipient_emails: Optional[List[str]] = None, + threshold_pct: int | None = None, + recipient_emails: list[str] | None = None, ): """ Send email to user when max budget alert threshold is reached. @@ -359,20 +366,14 @@ class BaseEmailLogger(CustomLogger): # Format budget values spend_str = f"${event.spend}" if event.spend is not None else "$0.00" - max_budget_str = ( - f"${event.max_budget}" if event.max_budget is not None else "N/A" - ) + max_budget_str = f"${event.max_budget}" if event.max_budget is not None else "N/A" # Calculate percentage and alert threshold - percentage = threshold_pct if threshold_pct is not None else int( - EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE * 100 + percentage = ( + threshold_pct if threshold_pct is not None else int(EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE * 100) ) threshold_fraction = percentage / 100.0 - alert_threshold_str = ( - f"${event.max_budget * threshold_fraction:.2f}" - if event.max_budget is not None - else "N/A" - ) + alert_threshold_str = f"${event.max_budget * threshold_fraction:.2f}" if event.max_budget is not None else "N/A" if recipient_emails: # Multi-threshold path: batch send with generic key-based greeting @@ -382,9 +383,7 @@ class BaseEmailLogger(CustomLogger): user_email=event.user_email or recipient_emails[0], event_message=event.event_message, ) - greeting = html.escape( - event.user_email or event.key_alias or event.token or "" - ) + greeting = html.escape(event.user_email or event.key_alias or event.token or "") if event.event_group == Litellm_EntityType.TEAM_MEMBER: email_html_content = TEAM_MEMBER_MAX_BUDGET_ALERT_EMAIL_TEMPLATE.format( email_logo_url=email_params.logo_url, @@ -488,10 +487,7 @@ class BaseEmailLogger(CustomLogger): # For non-team alerts, require either max_budget or soft_budget if user_info.max_budget is None and user_info.soft_budget is None: return - if ( - user_info.soft_budget is not None - and user_info.spend >= user_info.soft_budget - ): + if user_info.soft_budget is not None and user_info.spend >= user_info.soft_budget: # Generate cache key based on event type and identifier # Use appropriate ID based on event_group to ensure unique cache keys per entity type if user_info.event_group == Litellm_EntityType.TEAM: @@ -503,13 +499,8 @@ class BaseEmailLogger(CustomLogger): else: # For KEY and other types, use token or user_id _id = user_info.token or user_info.user_id or "default_id" - _cache_key = f"email_budget_alerts:soft_budget_crossed:{_id}" - - send_count = await _cache.async_increment_cache( - key=_cache_key, - value=1, - ttl=EMAIL_BUDGET_ALERT_TTL, - ) + soft_alert: Final = SoftBudgetAlert(_id) + send_count = await self._claim_budget_alert(_cache, soft_alert) if send_count is None or send_count <= 1: # Create WebhookEvent for soft budget alert event_message = f"Soft Budget Crossed - Total Soft Budget: ${user_info.soft_budget}" @@ -544,7 +535,7 @@ class BaseEmailLogger(CustomLogger): f"Error sending soft budget alert email: {e}", exc_info=True, ) - await self._release_budget_alert_claim(_cache, _cache_key) + await self._release_budget_alert_claim(_cache, soft_alert) return # For max_budget_alert, check if we've already sent an alert @@ -552,34 +543,20 @@ class BaseEmailLogger(CustomLogger): if user_info.max_budget is not None and user_info.spend is not None: if user_info.max_budget_alert_emails: # New path: multi-threshold alerts - await self._handle_multi_threshold_max_budget_alert( - user_info=user_info, _cache=_cache - ) + await self._handle_multi_threshold_max_budget_alert(user_info=user_info, _cache=_cache) return - alert_threshold = ( - user_info.max_budget * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE - ) + alert_threshold = user_info.max_budget * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE # Only alert if we've crossed the threshold but haven't exceeded max_budget yet - if ( - user_info.spend >= alert_threshold - and user_info.spend < user_info.max_budget - ): + if user_info.spend >= alert_threshold and user_info.spend < user_info.max_budget: # Generate cache key based on event type and identifier _id = user_info.token or user_info.user_id or "default_id" - _cache_key = f"email_budget_alerts:max_budget_alert:{_id}" - - send_count = await _cache.async_increment_cache( - key=_cache_key, - value=1, - ttl=EMAIL_BUDGET_ALERT_TTL, - ) + max_alert: Final = MaxBudgetAlert(_id) + send_count = await self._claim_budget_alert(_cache, max_alert) if send_count is None or send_count <= 1: # Calculate percentage - percentage = int( - EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE * 100 - ) + percentage = int(EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE * 100) # Create WebhookEvent for max budget alert event_message = f"Max Budget Alert - {percentage}% of Maximum Budget Reached" @@ -609,7 +586,7 @@ class BaseEmailLogger(CustomLogger): f"Error sending max budget alert email: {e}", exc_info=True, ) - await self._release_budget_alert_claim(_cache, _cache_key) + await self._release_budget_alert_claim(_cache, max_alert) return @with_service_target(_BUDGET_ALERT_CLAIMS_TARGET) @@ -636,9 +613,7 @@ class BaseEmailLogger(CustomLogger): continue _id = _max_budget_alert_id(user_info) - _cache_key = ( - f"email_budget_alerts:max_budget_alert:{threshold_pct}:{_id}" - ) + threshold_alert = ThresholdBudgetAlert(_id, threshold_pct) # Parse emails + auto-include owner emails = _parse_email_list(raw_emails) @@ -653,11 +628,7 @@ class BaseEmailLogger(CustomLogger): continue recipient_emails = list(set(emails)) - send_count = await _cache.async_increment_cache( - key=_cache_key, - value=1, - ttl=EMAIL_BUDGET_ALERT_TTL, - ) + send_count = await self._claim_budget_alert(_cache, threshold_alert) if send_count is not None and send_count > 1: continue @@ -696,24 +667,44 @@ class BaseEmailLogger(CustomLogger): f"Error sending multi-threshold max budget alert email for {threshold_pct}%: {e}", exc_info=True, ) - await self._release_budget_alert_claim(_cache, _cache_key) + await self._release_budget_alert_claim(_cache, threshold_alert) @with_service_target(_BUDGET_ALERT_CLAIMS_TARGET) - async def _release_budget_alert_claim(self, cache: DualCache, cache_key: str) -> None: + async def _claim_budget_alert(self, cache: _BudgetAlertCache, alert: EmailBudgetAlert) -> float | None: try: - await cache.async_delete_cache(key=cache_key) - except Exception: - verbose_proxy_logger.debug( - "Failed to release budget alert claim for %s; it expires with the TTL", - cache_key, + local_count: Final = ( + await cache.in_memory_cache.async_increment(alert_key(alert), 1, ttl=EMAIL_BUDGET_ALERT_TTL) + if cache.in_memory_cache is not None + else None ) + except Exception: + return None + if cache.redis_cache is None: + return local_count + try: + store: Final[EmailAlertClaimsStore] = TransportEmailAlertClaimsStore(cache.redis_cache.transport()) + count: Final = await store.claim(alert, ttl=timedelta(seconds=EMAIL_BUDGET_ALERT_TTL)) + return local_count if isinstance(count, RedisFailure) else count + except Exception: + return local_count + + @with_service_target(_BUDGET_ALERT_CLAIMS_TARGET) + async def _release_budget_alert_claim(self, cache: _BudgetAlertCache, alert: EmailBudgetAlert) -> None: + try: + if cache.in_memory_cache is not None: + cache.in_memory_cache.delete_cache(alert_key(alert)) + if cache.redis_cache is not None: + store: Final[EmailAlertClaimsStore] = TransportEmailAlertClaimsStore(cache.redis_cache.transport()) + await store.release_failed_send(alert) + except Exception: + return async def _get_email_params( self, email_event: EmailEvent, - user_id: Optional[str] = None, - user_email: Optional[str] = None, - event_message: Optional[str] = None, + user_id: str | None = None, + user_email: str | None = None, + event_message: str | None = None, ) -> EmailParams: """ Get common email parameters used across different email sending methods @@ -738,12 +729,8 @@ class BaseEmailLogger(CustomLogger): unused_custom_fields = [] # Function to safely get custom value or default - def get_custom_or_default( - custom_value: Optional[str], default_value: str, field_name: str - ) -> str: - if ( - custom_value is not None - ): # Only check premium if trying to use custom value + def get_custom_or_default(custom_value: str | None, default_value: str, field_name: str) -> str: + if custom_value is not None: # Only check premium if trying to use custom value from litellm.proxy.proxy_server import premium_user if premium_user is not True: @@ -754,15 +741,9 @@ class BaseEmailLogger(CustomLogger): # Get parameters, falling back to defaults if custom values aren't allowed logo_url = get_custom_or_default(custom_logo, LITELLM_LOGO_URL, "logo URL") - support_contact = get_custom_or_default( - custom_support, self.DEFAULT_SUPPORT_EMAIL, "support contact" - ) - base_url = os.getenv( - "PROXY_BASE_URL", "http://0.0.0.0:4000" - ) # Not a premium feature - signature = get_custom_or_default( - custom_signature, EMAIL_FOOTER, "email signature" - ) + support_contact = get_custom_or_default(custom_support, self.DEFAULT_SUPPORT_EMAIL, "support contact") + base_url = os.getenv("PROXY_BASE_URL", "http://0.0.0.0:4000") # Not a premium feature + signature = get_custom_or_default(custom_signature, EMAIL_FOOTER, "email signature") # Get custom subject template based on email event type if email_event == EmailEvent.new_user_invitation: @@ -787,25 +768,15 @@ class BaseEmailLogger(CustomLogger): else: subject_template = "LiteLLM: {event_message}" - subject = ( - subject_template.format(event_message=event_message) - if event_message - else "LiteLLM Notification" - ) + subject = subject_template.format(event_message=event_message) if event_message else "LiteLLM Notification" - recipient_email: Optional[str] = ( - user_email or await self._lookup_user_email_from_db(user_id=user_id) - ) + recipient_email: str | None = user_email or await self._lookup_user_email_from_db(user_id=user_id) if recipient_email is None: - raise ValueError( - f"User email not found for user_id: {user_id}. User email is required to send email." - ) + raise ValueError(f"User email not found for user_id: {user_id}. User email is required to send email.") # if user invited event then send invitation link if email_event == EmailEvent.new_user_invitation: - base_url = await self._get_invitation_link( - user_id=user_id, base_url=base_url - ) + base_url = await self._get_invitation_link(user_id=user_id, base_url=base_url) # If any custom fields were not applied, log a warning if unused_custom_fields: @@ -826,7 +797,7 @@ class BaseEmailLogger(CustomLogger): signature=signature, ) - def _format_key_budget(self, max_budget: Optional[float]) -> str: + def _format_key_budget(self, max_budget: float | None) -> str: """ Format the key budget to be displayed in the email """ @@ -834,27 +805,23 @@ class BaseEmailLogger(CustomLogger): return "No budget" return f"${max_budget}" - async def _lookup_user_email_from_db(self, user_id: Optional[str]) -> Optional[str]: + async def _lookup_user_email_from_db(self, user_id: str | None) -> str | None: """ Lookup user email from user_id """ from litellm.proxy.proxy_server import prisma_client if prisma_client is None: - verbose_proxy_logger.debug( - f"Prisma client not found. Unable to lookup user email for user_id: {user_id}" - ) + verbose_proxy_logger.debug(f"Prisma client not found. Unable to lookup user email for user_id: {user_id}") return None - user_row = await UserRepository(prisma_client).table.find_unique( - where={"user_id": user_id} - ) + user_row = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_id}) if user_row is not None: return user_row.user_email return None - async def _get_invitation_link(self, user_id: Optional[str], base_url: str) -> str: + async def _get_invitation_link(self, user_id: str | None, base_url: str) -> str: """ Get invitation link for the user """ @@ -883,9 +850,7 @@ class BaseEmailLogger(CustomLogger): return base_url invitation_link = self._construct_invitation_link(invitation.id, base_url) - verbose_proxy_logger.info( - f"Successfully created invitation link for user_id: {user_id}" - ) + verbose_proxy_logger.info(f"Successfully created invitation link for user_id: {user_id}") return invitation_link async def _is_prisma_client_available(self) -> bool: @@ -893,9 +858,7 @@ class BaseEmailLogger(CustomLogger): from litellm.proxy.proxy_server import prisma_client if prisma_client is None: - verbose_proxy_logger.debug( - "Prisma client not found. Unable to lookup invitation" - ) + verbose_proxy_logger.debug("Prisma client not found. Unable to lookup invitation") return False return True @@ -923,39 +886,29 @@ class BaseEmailLogger(CustomLogger): from litellm.proxy.proxy_server import prisma_client if prisma_client is None: - verbose_proxy_logger.error( - "Prisma client is None in _get_or_create_invitation" - ) + verbose_proxy_logger.error("Prisma client is None in _get_or_create_invitation") return None try: # Try to get existing invitation - existing_invitations = ( - await InvitationLinkRepository(prisma_client).table.find_many( - where={"user_id": user_id}, - order={"created_at": "desc"}, - ) + existing_invitations = await InvitationLinkRepository(prisma_client).table.find_many( + where={"user_id": user_id}, + order={"created_at": "desc"}, ) if existing_invitations and len(existing_invitations) > 0: - verbose_proxy_logger.debug( - f"Found existing invitation for user_id: {user_id}" - ) + verbose_proxy_logger.debug(f"Found existing invitation for user_id: {user_id}") return existing_invitations[0] # Create new invitation if none exists - verbose_proxy_logger.debug( - f"Creating new invitation for user_id: {user_id}" - ) + verbose_proxy_logger.debug(f"Creating new invitation for user_id: {user_id}") return await create_invitation_for_user( data=InvitationNew(user_id=user_id), user_api_key_dict=UserAPIKeyAuth(user_id=user_id), ) except Exception as e: - verbose_proxy_logger.error( - f"Error getting/creating invitation for user_id {user_id}: {e}" - ) + verbose_proxy_logger.error(f"Error getting/creating invitation for user_id {user_id}: {e}") return None def _construct_invitation_link(self, invitation_id: str, base_url: str) -> str: @@ -971,7 +924,7 @@ class BaseEmailLogger(CustomLogger): async def send_email( self, from_email: str, - to_email: List[str], + to_email: list[str], subject: str, html_body: str, ): diff --git a/litellm/redis/stores/email_alert_claims.py b/litellm/redis/stores/email_alert_claims.py new file mode 100644 index 00000000000..c44d1390332 --- /dev/null +++ b/litellm/redis/stores/email_alert_claims.py @@ -0,0 +1,77 @@ +from dataclasses import dataclass +from datetime import timedelta +from typing import Final, Protocol, TypeAlias + +from litellm.redis.transport.contract import RedisFailure, RedisTransport + + +@dataclass(frozen=True, slots=True) +class SoftBudgetAlert: + identifier: str + + +@dataclass(frozen=True, slots=True) +class MaxBudgetAlert: + identifier: str + + +@dataclass(frozen=True, slots=True) +class TeamMemberAlert: + user_id: str | None + team_id: str | None + + +@dataclass(frozen=True, slots=True) +class ThresholdBudgetAlert: + identifier: str | TeamMemberAlert + threshold_pct: int + + +EmailBudgetAlert: TypeAlias = SoftBudgetAlert | MaxBudgetAlert | ThresholdBudgetAlert + + +def alert_key(alert: EmailBudgetAlert) -> str: + match alert: + case SoftBudgetAlert(identifier): + return f"email_budget_alerts:soft_budget_crossed:{identifier}" + case MaxBudgetAlert(identifier): + return f"email_budget_alerts:max_budget_alert:{identifier}" + case ThresholdBudgetAlert(identifier, threshold_pct): + subject: Final = ( + f"team_member:{identifier.user_id}:{identifier.team_id}" + if isinstance(identifier, TeamMemberAlert) + else identifier + ) + return f"email_budget_alerts:max_budget_alert:{threshold_pct}:{subject}" + + +class EmailAlertClaimsStore(Protocol): + async def claim(self, alert: EmailBudgetAlert, *, ttl: timedelta) -> float | RedisFailure: ... + + async def release_failed_send(self, alert: EmailBudgetAlert) -> bool | RedisFailure: ... + + +@dataclass(frozen=True, slots=True) +class TransportEmailAlertClaimsStore: + transport: RedisTransport + + async def claim(self, alert: EmailBudgetAlert, *, ttl: timedelta) -> float | RedisFailure: + key: Final = alert_key(alert) + batch: Final = self.transport.batch(transaction=False) + increment: Final = batch.incrbyfloat(key, 1) + expiry: Final = batch.ttl(key) + count: Final = await increment + current_ttl: Final = await expiry + if isinstance(count, RedisFailure): + return count + if isinstance(current_ttl, RedisFailure): + return current_ttl + if current_ttl == -1: + expired: Final = await self.transport.expire_seconds(key, ttl) + if isinstance(expired, RedisFailure): + return expired + return count + + async def release_failed_send(self, alert: EmailBudgetAlert) -> bool | RedisFailure: + deleted: Final = await self.transport.delete(alert_key(alert)) + return deleted if isinstance(deleted, RedisFailure) else deleted > 0 diff --git a/tests/unit/enterprise/enterprise_callbacks/send_emails/test_base_email.py b/tests/unit/enterprise/enterprise_callbacks/send_emails/test_base_email.py index 6cfdb75fd01..11323a87ca8 100644 --- a/tests/unit/enterprise/enterprise_callbacks/send_emails/test_base_email.py +++ b/tests/unit/enterprise/enterprise_callbacks/send_emails/test_base_email.py @@ -1,33 +1,52 @@ import asyncio -import json import os -import unittest.mock as mock +from collections.abc import AsyncIterator from types import SimpleNamespace from typing import Final +from unittest import mock from unittest.mock import patch +import fakeredis import pytest -from fastapi.testclient import TestClient - -from litellm.caching.caching import DualCache +import pytest_asyncio from litellm_enterprise.enterprise_callbacks.send_emails.base_email import ( BaseEmailLogger, ) - from litellm_enterprise.types.enterprise_callbacks.send_emails import ( EmailEvent, SendKeyCreatedEmailEvent, SendKeyRotatedEmailEvent, ) +from litellm.caching.caching import DualCache +from litellm.caching.redis_cache import RedisCache +from litellm.constants import EMAIL_BUDGET_ALERT_TTL from litellm.integrations.email_templates.email_footer import EMAIL_FOOTER from litellm.proxy._types import ( CallInfo, - LiteLLM_UserTable, Litellm_EntityType, + LiteLLM_UserTable, WebhookEvent, ) -from litellm.constants import EMAIL_BUDGET_ALERT_TTL +from tests._support.fake_redis import fake_redis_cache, local_get + + +@pytest_asyncio.fixture(loop_scope="function") +async def alert_redis() -> AsyncIterator[tuple[RedisCache, fakeredis.FakeAsyncRedis, fakeredis.FakeServer]]: + server: Final = fakeredis.FakeServer() + cache: Final = fake_redis_cache(server, namespace="alerts") + raw: Final = fakeredis.FakeAsyncRedis(server=server) + try: + yield cache, raw, server + finally: + await raw.aclose() # pyright: ignore[reportUnknownMemberType, reportAttributeAccessIssue] # fakeredis stubs predate aclose + + +@pytest.fixture +def redis_email_logger( + alert_redis: tuple[RedisCache, fakeredis.FakeAsyncRedis, fakeredis.FakeServer], +) -> BaseEmailLogger: + return BaseEmailLogger(internal_usage_cache=DualCache(redis_cache=alert_redis[0])) @pytest.fixture(autouse=True) @@ -51,9 +70,7 @@ def mock_send_email(): @pytest.fixture def mock_lookup_user_email(): - with mock.patch.object( - BaseEmailLogger, "_lookup_user_email_from_db" - ) as mock_lookup: + with mock.patch.object(BaseEmailLogger, "_lookup_user_email_from_db") as mock_lookup: yield mock_lookup @@ -66,9 +83,7 @@ def test_format_key_budget(base_email_logger): @pytest.mark.asyncio -async def test_send_key_created_email( - base_email_logger, mock_send_email, mock_lookup_user_email -): +async def test_send_key_created_email(base_email_logger, mock_send_email, mock_lookup_user_email): # Setup test data event = SendKeyCreatedEmailEvent( user_id="test_user", @@ -104,9 +119,7 @@ async def test_send_key_created_email( @pytest.mark.asyncio -async def test_send_user_invitation_email( - base_email_logger, mock_send_email, mock_lookup_user_email -): +async def test_send_user_invitation_email(base_email_logger, mock_send_email, mock_lookup_user_email): # Setup test data event = WebhookEvent( user_id="test_user", @@ -139,9 +152,7 @@ async def test_send_user_invitation_email( @pytest.mark.asyncio -async def test_send_user_invitation_email_from_db( - base_email_logger, mock_send_email, mock_lookup_user_email -): +async def test_send_user_invitation_email_from_db(base_email_logger, mock_send_email, mock_lookup_user_email): # Setup test data with no direct email but one in the database event = WebhookEvent( user_id="test_user", @@ -177,9 +188,7 @@ async def test_send_user_invitation_email_from_db( @pytest.mark.asyncio -async def test_send_user_invitation_email_no_email( - base_email_logger, mock_lookup_user_email -): +async def test_send_user_invitation_email_no_email(base_email_logger, mock_lookup_user_email): # Setup test data with no email event = WebhookEvent( user_id="test_user", @@ -198,9 +207,7 @@ async def test_send_user_invitation_email_no_email( @pytest.mark.asyncio -async def test_send_key_created_email_no_email( - base_email_logger, mock_lookup_user_email -): +async def test_send_key_created_email_no_email(base_email_logger, mock_lookup_user_email): # Setup test data with no email event = SendKeyCreatedEmailEvent( user_id="test_user", @@ -222,9 +229,7 @@ async def test_send_key_created_email_no_email( @pytest.mark.asyncio -async def test_send_key_rotated_email( - base_email_logger, mock_send_email, mock_lookup_user_email -): +async def test_send_key_rotated_email(base_email_logger, mock_send_email, mock_lookup_user_email): """ Test that send_key_rotated_email sends an email with the correct parameters and content """ @@ -262,9 +267,7 @@ async def test_send_key_rotated_email( @pytest.mark.asyncio -async def test_send_key_created_email_without_key( - base_email_logger, mock_send_email, mock_lookup_user_email -): +async def test_send_key_created_email_without_key(base_email_logger, mock_send_email, mock_lookup_user_email): """ Test that send_key_created_email hides the API key when EMAIL_INCLUDE_API_KEY is false """ @@ -291,16 +294,11 @@ async def test_send_key_created_email_without_key( mock_send_email.assert_called_once() call_args = mock_send_email.call_args[1] assert "sk-secret-key-456" not in call_args["html_body"] - assert ( - "[Key hidden for security - retrieve from dashboard]" - in call_args["html_body"] - ) + assert "[Key hidden for security - retrieve from dashboard]" in call_args["html_body"] @pytest.mark.asyncio -async def test_send_key_rotated_email_without_key( - base_email_logger, mock_send_email, mock_lookup_user_email -): +async def test_send_key_rotated_email_without_key(base_email_logger, mock_send_email, mock_lookup_user_email): """ Test that send_key_rotated_email hides the API key when EMAIL_INCLUDE_API_KEY is false """ @@ -328,10 +326,7 @@ async def test_send_key_rotated_email_without_key( mock_send_email.assert_called_once() call_args = mock_send_email.call_args[1] assert "sk-secret-rotated-789" not in call_args["html_body"] - assert ( - "[Key hidden for security - retrieve from dashboard]" - in call_args["html_body"] - ) + assert "[Key hidden for security - retrieve from dashboard]" in call_args["html_body"] @pytest.mark.asyncio @@ -350,17 +345,11 @@ async def test_get_invitation_link(base_email_logger): with mock.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): # Test with valid user_id - result = await base_email_logger._get_invitation_link( - user_id="test-user", base_url="http://test.com" - ) - assert ( - result == "http://test.com/ui/onboarding?invitation_id=test-invitation-id" - ) + result = await base_email_logger._get_invitation_link(user_id="test-user", base_url="http://test.com") + assert result == "http://test.com/ui/onboarding?invitation_id=test-invitation-id" # Test with None user_id - result = await base_email_logger._get_invitation_link( - user_id=None, base_url="http://test.com" - ) + result = await base_email_logger._get_invitation_link(user_id=None, base_url="http://test.com") assert result == "http://test.com" # Test with no invitation links @@ -368,17 +357,13 @@ async def test_get_invitation_link(base_email_logger): return [] mock_prisma.db.litellm_invitationlink.find_many = mock_find_many_empty - result = await base_email_logger._get_invitation_link( - user_id="test-user", base_url="http://test.com" - ) + result = await base_email_logger._get_invitation_link(user_id="test-user", base_url="http://test.com") assert result == "http://test.com" def test_construct_invitation_link(base_email_logger): # Test invitation link construction - result = base_email_logger._construct_invitation_link( - invitation_id="test-id-123", base_url="http://test.com" - ) + result = base_email_logger._construct_invitation_link(invitation_id="test-id-123", base_url="http://test.com") assert result == "http://test.com/ui/onboarding?invitation_id=test-id-123" @@ -404,9 +389,7 @@ async def test_get_invitation_link_creates_new_when_none_exist(base_email_logger return_value=mock_created_invitation, ) as mock_create_invitation: # Execute - result = await base_email_logger._get_invitation_link( - user_id="test-user", base_url="http://test.com" - ) + result = await base_email_logger._get_invitation_link(user_id="test-user", base_url="http://test.com") # Verify that create_invitation_for_user was called mock_create_invitation.assert_called_once() @@ -415,10 +398,7 @@ async def test_get_invitation_link_creates_new_when_none_exist(base_email_logger assert call_args["user_api_key_dict"].user_id == "test-user" # Verify the returned link uses the new invitation ID - assert ( - result - == "http://test.com/ui/onboarding?invitation_id=new-invitation-id" - ) + assert result == "http://test.com/ui/onboarding?invitation_id=new-invitation-id" @pytest.mark.asyncio @@ -441,18 +421,13 @@ async def test_get_invitation_link_uses_existing_when_available(base_email_logge "litellm.proxy.management_helpers.user_invitation.create_invitation_for_user" ) as mock_create_invitation: # Execute - result = await base_email_logger._get_invitation_link( - user_id="test-user", base_url="http://test.com" - ) + result = await base_email_logger._get_invitation_link(user_id="test-user", base_url="http://test.com") # Verify that create_invitation_for_user was NOT called mock_create_invitation.assert_not_called() # Verify the returned link uses the existing invitation ID - assert ( - result - == "http://test.com/ui/onboarding?invitation_id=existing-invitation-id" - ) + assert result == "http://test.com/ui/onboarding?invitation_id=existing-invitation-id" @pytest.mark.asyncio @@ -477,9 +452,7 @@ async def test_get_invitation_link_creates_new_when_list_is_none(base_email_logg return_value=mock_created_invitation, ) as mock_create_invitation: # Execute - result = await base_email_logger._get_invitation_link( - user_id="test-user", base_url="http://test.com" - ) + result = await base_email_logger._get_invitation_link(user_id="test-user", base_url="http://test.com") # Verify that create_invitation_for_user was called mock_create_invitation.assert_called_once() @@ -488,16 +461,11 @@ async def test_get_invitation_link_creates_new_when_list_is_none(base_email_logg assert call_args["user_api_key_dict"].user_id == "test-user" # Verify the returned link uses the new invitation ID - assert ( - result - == "http://test.com/ui/onboarding?invitation_id=new-invitation-from-none" - ) + assert result == "http://test.com/ui/onboarding?invitation_id=new-invitation-from-none" @pytest.mark.asyncio -async def test_get_email_params_user_invitation( - base_email_logger, mock_lookup_user_email -): +async def test_get_email_params_user_invitation(base_email_logger, mock_lookup_user_email): # Mock environment variables with mock.patch.dict( os.environ, @@ -520,14 +488,9 @@ async def test_get_email_params_user_invitation( user_email="test@example.com", ) - assert ( - result.logo_url - == "https://litellm-listing.s3.amazonaws.com/litellm_logo.png" - ) + assert result.logo_url == "https://litellm-listing.s3.amazonaws.com/litellm_logo.png" assert result.support_contact == "support@berri.ai" - assert ( - result.base_url == "http://test.com/ui/onboarding?invitation_id=test-id" - ) + assert result.base_url == "http://test.com/ui/onboarding?invitation_id=test-id" assert result.recipient_email == "test@example.com" @@ -593,10 +556,7 @@ async def test_get_email_params_non_premium_user(mock_env_vars): # Should use default values even though custom values are set in env assert email_params.subject == "LiteLLM: New User Invitation" assert email_params.signature == EMAIL_FOOTER - assert ( - email_params.logo_url - == "https://litellm-listing.s3.amazonaws.com/litellm_logo.png" - ) + assert email_params.logo_url == "https://litellm-listing.s3.amazonaws.com/litellm_logo.png" assert email_params.support_contact == "support@berri.ai" # Test key created email params @@ -644,9 +604,7 @@ async def test_get_email_params_default_templates(monkeypatch): @pytest.mark.asyncio -async def test_send_soft_budget_alert_email( - base_email_logger, mock_send_email, mock_lookup_user_email -): +async def test_send_soft_budget_alert_email(base_email_logger, mock_send_email, mock_lookup_user_email): """Test that send_soft_budget_alert_email sends an email with the correct parameters and content""" event = WebhookEvent( user_id="test_user", @@ -673,19 +631,14 @@ async def test_send_soft_budget_alert_email( call_args = mock_send_email.call_args[1] assert call_args["from_email"] == BaseEmailLogger.DEFAULT_LITELLM_EMAIL assert call_args["to_email"] == ["test@example.com"] - assert ( - call_args["subject"] - == "LiteLLM: Soft Budget Crossed - Total Soft Budget: $100.0" - ) + assert call_args["subject"] == "LiteLLM: Soft Budget Crossed - Total Soft Budget: $100.0" assert "$100.0" in call_args["html_body"] # soft_budget assert "$105.0" in call_args["html_body"] # spend assert "$200.0" in call_args["html_body"] # max_budget @pytest.mark.asyncio -async def test_send_soft_budget_alert_email_no_max_budget( - base_email_logger, mock_send_email, mock_lookup_user_email -): +async def test_send_soft_budget_alert_email_no_max_budget(base_email_logger, mock_send_email, mock_lookup_user_email): """Test that send_soft_budget_alert_email handles missing max_budget correctly""" event = WebhookEvent( user_id="test_user", @@ -710,13 +663,15 @@ async def test_send_soft_budget_alert_email_no_max_budget( call_args = mock_send_email.call_args[1] assert "$100.0" in call_args["html_body"] # soft_budget assert "$105.0" in call_args["html_body"] # spend - assert ( - "Maximum Budget" not in call_args["html_body"] - ) # max_budget should not be shown + assert "Maximum Budget" not in call_args["html_body"] # max_budget should not be shown @pytest.mark.asyncio -async def test_budget_alerts_soft_budget_crossed(base_email_logger, mock_send_email): +async def test_budget_alerts_soft_budget_crossed( + redis_email_logger: BaseEmailLogger, + mock_send_email, + alert_redis: tuple[RedisCache, fakeredis.FakeAsyncRedis, fakeredis.FakeServer], +): """Test that budget_alerts sends email when soft budget is crossed""" user_info = CallInfo( user_id="test_user", @@ -727,39 +682,29 @@ async def test_budget_alerts_soft_budget_crossed(base_email_logger, mock_send_em event_group=Litellm_EntityType.USER, ) - # Mock the cache so the claim is won (increment returns 1) - mock_cache = mock.AsyncMock() - mock_cache.async_increment_cache = mock.AsyncMock(return_value=1) - base_email_logger.internal_usage_cache = mock_cache - with mock.patch.dict( os.environ, { "PROXY_BASE_URL": "http://test.com", }, ): - await base_email_logger.budget_alerts(type="soft_budget", user_info=user_info) + await redis_email_logger.budget_alerts(type="soft_budget", user_info=user_info) # Verify email was sent mock_send_email.assert_called_once() call_args = mock_send_email.call_args[1] assert call_args["to_email"] == ["test@example.com"] - # Verify the send slot was claimed to prevent duplicate alerts - mock_cache.async_increment_cache.assert_called_once() - cache_call_args = mock_cache.async_increment_cache.call_args[1] + assert await alert_redis[1].get("alerts:email_budget_alerts:soft_budget_crossed:test_user") == b"1" assert ( - cache_call_args["key"] - == "email_budget_alerts:soft_budget_crossed:test_user" + 0 + < await alert_redis[1].pttl("alerts:email_budget_alerts:soft_budget_crossed:test_user") + <= EMAIL_BUDGET_ALERT_TTL * 1000 ) - assert cache_call_args["value"] == 1 - assert cache_call_args["ttl"] == EMAIL_BUDGET_ALERT_TTL @pytest.mark.asyncio -async def test_budget_alerts_soft_budget_not_crossed( - base_email_logger, mock_send_email -): +async def test_budget_alerts_soft_budget_not_crossed(base_email_logger, mock_send_email): """Test that budget_alerts does not send email when soft budget is not crossed""" user_info = CallInfo( user_id="test_user", @@ -770,9 +715,6 @@ async def test_budget_alerts_soft_budget_not_crossed( event_group=Litellm_EntityType.USER, ) - mock_cache = mock.AsyncMock() - base_email_logger.internal_usage_cache = mock_cache - await base_email_logger.budget_alerts(type="soft_budget", user_info=user_info) # Verify email was NOT sent @@ -781,7 +723,9 @@ async def test_budget_alerts_soft_budget_not_crossed( @pytest.mark.asyncio async def test_budget_alerts_soft_budget_duplicate_prevention( - base_email_logger, mock_send_email + redis_email_logger: BaseEmailLogger, + mock_send_email, + alert_redis: tuple[RedisCache, fakeredis.FakeAsyncRedis, fakeredis.FakeServer], ): """Test that budget_alerts does not send duplicate alerts within TTL period""" user_info = CallInfo( @@ -793,12 +737,10 @@ async def test_budget_alerts_soft_budget_duplicate_prevention( event_group=Litellm_EntityType.USER, ) - # Mock the cache so the slot is already claimed (increment returns > 1) - mock_cache = mock.AsyncMock() - mock_cache.async_increment_cache = mock.AsyncMock(return_value=2) - base_email_logger.internal_usage_cache = mock_cache - - await base_email_logger.budget_alerts(type="soft_budget", user_info=user_info) + await alert_redis[0].async_increment( + "email_budget_alerts:soft_budget_crossed:test_user", 1, ttl=EMAIL_BUDGET_ALERT_TTL + ) + await redis_email_logger.budget_alerts(type="soft_budget", user_info=user_info) # Verify email was NOT sent (duplicate prevention) mock_send_email.assert_not_called() @@ -824,7 +766,9 @@ async def test_budget_alerts_no_budgets(base_email_logger, mock_send_email): @pytest.mark.asyncio async def test_budget_alerts_uses_token_for_cache_key( - base_email_logger, mock_send_email + redis_email_logger: BaseEmailLogger, + mock_send_email, + alert_redis: tuple[RedisCache, fakeredis.FakeAsyncRedis, fakeredis.FakeServer], ): """Test that budget_alerts uses token for cache key when available""" user_info = CallInfo( @@ -837,32 +781,19 @@ async def test_budget_alerts_uses_token_for_cache_key( event_group=Litellm_EntityType.KEY, ) - # Mock the cache so the claim is won (increment returns 1) - mock_cache = mock.AsyncMock() - mock_cache.async_increment_cache = mock.AsyncMock(return_value=1) - base_email_logger.internal_usage_cache = mock_cache - with mock.patch.dict( os.environ, { "PROXY_BASE_URL": "http://test.com", }, ): - await base_email_logger.budget_alerts(type="soft_budget", user_info=user_info) + await redis_email_logger.budget_alerts(type="soft_budget", user_info=user_info) - # Verify cache key uses token instead of user_id - mock_cache.async_increment_cache.assert_called_once() - cache_call_args = mock_cache.async_increment_cache.call_args[1] - assert ( - cache_call_args["key"] - == "email_budget_alerts:soft_budget_crossed:hashed_token_123" - ) + assert await alert_redis[1].get("alerts:email_budget_alerts:soft_budget_crossed:hashed_token_123") == b"1" @pytest.mark.asyncio -async def test_get_email_params_soft_budget_crossed( - base_email_logger, mock_lookup_user_email -): +async def test_get_email_params_soft_budget_crossed(base_email_logger, mock_lookup_user_email): """Test that _get_email_params handles soft_budget_crossed event correctly""" with mock.patch.dict( os.environ, @@ -877,16 +808,16 @@ async def test_get_email_params_soft_budget_crossed( ) # Should use default subject template for soft_budget_crossed - assert ( - result.subject == "LiteLLM: Soft Budget Crossed - Total Soft Budget: $100.0" - ) + assert result.subject == "LiteLLM: Soft Budget Crossed - Total Soft Budget: $100.0" assert result.recipient_email == "test@example.com" assert result.base_url == "http://test.com" @pytest.mark.asyncio async def test_budget_alerts_max_budget_alert_crossed( - base_email_logger, mock_send_email + redis_email_logger: BaseEmailLogger, + mock_send_email, + alert_redis: tuple[RedisCache, fakeredis.FakeAsyncRedis, fakeredis.FakeServer], ): """Test that budget_alerts sends email when max budget alert threshold is crossed""" user_info = CallInfo( @@ -897,37 +828,32 @@ async def test_budget_alerts_max_budget_alert_crossed( event_group=Litellm_EntityType.USER, ) - mock_cache = mock.AsyncMock() - mock_cache.async_increment_cache = mock.AsyncMock(return_value=1) - base_email_logger.internal_usage_cache = mock_cache - with mock.patch.dict( os.environ, { "PROXY_BASE_URL": "http://test.com", }, ): - await base_email_logger.budget_alerts( - type="max_budget_alert", user_info=user_info - ) + await redis_email_logger.budget_alerts(type="max_budget_alert", user_info=user_info) mock_send_email.assert_called_once() call_args = mock_send_email.call_args[1] assert call_args["to_email"] == ["test@example.com"] assert "Max Budget Alert" in call_args["subject"] - mock_cache.async_increment_cache.assert_called_once() - cache_call_args = mock_cache.async_increment_cache.call_args[1] + assert await alert_redis[1].get("alerts:email_budget_alerts:max_budget_alert:test_user") == b"1" assert ( - cache_call_args["key"] == "email_budget_alerts:max_budget_alert:test_user" + 0 + < await alert_redis[1].pttl("alerts:email_budget_alerts:max_budget_alert:test_user") + <= EMAIL_BUDGET_ALERT_TTL * 1000 ) - assert cache_call_args["value"] == 1 - assert cache_call_args["ttl"] == EMAIL_BUDGET_ALERT_TTL @pytest.mark.asyncio async def test_multi_threshold_sends_crossed_thresholds( - base_email_logger, mock_send_email + redis_email_logger: BaseEmailLogger, + mock_send_email, + alert_redis: tuple[RedisCache, fakeredis.FakeAsyncRedis, fakeredis.FakeServer], ): """Test that multi-threshold path sends emails for all crossed thresholds""" user_info = CallInfo( @@ -944,29 +870,23 @@ async def test_multi_threshold_sends_crossed_thresholds( }, ) - mock_cache = mock.AsyncMock() - mock_cache.async_increment_cache = mock.AsyncMock(return_value=1) - base_email_logger.internal_usage_cache = mock_cache - with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}): - await base_email_logger.budget_alerts( - type="max_budget_alert", user_info=user_info - ) + await redis_email_logger.budget_alerts(type="max_budget_alert", user_info=user_info) # spend=80 crosses 50% ($50) and 75% ($75), but not 100% ($100) assert mock_send_email.call_count == 2 - # Check cache keys include threshold percentage - cache_keys = [ - c[1]["key"] for c in mock_cache.async_increment_cache.call_args_list + assert sorted(await alert_redis[1].keys("*")) == [ + b"alerts:email_budget_alerts:max_budget_alert:50:hashed_key_1", + b"alerts:email_budget_alerts:max_budget_alert:75:hashed_key_1", ] - assert "email_budget_alerts:max_budget_alert:50:hashed_key_1" in cache_keys - assert "email_budget_alerts:max_budget_alert:75:hashed_key_1" in cache_keys @pytest.mark.asyncio async def test_multi_threshold_dedup_cache_prevents_resend( - base_email_logger, mock_send_email + redis_email_logger: BaseEmailLogger, + mock_send_email, + alert_redis: tuple[RedisCache, fakeredis.FakeAsyncRedis, fakeredis.FakeServer], ): """Test that cached thresholds are not re-sent""" user_info = CallInfo( @@ -982,31 +902,22 @@ async def test_multi_threshold_dedup_cache_prevents_resend( }, ) - # Simulate 50% already claimed (increment returns >1), 75% first send (returns 1) - async def cache_increment(key, value, ttl=None): - if "50:" in key: - return 2 - return 1 - - mock_cache = mock.AsyncMock() - mock_cache.async_increment_cache = mock.AsyncMock(side_effect=cache_increment) - base_email_logger.internal_usage_cache = mock_cache + await alert_redis[0].async_increment( + "email_budget_alerts:max_budget_alert:50:hashed_key_1", 1, ttl=EMAIL_BUDGET_ALERT_TTL + ) with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}): - await base_email_logger.budget_alerts( - type="max_budget_alert", user_info=user_info - ) + await redis_email_logger.budget_alerts(type="max_budget_alert", user_info=user_info) # Only 75% should fire assert mock_send_email.call_count == 1 - cache_key = mock_cache.async_increment_cache.call_args[1]["key"] - assert "75:" in cache_key + assert "75%" in mock_send_email.call_args.kwargs["subject"] + assert await alert_redis[1].get("alerts:email_budget_alerts:max_budget_alert:50:hashed_key_1") == b"2" + assert await alert_redis[1].get("alerts:email_budget_alerts:max_budget_alert:75:hashed_key_1") == b"1" @pytest.mark.asyncio -async def test_multi_threshold_owner_email_auto_included( - base_email_logger, mock_send_email -): +async def test_multi_threshold_owner_email_auto_included(redis_email_logger: BaseEmailLogger, mock_send_email): """Test that the owner email is auto-appended and deduplicated""" user_info = CallInfo( token="hashed_key_1", @@ -1020,14 +931,8 @@ async def test_multi_threshold_owner_email_auto_included( }, ) - mock_cache = mock.AsyncMock() - mock_cache.async_increment_cache = mock.AsyncMock(return_value=1) - base_email_logger.internal_usage_cache = mock_cache - with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}): - await base_email_logger.budget_alerts( - type="max_budget_alert", user_info=user_info - ) + await redis_email_logger.budget_alerts(type="max_budget_alert", user_info=user_info) mock_send_email.assert_called_once() to_emails = mock_send_email.call_args[1]["to_email"] @@ -1036,9 +941,7 @@ async def test_multi_threshold_owner_email_auto_included( @pytest.mark.asyncio -async def test_multi_threshold_malformed_keys_skipped( - base_email_logger, mock_send_email -): +async def test_multi_threshold_malformed_keys_skipped(redis_email_logger: BaseEmailLogger, mock_send_email): """Test that non-numeric threshold keys are skipped""" user_info = CallInfo( token="hashed_key_1", @@ -1053,23 +956,15 @@ async def test_multi_threshold_malformed_keys_skipped( }, ) - mock_cache = mock.AsyncMock() - mock_cache.async_increment_cache = mock.AsyncMock(return_value=1) - base_email_logger.internal_usage_cache = mock_cache - with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}): - await base_email_logger.budget_alerts( - type="max_budget_alert", user_info=user_info - ) + await redis_email_logger.budget_alerts(type="max_budget_alert", user_info=user_info) # Only the valid "50" threshold should fire assert mock_send_email.call_count == 1 @pytest.mark.asyncio -async def test_multi_threshold_empty_emails_only_owner( - base_email_logger, mock_send_email -): +async def test_multi_threshold_empty_emails_only_owner(redis_email_logger: BaseEmailLogger, mock_send_email): """Test that empty email list for a threshold sends only to owner""" user_info = CallInfo( token="hashed_key_1", @@ -1083,14 +978,8 @@ async def test_multi_threshold_empty_emails_only_owner( }, ) - mock_cache = mock.AsyncMock() - mock_cache.async_increment_cache = mock.AsyncMock(return_value=1) - base_email_logger.internal_usage_cache = mock_cache - with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}): - await base_email_logger.budget_alerts( - type="max_budget_alert", user_info=user_info - ) + await redis_email_logger.budget_alerts(type="max_budget_alert", user_info=user_info) mock_send_email.assert_called_once() to_emails = mock_send_email.call_args[1]["to_email"] @@ -1099,7 +988,7 @@ async def test_multi_threshold_empty_emails_only_owner( @pytest.mark.asyncio async def test_multi_threshold_team_member_alert_renders_member_template_per_team( - base_email_logger, mock_send_email + redis_email_logger: BaseEmailLogger, mock_send_email, alert_redis ): """A team member budget alert is keyed per member and team, names the member and team, and goes to the member plus the threshold's configured recipients""" @@ -1114,22 +1003,15 @@ async def test_multi_threshold_team_member_alert_renders_member_template_per_tea max_budget_alert_emails={"50": [], "100": ["finance@co.com"]}, ) - mock_cache = mock.AsyncMock() - mock_cache.async_increment_cache = mock.AsyncMock(return_value=1) - base_email_logger.internal_usage_cache = mock_cache - with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}): - await base_email_logger.budget_alerts(type="max_budget_alert", user_info=user_info) + await redis_email_logger.budget_alerts(type="max_budget_alert", user_info=user_info) - cache_keys = sorted(c[1]["key"] for c in mock_cache.async_increment_cache.call_args_list) - assert cache_keys == [ - "email_budget_alerts:max_budget_alert:100:team_member:member_1:team_a", - "email_budget_alerts:max_budget_alert:50:team_member:member_1:team_a", + assert sorted(await alert_redis[1].keys("*")) == [ + b"alerts:email_budget_alerts:max_budget_alert:100:team_member:member_1:team_a", + b"alerts:email_budget_alerts:max_budget_alert:50:team_member:member_1:team_a", ] assert mock_send_email.call_count == 2 - hundred = next( - c.kwargs for c in mock_send_email.call_args_list if "100%" in c.kwargs["subject"] - ) + hundred = next(c.kwargs for c in mock_send_email.call_args_list if "100%" in c.kwargs["subject"]) assert hundred["subject"] == "LiteLLM: Team Member Budget Alert - 100% of Team Member Budget Reached" assert sorted(hundred["to_email"]) == ["finance@co.com", "member@co.com"] assert "member@co.com" in hundred["html_body"] and "Platform" in hundred["html_body"] @@ -1140,7 +1022,9 @@ async def test_multi_threshold_team_member_alert_renders_member_template_per_tea @pytest.mark.asyncio async def test_no_map_preserves_old_single_threshold( - base_email_logger, mock_send_email + redis_email_logger: BaseEmailLogger, + mock_send_email, + alert_redis: tuple[RedisCache, fakeredis.FakeAsyncRedis, fakeredis.FakeServer], ): """Test that without max_budget_alert_emails, the old 80% single-threshold path works""" user_info = CallInfo( @@ -1151,21 +1035,13 @@ async def test_no_map_preserves_old_single_threshold( event_group=Litellm_EntityType.USER, ) - mock_cache = mock.AsyncMock() - mock_cache.async_increment_cache = mock.AsyncMock(return_value=1) - base_email_logger.internal_usage_cache = mock_cache - with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}): - await base_email_logger.budget_alerts( - type="max_budget_alert", user_info=user_info - ) + await redis_email_logger.budget_alerts(type="max_budget_alert", user_info=user_info) mock_send_email.assert_called_once() call_args = mock_send_email.call_args[1] assert call_args["to_email"] == ["test@example.com"] - # Old path cache key has no threshold percentage - cache_key = mock_cache.async_increment_cache.call_args[1]["key"] - assert cache_key == "email_budget_alerts:max_budget_alert:test_user" + assert await alert_redis[1].get("alerts:email_budget_alerts:max_budget_alert:test_user") == b"1" CUSTOM_SIGNATURE = "
Best,
The Acme Platform Team
" @@ -1186,10 +1062,13 @@ async def test_send_soft_budget_alert_email_uses_custom_signature( max_budget=200.0, soft_budget=100.0, ) - with mock.patch.dict( - os.environ, - {"PROXY_BASE_URL": "http://test.com", "EMAIL_SIGNATURE": CUSTOM_SIGNATURE}, - ), patch("litellm.proxy.proxy_server.premium_user", True): + with ( + mock.patch.dict( + os.environ, + {"PROXY_BASE_URL": "http://test.com", "EMAIL_SIGNATURE": CUSTOM_SIGNATURE}, + ), + patch("litellm.proxy.proxy_server.premium_user", True), + ): await base_email_logger.send_soft_budget_alert_email(event) html_body = mock_send_email.call_args[1]["html_body"] @@ -1213,10 +1092,13 @@ async def test_send_team_soft_budget_alert_email_uses_custom_signature( team_alias="Acme", alert_emails=["teamlead@example.com"], ) - with mock.patch.dict( - os.environ, - {"PROXY_BASE_URL": "http://test.com", "EMAIL_SIGNATURE": CUSTOM_SIGNATURE}, - ), patch("litellm.proxy.proxy_server.premium_user", True): + with ( + mock.patch.dict( + os.environ, + {"PROXY_BASE_URL": "http://test.com", "EMAIL_SIGNATURE": CUSTOM_SIGNATURE}, + ), + patch("litellm.proxy.proxy_server.premium_user", True), + ): await base_email_logger.send_team_soft_budget_alert_email(event) html_body = mock_send_email.call_args[1]["html_body"] @@ -1238,10 +1120,13 @@ async def test_send_max_budget_alert_email_single_recipient_uses_custom_signatur spend=165.0, max_budget=200.0, ) - with mock.patch.dict( - os.environ, - {"PROXY_BASE_URL": "http://test.com", "EMAIL_SIGNATURE": CUSTOM_SIGNATURE}, - ), patch("litellm.proxy.proxy_server.premium_user", True): + with ( + mock.patch.dict( + os.environ, + {"PROXY_BASE_URL": "http://test.com", "EMAIL_SIGNATURE": CUSTOM_SIGNATURE}, + ), + patch("litellm.proxy.proxy_server.premium_user", True), + ): await base_email_logger.send_max_budget_alert_email(event) html_body = mock_send_email.call_args[1]["html_body"] @@ -1263,10 +1148,13 @@ async def test_send_max_budget_alert_email_multi_recipient_uses_custom_signature spend=165.0, max_budget=200.0, ) - with mock.patch.dict( - os.environ, - {"PROXY_BASE_URL": "http://test.com", "EMAIL_SIGNATURE": CUSTOM_SIGNATURE}, - ), patch("litellm.proxy.proxy_server.premium_user", True): + with ( + mock.patch.dict( + os.environ, + {"PROXY_BASE_URL": "http://test.com", "EMAIL_SIGNATURE": CUSTOM_SIGNATURE}, + ), + patch("litellm.proxy.proxy_server.premium_user", True), + ): await base_email_logger.send_max_budget_alert_email( event, threshold_pct=75, recipient_emails=["a@example.com", "b@example.com"] ) @@ -1354,12 +1242,8 @@ async def test_budget_alert_no_duplicate_on_concurrent_crossing( with mock.patch.object(base_email_logger, send_method, side_effect=slow_send): with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}): await asyncio.gather( - base_email_logger.budget_alerts( - type=alert_type, user_info=_budget_alert_user_info(ci_kwargs) - ), - base_email_logger.budget_alerts( - type=alert_type, user_info=_budget_alert_user_info(ci_kwargs) - ), + base_email_logger.budget_alerts(type=alert_type, user_info=_budget_alert_user_info(ci_kwargs)), + base_email_logger.budget_alerts(type=alert_type, user_info=_budget_alert_user_info(ci_kwargs)), ) assert len(sends) == 1 @@ -1388,42 +1272,37 @@ async def test_budget_alert_failed_send_releases_claim_for_retry( with mock.patch.object(base_email_logger, send_method, side_effect=flaky_send): with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}): - await base_email_logger.budget_alerts( - type=alert_type, user_info=_budget_alert_user_info(ci_kwargs) - ) - await base_email_logger.budget_alerts( - type=alert_type, user_info=_budget_alert_user_info(ci_kwargs) - ) + await base_email_logger.budget_alerts(type=alert_type, user_info=_budget_alert_user_info(ci_kwargs)) + await base_email_logger.budget_alerts(type=alert_type, user_info=_budget_alert_user_info(ci_kwargs)) assert len(attempts) == 2 @pytest.mark.asyncio -async def test_budget_alert_release_failure_does_not_propagate(base_email_logger): +async def test_budget_alert_release_failure_does_not_propagate( + redis_email_logger: BaseEmailLogger, alert_redis: tuple[RedisCache, fakeredis.FakeAsyncRedis, fakeredis.FakeServer] +): """If the send fails and releasing the claim also fails (transient cache error), budget_alerts must swallow it and still log the send failure rather than letting the exception escape the fire-and-forget task.""" - mock_cache = mock.AsyncMock() - mock_cache.async_increment_cache = mock.AsyncMock(return_value=1) - mock_cache.async_delete_cache = mock.AsyncMock( - side_effect=RuntimeError("cache backend unavailable") - ) - base_email_logger.internal_usage_cache = mock_cache async def failing_send(*args, **kwargs): + alert_redis[2].connected = False raise ValueError("smtp backend down") - with mock.patch.object( - base_email_logger, "send_max_budget_alert_email", side_effect=failing_send - ): + with mock.patch.object(redis_email_logger, "send_email", side_effect=failing_send): with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}): # Must not raise even though both the send and the release fail. - await base_email_logger.budget_alerts( + await redis_email_logger.budget_alerts( type="max_budget_alert", user_info=_budget_alert_user_info(dict(max_budget=100.0, spend=85.0)), ) - mock_cache.async_delete_cache.assert_awaited_once() + alert_redis[2].connected = True + assert await alert_redis[1].get("alerts:email_budget_alerts:max_budget_alert:hashed_key_1") == b"1" + assert ( + local_get(redis_email_logger.internal_usage_cache, "email_budget_alerts:max_budget_alert:hashed_key_1") is None + ) class _RecordingEmailLogger(BaseEmailLogger): @@ -1443,12 +1322,12 @@ class _UserTable: return self._rows.get(where["user_id"]) -def _prisma_client_with_users(*rows): +def _prisma_client_with_users(*rows: LiteLLM_UserTable) -> SimpleNamespace: table: Final = _UserTable({row.user_id: row for row in rows}) return SimpleNamespace(db=SimpleNamespace(litellm_usertable=table)) -def _key_created_event(user_id): +def _key_created_event(user_id: str) -> SendKeyCreatedEmailEvent: return SendKeyCreatedEmailEvent( user_id=user_id, user_email=None, @@ -1463,9 +1342,7 @@ def _key_created_event(user_id): @pytest.mark.asyncio async def test_key_created_email_goes_to_the_address_stored_for_the_user(monkeypatch): - stored_user: Final = LiteLLM_UserTable( - user_id="user-1", user_email="stored@example.com" - ) + stored_user: Final = LiteLLM_UserTable(user_id="user-1", user_email="stored@example.com") monkeypatch.setattr( "litellm.proxy.proxy_server.prisma_client", _prisma_client_with_users(stored_user), @@ -1481,9 +1358,7 @@ async def test_key_created_email_goes_to_the_address_stored_for_the_user(monkeyp async def test_key_created_email_is_refused_for_a_user_the_database_does_not_know( monkeypatch, ): - monkeypatch.setattr( - "litellm.proxy.proxy_server.prisma_client", _prisma_client_with_users() - ) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _prisma_client_with_users()) email_logger: Final = _RecordingEmailLogger() with pytest.raises(ValueError, match="User email not found for user_id: user-2"): diff --git a/tests/unit/redis/stores/test_email_alert_claims.py b/tests/unit/redis/stores/test_email_alert_claims.py new file mode 100644 index 00000000000..71d80ad93e1 --- /dev/null +++ b/tests/unit/redis/stores/test_email_alert_claims.py @@ -0,0 +1,293 @@ +import asyncio +from collections.abc import AsyncIterator, Awaitable, Callable +from dataclasses import dataclass +from datetime import timedelta +from typing import Final, Literal +from unittest.mock import patch + +import fakeredis +import pytest +import pytest_asyncio +from litellm_enterprise.enterprise_callbacks.send_emails.base_email import BaseEmailLogger +from redis.exceptions import ConnectionError as RedisConnectionError +from redis.exceptions import ResponseError + +from litellm.caching.dual_cache import DualCache +from litellm.caching.redis_cache import RedisCache +from litellm.constants import EMAIL_BUDGET_ALERT_TTL +from litellm.proxy._types import CallInfo, Litellm_EntityType +from litellm.redis.stores.email_alert_claims import SoftBudgetAlert, TransportEmailAlertClaimsStore +from litellm.redis.transport.contract import RedisFailure +from tests._support.fake_redis import async_client, fake_redis_cache, legacy_increment, local_get + + +@dataclass(frozen=True, slots=True) +class SentEmail: + recipients: tuple[str, ...] + subject: str + + +class EmailSink(BaseEmailLogger): + def __init__(self, cache: DualCache, on_send: Callable[[], Awaitable[None]] | None = None) -> None: + super().__init__(internal_usage_cache=cache) # pyright: ignore[reportUnknownMemberType] # legacy logger init is untyped + self.sent: tuple[SentEmail, ...] = () + self.on_send = on_send + + async def send_email(self, from_email: str, to_email: list[str], subject: str, html_body: str) -> None: + self.sent += (SentEmail(tuple(sorted(to_email)), subject),) + if self.on_send is not None: + await self.on_send() + + +@pytest_asyncio.fixture(loop_scope="function") +async def redis_env() -> AsyncIterator[tuple[RedisCache, fakeredis.FakeAsyncRedis, fakeredis.FakeServer]]: + with patch("time.time", return_value=1700000000.0), patch.dict("os.environ", {"PROXY_BASE_URL": "http://test.com"}): + server: Final = fakeredis.FakeServer() + cache: Final = fake_redis_cache(server, namespace="alerts") + raw: Final = fakeredis.FakeAsyncRedis(server=server) + try: + yield cache, raw, server + finally: + await raw.aclose() # pyright: ignore[reportUnknownMemberType, reportAttributeAccessIssue] # fakeredis stubs predate aclose + + +_AlertType = Literal[ + "token_budget", "soft_budget", "max_budget_alert", "user_budget", "team_budget", "organization_budget" +] +_ALERTS: Final[tuple[tuple[_AlertType, CallInfo, str], ...]] = ( + ( + "soft_budget", + CallInfo( + user_id="user:雪", + user_email="owner@example.com", + spend=60, + soft_budget=50, + event_group=Litellm_EntityType.USER, + ), + "email_budget_alerts:soft_budget_crossed:user:雪", + ), + ( + "max_budget_alert", + CallInfo( + token="key:雪", user_email="owner@example.com", spend=85, max_budget=100, event_group=Litellm_EntityType.KEY + ), + "email_budget_alerts:max_budget_alert:key:雪", + ), + ( + "max_budget_alert", + CallInfo( + user_id="member:雪", + team_id="team:a", + user_email="owner@example.com", + spend=85, + max_budget=100, + max_budget_alert_emails={"50": ["finance@example.com"]}, + event_group=Litellm_EntityType.TEAM_MEMBER, + ), + "email_budget_alerts:max_budget_alert:50:team_member:member:雪:team:a", + ), +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("alert_type,info,key", _ALERTS) +async def test_old_and_new_claims_share_numeric_bytes_without_refreshing_expiry( + redis_env: tuple[RedisCache, fakeredis.FakeAsyncRedis, fakeredis.FakeServer], + alert_type: Literal["soft_budget", "max_budget_alert"], + info: CallInfo, + key: str, +) -> None: + cache, raw, _ = redis_env + logger: Final = EmailSink(DualCache(redis_cache=cache)) + assert await legacy_increment(cache, key, 0.25, ttl=37) == 0.25 + await logger.budget_alerts(type=alert_type, user_info=info) + assert logger.sent == () + assert await raw.get(f"alerts:{key}") == b"1.25" + assert await raw.pttl(f"alerts:{key}") == 37000 + await raw.delete(f"alerts:{key}") + await logger.budget_alerts(type=alert_type, user_info=info) + assert len(logger.sent) == 1 + assert await raw.get(f"alerts:{key}") == b"1" + assert await raw.pttl(f"alerts:{key}") == EMAIL_BUDGET_ALERT_TTL * 1000 + assert await legacy_increment(cache, key, 1, ttl=2) == 2 + assert await raw.get(f"alerts:{key}") == b"2" + assert await raw.pttl(f"alerts:{key}") == EMAIL_BUDGET_ALERT_TTL * 1000 + assert await raw.keys("*") == [f"alerts:{key}".encode()] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("alert_type,info,key", _ALERTS) +async def test_concurrent_loggers_claim_once_and_failed_send_deletes_both_tiers( + redis_env: tuple[RedisCache, fakeredis.FakeAsyncRedis, fakeredis.FakeServer], + alert_type: Literal["soft_budget", "max_budget_alert"], + info: CallInfo, + key: str, +) -> None: + cache, raw, _ = redis_env + first: Final = EmailSink(DualCache(redis_cache=cache)) + second: Final = EmailSink(DualCache(redis_cache=cache)) + await asyncio.gather( + first.budget_alerts(type=alert_type, user_info=info), second.budget_alerts(type=alert_type, user_info=info) + ) + assert len(first.sent) + len(second.sent) == 1 + assert await raw.get(f"alerts:{key}") == b"2" + await raw.delete(f"alerts:{key}") + + async def fail_send() -> None: + raise ValueError("email backend failed") + + failed: Final = EmailSink(DualCache(redis_cache=cache), fail_send) + await failed.budget_alerts(type=alert_type, user_info=info) + assert len(failed.sent) == 1 + assert await raw.get(f"alerts:{key}") is None + assert local_get(failed.internal_usage_cache, key) is None + retry: Final = EmailSink(failed.internal_usage_cache) + await retry.budget_alerts(type=alert_type, user_info=info) + assert len(retry.sent) == 1 + assert await raw.get(f"alerts:{key}") == b"1" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "seed,sends,stored", [(b"-1", 1, b"0"), (b"0", 1, b"1"), (b"1", 0, b"2"), (b"broken", 1, b"broken")] +) +async def test_counter_result_and_server_error_keep_existing_send_policy( + redis_env: tuple[RedisCache, fakeredis.FakeAsyncRedis, fakeredis.FakeServer], + seed: bytes, + sends: int, + stored: bytes, +) -> None: + cache, raw, _ = redis_env + logger: Final = EmailSink(DualCache(redis_cache=cache)) + alert_type, info, key = _ALERTS[0] + await raw.set(f"alerts:{key}", seed) + await logger.budget_alerts(type=alert_type, user_info=info) + assert len(logger.sent) == sends + assert await raw.get(f"alerts:{key}") == stored + assert await cache.transport().ping() is True + assert await raw.pttl(f"alerts:{key}") == (-1 if seed == b"broken" else EMAIL_BUDGET_ALERT_TTL * 1000) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("local_cache", [True, False]) +async def test_outage_and_open_breaker_use_local_count_or_none( + redis_env: tuple[RedisCache, fakeredis.FakeAsyncRedis, fakeredis.FakeServer], + local_cache: bool, +) -> None: + cache, raw, server = redis_env + dual: Final = DualCache(redis_cache=cache) + if not local_cache: + dual.in_memory_cache = None # pyright: ignore[reportAttributeAccessIssue] # the logger tolerates a cache without a local layer + logger: Final = EmailSink(dual) + alert_type, info, key = _ALERTS[0] + server.connected = False + for _ in range(6): + await logger.budget_alerts(type=alert_type, user_info=info) + assert len(logger.sent) == (1 if local_cache else 6) + result: Final = await cache.transport().get(key) + assert isinstance(result, RedisFailure) + assert result.kind == "breaker_open" + server.connected = True + assert await raw.get(f"alerts:{key}") is None + + +@pytest.mark.asyncio +async def test_applied_increment_with_lost_reply_falls_back_without_retrying_or_rolling_back( + redis_env: tuple[RedisCache, fakeredis.FakeAsyncRedis, fakeredis.FakeServer], +) -> None: + cache, raw, _ = redis_env + logger: Final = EmailSink(DualCache(redis_cache=cache)) + alert_type, info, key = _ALERTS[0] + client: Final = async_client(cache) + pipeline: Final = client.pipeline(transaction=False) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # fakeredis pipeline is untyped + execute: Final = pipeline.execute + + async def lose_reply(raise_on_error: bool = True) -> list[object]: + await execute(raise_on_error=raise_on_error) + assert local_get(logger.internal_usage_cache, key) == 1 + raise RedisConnectionError("reply lost after write applied") + + with ( + patch.object(client, "pipeline", return_value=pipeline), + patch.object(pipeline, "execute", side_effect=lose_reply), + ): + await logger.budget_alerts(type=alert_type, user_info=info) + assert len(logger.sent) == 1 + assert await raw.get(f"alerts:{key}") == b"1" + assert await raw.pttl(f"alerts:{key}") == -1 + await logger.budget_alerts(type=alert_type, user_info=info) + assert len(logger.sent) == 1 + assert await raw.get(f"alerts:{key}") == b"2" + assert await raw.pttl(f"alerts:{key}") == EMAIL_BUDGET_ALERT_TTL * 1000 + + +@pytest.mark.asyncio +async def test_expiry_failure_leaves_the_increment_and_uses_local_fallback( + redis_env: tuple[RedisCache, fakeredis.FakeAsyncRedis, fakeredis.FakeServer], +) -> None: + cache, raw, _ = redis_env + logger: Final = EmailSink(DualCache(redis_cache=cache)) + alert_type, info, key = _ALERTS[0] + client: Final = async_client(cache) + with patch.object(client, "expire", side_effect=RedisConnectionError("lost before expiry")): + await logger.budget_alerts(type=alert_type, user_info=info) + assert len(logger.sent) == 1 + assert await raw.get(f"alerts:{key}") == b"1" + assert await raw.pttl(f"alerts:{key}") == -1 + + +@pytest.mark.asyncio +async def test_rejected_expiry_uses_local_result_instead_of_applied_redis_count( + redis_env: tuple[RedisCache, fakeredis.FakeAsyncRedis, fakeredis.FakeServer], +) -> None: + cache, raw, _ = redis_env + logger: Final = EmailSink(DualCache(redis_cache=cache)) + alert_type, info, key = _ALERTS[0] + await raw.set(f"alerts:{key}", b"1") + with patch.object(cache.init_async_client(), "expire", side_effect=ResponseError("NOPERM expire denied")): + await logger.budget_alerts(type=alert_type, user_info=info) + assert len(logger.sent) == 1 + assert await raw.get(f"alerts:{key}") == b"2" + assert await raw.pttl(f"alerts:{key}") == -1 + + +@pytest.mark.asyncio +async def test_unknown_applied_release_deletes_locally_and_keeps_retry_policy( + redis_env: tuple[RedisCache, fakeredis.FakeAsyncRedis, fakeredis.FakeServer], +) -> None: + cache, raw, _ = redis_env + + async def fail_send() -> None: + raise ValueError("email backend failed") + + logger: Final = EmailSink(DualCache(redis_cache=cache), fail_send) + alert_type, info, key = _ALERTS[0] + client: Final = async_client(cache) + delete: Final = client.delete + + async def lose_reply(*keys: str) -> int: + await delete(*keys) + raise RedisConnectionError("release reply lost after deletion") + + with patch.object(client, "delete", side_effect=lose_reply): + await logger.budget_alerts(type=alert_type, user_info=info) + assert len(logger.sent) == 1 + assert await raw.get(f"alerts:{key}") is None + assert local_get(logger.internal_usage_cache, key) is None + retry: Final = EmailSink(logger.internal_usage_cache) + await retry.budget_alerts(type=alert_type, user_info=info) + assert len(retry.sent) == 1 + assert await raw.get(f"alerts:{key}") == b"1" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("seconds", [0, -1]) +async def test_nonpositive_expiry_deletes_the_counter_like_legacy_expire( + redis_env: tuple[RedisCache, fakeredis.FakeAsyncRedis, fakeredis.FakeServer], + seconds: int, +) -> None: + cache, raw, _ = redis_env + store: Final = TransportEmailAlertClaimsStore(cache.transport()) + assert await store.claim(SoftBudgetAlert("expired"), ttl=timedelta(seconds=seconds)) == 1 + with patch("time.time", return_value=1700000000.001): + assert await raw.get("alerts:email_budget_alerts:soft_budget_crossed:expired") is None