diff --git a/litellm/constants.py b/litellm/constants.py index 4c38ecd74b5..26bba4670f0 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1307,6 +1307,9 @@ LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS = int( LITELLM_KEY_ROTATION_GRACE_PERIOD: str = os.getenv( "LITELLM_KEY_ROTATION_GRACE_PERIOD", "" ) # Duration to keep old key valid after rotation (e.g. "24h", "2d"); empty = immediate revoke (default) +LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS = int( + os.getenv("LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS", 600) +) # 10 minutes default — caps the deadlock window if a pod crashes mid-rotation UI_SESSION_TOKEN_TEAM_ID = "litellm-dashboard" LITELLM_PROXY_ADMIN_NAME = "default_user_id" @@ -1330,6 +1333,7 @@ CLOUDZERO_MAX_FETCHED_DATA_RECORDS = int( os.getenv("CLOUDZERO_MAX_FETCHED_DATA_RECORDS", 50000) ) SPEND_LOG_CLEANUP_JOB_NAME = "spend_log_cleanup" +KEY_ROTATION_JOB_NAME = "litellm_key_rotation_job" SPEND_LOG_RUN_LOOPS = int(os.getenv("SPEND_LOG_RUN_LOOPS", 500)) SPEND_LOG_CLEANUP_BATCH_SIZE = int(os.getenv("SPEND_LOG_CLEANUP_BATCH_SIZE", 1000)) SPEND_LOG_QUEUE_SIZE_THRESHOLD = int(os.getenv("SPEND_LOG_QUEUE_SIZE_THRESHOLD", 100)) diff --git a/litellm/proxy/common_utils/key_rotation_manager.py b/litellm/proxy/common_utils/key_rotation_manager.py index 5a0a1fabc7d..aaf39a7a19d 100644 --- a/litellm/proxy/common_utils/key_rotation_manager.py +++ b/litellm/proxy/common_utils/key_rotation_manager.py @@ -11,6 +11,7 @@ from litellm._logging import verbose_proxy_logger from litellm.constants import ( LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME, LITELLM_KEY_ROTATION_GRACE_PERIOD, + LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS, ) from litellm.proxy._types import ( GenerateKeyResponse, @@ -30,14 +31,42 @@ class KeyRotationManager: Manages automated key rotation based on individual key rotation schedules. """ - def __init__(self, prisma_client: PrismaClient): + def __init__(self, prisma_client: PrismaClient, pod_lock_manager=None): self.prisma_client = prisma_client + self.pod_lock_manager = pod_lock_manager async def process_rotations(self): """ - Main entry point - find and rotate keys that are due for rotation + Main entry point - find and rotate keys that are due for rotation. + Uses PodLockManager to ensure only one pod runs rotation in multi-pod deployments. """ + from litellm.constants import KEY_ROTATION_JOB_NAME + + lock_acquired = False try: + # If we have a pod lock manager with Redis, try to acquire the lock + if self.pod_lock_manager and self.pod_lock_manager.redis_cache: + # Use a dedicated lock TTL (default 600s) instead of the check interval + # (which defaults to 86400s / 24h). Using the check interval would create + # a 24-hour deadlock window if a pod crashes before releasing the lock. + lock_ttl = max( + LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS, 300 + ) # At least 5 minutes, configurable via LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS + lock_acquired = ( + await self.pod_lock_manager.acquire_lock( + cronjob_id=KEY_ROTATION_JOB_NAME, + ttl=lock_ttl, + ) + or False + ) + if not lock_acquired: + verbose_proxy_logger.warning( + "Key rotation: another pod is already running rotation " + "or Redis lock acquisition failed — skipping this cycle. " + "Keys will be rotated on the next cycle." + ) + return + verbose_proxy_logger.info("Starting scheduled key rotation check...") # Clean up expired deprecated keys first @@ -74,6 +103,16 @@ class KeyRotationManager: except Exception as e: verbose_proxy_logger.error(f"Key rotation process failed: {e}") + finally: + # Only release the lock if it was actually acquired + if ( + lock_acquired + and self.pod_lock_manager + and self.pod_lock_manager.redis_cache + ): + await self.pod_lock_manager.release_lock( + cronjob_id=KEY_ROTATION_JOB_NAME, + ) async def _find_keys_needing_rotation(self) -> List[LiteLLM_VerificationToken]: """ diff --git a/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py b/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py index 6f86e82cf29..6435498ae03 100644 --- a/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py +++ b/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py @@ -32,22 +32,28 @@ class PodLockManager: async def acquire_lock( self, cronjob_id: str, + ttl: Optional[int] = None, ) -> Optional[bool]: """ Attempt to acquire the lock for a specific cron job using Redis. Uses the SET command with NX and EX options to ensure atomicity. - + Args: cronjob_id: The ID of the cron job to lock + ttl: Optional custom TTL in seconds. Defaults to DEFAULT_CRON_JOB_LOCK_TTL_SECONDS. + Use a longer TTL for jobs that may take longer than the default 60s + (e.g. key rotation with many keys). """ if self.redis_cache is None: verbose_proxy_logger.debug("redis_cache is None, skipping acquire_lock") return None try: + lock_ttl = ttl or DEFAULT_CRON_JOB_LOCK_TTL_SECONDS verbose_proxy_logger.debug( - "Pod %s attempting to acquire Redis lock for cronjob_id=%s", + "Pod %s attempting to acquire Redis lock for cronjob_id=%s (ttl=%ds)", self.pod_id, cronjob_id, + lock_ttl, ) # Try to set the lock key with the pod_id as its value, only if it doesn't exist (NX) # and with an expiration (EX) to avoid deadlocks. @@ -56,7 +62,7 @@ class PodLockManager: lock_key, self.pod_id, nx=True, - ttl=DEFAULT_CRON_JOB_LOCK_TTL_SECONDS, + ttl=lock_ttl, ) if acquired: verbose_proxy_logger.info( @@ -133,11 +139,10 @@ class PodLockManager: ) else: verbose_proxy_logger.warning( - "Spend tracking - pod %s failed to release Redis lock for cronjob_id=%s. " - "Lock will expire after TTL=%ds.", + "Pod %s failed to release Redis lock for cronjob_id=%s. " + "Lock will expire after its TTL.", self.pod_id, cronjob_id, - DEFAULT_CRON_JOB_LOCK_TTL_SECONDS, ) else: verbose_proxy_logger.debug( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index bc2728c2203..3b92daa73d1 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -5955,10 +5955,24 @@ class ProxyStartupEvent: KeyRotationManager, ) - # Get prisma_client from global scope + # Get prisma_client and proxy_logging_obj from global scope global prisma_client + global proxy_logging_obj if prisma_client is not None: - key_rotation_manager = KeyRotationManager(prisma_client) + # Reuse the PodLockManager from db_spend_update_writer + pod_lock_manager = ( + getattr( + getattr(proxy_logging_obj, "db_spend_update_writer", None), + "pod_lock_manager", + None, + ) + if proxy_logging_obj is not None + else None + ) + key_rotation_manager = KeyRotationManager( + prisma_client, + pod_lock_manager=pod_lock_manager, + ) verbose_proxy_logger.debug( f"Key rotation background job scheduled every {LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS} seconds (LITELLM_KEY_ROTATION_ENABLED=true)" ) diff --git a/litellm/router.py b/litellm/router.py index d89a5099b01..9bab9738b30 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -8458,7 +8458,9 @@ class Router: ## get healthy deployments ### get all deployments - healthy_deployments = self._get_all_deployments(model_name=model) + healthy_deployments = self._get_all_deployments( + model_name=model, team_id=request_team_id + ) if len(healthy_deployments) == 0: # check if the user sent in a deployment name instead @@ -8479,7 +8481,9 @@ class Router: ) # Re-assign model to the fallback and try to get deployments again model = fallback_model - healthy_deployments = self._get_all_deployments(model_name=model) + healthy_deployments = self._get_all_deployments( + model_name=model, team_id=request_team_id + ) # If still no deployments after checking for fallbacks, raise an error if len(healthy_deployments) == 0: diff --git a/tests/test_litellm/proxy/common_utils/test_key_rotation_e2e.py b/tests/test_litellm/proxy/common_utils/test_key_rotation_e2e.py new file mode 100644 index 00000000000..f6ef02a86de --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_key_rotation_e2e.py @@ -0,0 +1,559 @@ +""" +End-to-end tests for key rotation feature. + +Covers the critical gaps: +1. Multi-pod simulation: two KeyRotationManagers sharing one PodLockManager +2. Error resilience: partial failures, regenerate_key_fn failures, hook failures +3. Full process_rotations flow with actual key finding + rotation + lock +4. Initialization wiring: PodLockManager is correctly passed +5. Multiple keys: some succeed, some fail, all are attempted +6. Rotation count increments correctly over multiple rotations +""" + +import os +import sys +from datetime import datetime, timedelta, timezone +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.proxy._types import ( + GenerateKeyResponse, + LiteLLM_VerificationToken, +) +from litellm.proxy.common_utils.key_rotation_manager import KeyRotationManager + + +class TestMultiPodKeyRotation: + """ + Simulate two pods sharing one Redis lock to verify only one pod + runs key rotation at a time. + """ + + @pytest.mark.asyncio + async def test_two_pods_only_one_rotates(self): + """ + Two KeyRotationManagers with separate pod_lock_managers but + the same Redis backend. Only the first to acquire the lock + should rotate; the second should skip. + """ + mock_prisma = AsyncMock() + + # Shared state to simulate Redis SET NX behavior + redis_lock = {"holder": None} + + async def make_acquire_lock(pod_id): + async def acquire(cronjob_id, **kwargs): + if redis_lock["holder"] is None: + redis_lock["holder"] = pod_id + return True + return redis_lock["holder"] == pod_id + + return acquire + + async def make_release_lock(pod_id): + async def release(cronjob_id): + if redis_lock["holder"] == pod_id: + redis_lock["holder"] = None + + return release + + # Pod A + pod_a_lock_mgr = MagicMock() + pod_a_lock_mgr.redis_cache = MagicMock() + pod_a_lock_mgr.acquire_lock = AsyncMock( + side_effect=await make_acquire_lock("pod-a") + ) + pod_a_lock_mgr.release_lock = AsyncMock( + side_effect=await make_release_lock("pod-a") + ) + + # Pod B + pod_b_lock_mgr = MagicMock() + pod_b_lock_mgr.redis_cache = MagicMock() + pod_b_lock_mgr.acquire_lock = AsyncMock( + side_effect=await make_acquire_lock("pod-b") + ) + pod_b_lock_mgr.release_lock = AsyncMock( + side_effect=await make_release_lock("pod-b") + ) + + manager_a = KeyRotationManager(mock_prisma, pod_lock_manager=pod_a_lock_mgr) + manager_b = KeyRotationManager(mock_prisma, pod_lock_manager=pod_b_lock_mgr) + + # Both share the same mock methods for rotation logic + for mgr in [manager_a, manager_b]: + mgr._cleanup_expired_deprecated_keys = AsyncMock() + mgr._find_keys_needing_rotation = AsyncMock(return_value=[]) + + # Pod A acquires lock first + await manager_a.process_rotations() + # Pod A should have run rotation + manager_a._cleanup_expired_deprecated_keys.assert_called_once() + manager_a._find_keys_needing_rotation.assert_called_once() + + # Lock is released after pod A finishes, so pod B can now acquire + # But let's simulate pod B trying WHILE pod A holds the lock + # Reset the lock state to simulate concurrent access + redis_lock["holder"] = "pod-a" # Pod A holds the lock + + await manager_b.process_rotations() + # Pod B should NOT have run rotation (lock held by pod-a) + manager_b._cleanup_expired_deprecated_keys.assert_not_called() + manager_b._find_keys_needing_rotation.assert_not_called() + + @pytest.mark.asyncio + async def test_second_pod_runs_after_first_releases(self): + """ + After the first pod releases the lock, the second pod should + be able to acquire and run rotation. + """ + mock_prisma = AsyncMock() + + call_order = [] + + # Pod A - always gets the lock + pod_a_lock = MagicMock() + pod_a_lock.redis_cache = MagicMock() + pod_a_lock.acquire_lock = AsyncMock(return_value=True) + pod_a_lock.release_lock = AsyncMock() + + # Pod B - also gets the lock (simulating after A releases) + pod_b_lock = MagicMock() + pod_b_lock.redis_cache = MagicMock() + pod_b_lock.acquire_lock = AsyncMock(return_value=True) + pod_b_lock.release_lock = AsyncMock() + + manager_a = KeyRotationManager(mock_prisma, pod_lock_manager=pod_a_lock) + manager_b = KeyRotationManager(mock_prisma, pod_lock_manager=pod_b_lock) + + async def cleanup_a(): + call_order.append("a_cleanup") + + async def cleanup_b(): + call_order.append("b_cleanup") + + manager_a._cleanup_expired_deprecated_keys = AsyncMock(side_effect=cleanup_a) + manager_a._find_keys_needing_rotation = AsyncMock(return_value=[]) + manager_b._cleanup_expired_deprecated_keys = AsyncMock(side_effect=cleanup_b) + manager_b._find_keys_needing_rotation = AsyncMock(return_value=[]) + + # Run sequentially: A then B + await manager_a.process_rotations() + await manager_b.process_rotations() + + # Both should have run + assert call_order == ["a_cleanup", "b_cleanup"] + pod_a_lock.release_lock.assert_called_once() + pod_b_lock.release_lock.assert_called_once() + + +class TestKeyRotationErrorResilience: + """ + Tests that key rotation handles errors gracefully: + - regenerate_key_fn failure for one key doesn't block others + - Hook failure doesn't crash the process + - Database update failure is handled + """ + + @pytest.mark.asyncio + async def test_one_key_fails_others_still_rotate(self): + """ + If rotation fails for one key, the remaining keys should still + be attempted. No key should be silently skipped. + """ + mock_prisma = AsyncMock() + manager = KeyRotationManager(mock_prisma) + + key1 = LiteLLM_VerificationToken( + token="token-1", + auto_rotate=True, + rotation_interval="30s", + key_rotation_at=None, + rotation_count=0, + key_name="key-1", + ) + key2 = LiteLLM_VerificationToken( + token="token-2", + auto_rotate=True, + rotation_interval="30s", + key_rotation_at=None, + rotation_count=0, + key_name="key-2", + ) + key3 = LiteLLM_VerificationToken( + token="token-3", + auto_rotate=True, + rotation_interval="30s", + key_rotation_at=None, + rotation_count=0, + key_name="key-3", + ) + + manager._cleanup_expired_deprecated_keys = AsyncMock() + manager._find_keys_needing_rotation = AsyncMock(return_value=[key1, key2, key3]) + + rotate_calls = [] + + async def mock_rotate(key): + rotate_calls.append(key.token) + if key.token == "token-2": + raise Exception("Database connection lost") + + manager._rotate_key = AsyncMock(side_effect=mock_rotate) + + await manager.process_rotations() + + # All 3 keys should have been attempted + assert rotate_calls == ["token-1", "token-2", "token-3"] + + @pytest.mark.asyncio + async def test_regenerate_key_fn_failure_is_caught(self): + """ + If regenerate_key_fn throws, _rotate_key should propagate the error + but process_rotations should catch it per-key. + """ + mock_prisma = AsyncMock() + manager = KeyRotationManager(mock_prisma) + + key = LiteLLM_VerificationToken( + token="test-token", + auto_rotate=True, + rotation_interval="30s", + key_rotation_at=None, + rotation_count=0, + key_name="test-key", + ) + + with patch( + "litellm.proxy.common_utils.key_rotation_manager.regenerate_key_fn", + new_callable=AsyncMock, + side_effect=Exception("regenerate failed: DB timeout"), + ): + # _rotate_key should raise + with pytest.raises(Exception, match="regenerate failed"): + await manager._rotate_key(key) + + # But process_rotations should catch per-key errors + manager._cleanup_expired_deprecated_keys = AsyncMock() + manager._find_keys_needing_rotation = AsyncMock(return_value=[key]) + + with patch( + "litellm.proxy.common_utils.key_rotation_manager.regenerate_key_fn", + new_callable=AsyncMock, + side_effect=Exception("regenerate failed: DB timeout"), + ): + # Should NOT raise - error is caught per-key + await manager.process_rotations() + + @pytest.mark.asyncio + async def test_hook_failure_does_not_prevent_db_update(self): + """ + If the rotation hook (async_key_rotated_hook) fails, the database + update for rotation_count should still have succeeded (it runs before the hook). + """ + mock_prisma = AsyncMock() + manager = KeyRotationManager(mock_prisma) + + key = LiteLLM_VerificationToken( + token="test-token", + auto_rotate=True, + rotation_interval="30s", + key_rotation_at=None, + rotation_count=0, + ) + + mock_response = GenerateKeyResponse( + key="new-key", token_id="new-token-id", user_id="test-user" + ) + + with patch( + "litellm.proxy.common_utils.key_rotation_manager.regenerate_key_fn", + new_callable=AsyncMock, + return_value=mock_response, + ): + with patch( + "litellm.proxy.common_utils.key_rotation_manager.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + side_effect=Exception("Hook failed: secret manager down"), + ): + # This will raise because the hook fails + with pytest.raises(Exception, match="Hook failed"): + await manager._rotate_key(key) + + # The DB update should have been called BEFORE the hook + mock_prisma.db.litellm_verificationtoken.update.assert_called_once() + update_data = mock_prisma.db.litellm_verificationtoken.update.call_args[1][ + "data" + ] + assert update_data["rotation_count"] == 1 + + @pytest.mark.asyncio + async def test_cleanup_failure_does_not_prevent_rotation(self): + """ + If deprecated key cleanup fails, the rotation should still proceed. + """ + mock_prisma = AsyncMock() + mock_pod_lock = MagicMock() + mock_pod_lock.redis_cache = MagicMock() + mock_pod_lock.acquire_lock = AsyncMock(return_value=True) + mock_pod_lock.release_lock = AsyncMock() + + manager = KeyRotationManager(mock_prisma, pod_lock_manager=mock_pod_lock) + + # Cleanup fails + manager._cleanup_expired_deprecated_keys = AsyncMock( + side_effect=Exception("Deprecated table doesn't exist") + ) + + # process_rotations catches the exception internally (try/except), + # but the lock must still be released in the finally block. + await manager.process_rotations() + + # Lock should still be released in finally block + mock_pod_lock.release_lock.assert_called_once() + + +class TestKeyRotationFullFlow: + """ + Full end-to-end flow tests: find keys -> rotate -> update DB -> release lock + """ + + @pytest.mark.asyncio + async def test_full_rotation_flow_with_lock(self): + """ + Full flow: acquire lock -> cleanup -> find keys -> rotate -> update DB -> release lock + """ + mock_prisma = AsyncMock() + + # Setup lock manager + mock_lock = MagicMock() + mock_lock.redis_cache = MagicMock() + mock_lock.acquire_lock = AsyncMock(return_value=True) + mock_lock.release_lock = AsyncMock() + + manager = KeyRotationManager(mock_prisma, pod_lock_manager=mock_lock) + + key = LiteLLM_VerificationToken( + token="old-token-hash", + auto_rotate=True, + rotation_interval="30s", + key_rotation_at=datetime.now(timezone.utc) - timedelta(seconds=60), + rotation_count=2, + key_name="my-key", + key_alias="prod/my-key", + ) + + mock_response = GenerateKeyResponse( + key="sk-new-key-value", + token_id="new-token-hash", + user_id="system", + ) + + # Mock cleanup + mock_prisma.db.litellm_deprecatedverificationtoken.delete_many.return_value = 1 + # Mock find keys + mock_prisma.db.litellm_verificationtoken.find_many.return_value = [key] + + with patch( + "litellm.proxy.common_utils.key_rotation_manager.regenerate_key_fn", + new_callable=AsyncMock, + return_value=mock_response, + ): + with patch( + "litellm.proxy.common_utils.key_rotation_manager.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + ): + await manager.process_rotations() + + # Verify full flow executed: + # 1. Lock acquired + mock_lock.acquire_lock.assert_called_once() + + # 2. Cleanup ran + mock_prisma.db.litellm_deprecatedverificationtoken.delete_many.assert_called_once() + + # 3. Keys were queried + mock_prisma.db.litellm_verificationtoken.find_many.assert_called_once() + + # 4. DB was updated with new rotation info + mock_prisma.db.litellm_verificationtoken.update.assert_called_once() + update_args = mock_prisma.db.litellm_verificationtoken.update.call_args[1] + assert update_args["where"]["token"] == "new-token-hash" + assert update_args["data"]["rotation_count"] == 3 # was 2, now 3 + + # 5. Lock released + mock_lock.release_lock.assert_called_once() + + @pytest.mark.asyncio + async def test_rotation_count_increments_across_multiple_rotations(self): + """ + Simulate 3 consecutive rotations and verify rotation_count increments + correctly each time: 0 -> 1 -> 2 -> 3 + """ + mock_prisma = AsyncMock() + manager = KeyRotationManager(mock_prisma) + + rotation_counts_seen = [] + + for expected_count in range(3): + key = LiteLLM_VerificationToken( + token=f"token-v{expected_count}", + auto_rotate=True, + rotation_interval="30s", + key_rotation_at=None, + rotation_count=expected_count, + ) + + mock_response = GenerateKeyResponse( + key=f"sk-new-v{expected_count + 1}", + token_id=f"token-v{expected_count + 1}", + user_id="system", + ) + + mock_prisma.db.litellm_verificationtoken.update.reset_mock() + + with patch( + "litellm.proxy.common_utils.key_rotation_manager.regenerate_key_fn", + new_callable=AsyncMock, + return_value=mock_response, + ): + with patch( + "litellm.proxy.common_utils.key_rotation_manager.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + ): + await manager._rotate_key(key) + + update_data = mock_prisma.db.litellm_verificationtoken.update.call_args[1][ + "data" + ] + rotation_counts_seen.append(update_data["rotation_count"]) + + assert rotation_counts_seen == [1, 2, 3] + + @pytest.mark.asyncio + async def test_no_keys_to_rotate_skips_gracefully(self): + """ + When no keys need rotation, process should complete without errors. + """ + mock_prisma = AsyncMock() + mock_prisma.db.litellm_deprecatedverificationtoken.delete_many.return_value = 0 + mock_prisma.db.litellm_verificationtoken.find_many.return_value = [] + + mock_lock = MagicMock() + mock_lock.redis_cache = MagicMock() + mock_lock.acquire_lock = AsyncMock(return_value=True) + mock_lock.release_lock = AsyncMock() + + manager = KeyRotationManager(mock_prisma, pod_lock_manager=mock_lock) + + await manager.process_rotations() + + # Verify no rotation was attempted + mock_prisma.db.litellm_verificationtoken.update.assert_not_called() + # But lock was still properly released + mock_lock.release_lock.assert_called_once() + + @pytest.mark.asyncio + async def test_regenerate_response_missing_token_id_skips_db_update(self): + """ + If regenerate_key_fn returns a response without token_id, + the DB update for rotation metadata should be skipped. + """ + mock_prisma = AsyncMock() + manager = KeyRotationManager(mock_prisma) + + key = LiteLLM_VerificationToken( + token="old-token", + auto_rotate=True, + rotation_interval="30s", + key_rotation_at=None, + rotation_count=0, + ) + + # Response with no token_id + mock_response = GenerateKeyResponse( + key="sk-new", + token_id=None, + user_id="system", + ) + + with patch( + "litellm.proxy.common_utils.key_rotation_manager.regenerate_key_fn", + new_callable=AsyncMock, + return_value=mock_response, + ): + with patch( + "litellm.proxy.common_utils.key_rotation_manager.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + ): + await manager._rotate_key(key) + + # DB update should NOT have been called (no token_id) + mock_prisma.db.litellm_verificationtoken.update.assert_not_called() + + +class TestKeyRotationInitialization: + """ + Tests that the PodLockManager wiring in proxy_server.py is correct. + """ + + @pytest.mark.asyncio + async def test_key_rotation_manager_receives_pod_lock_manager(self): + """ + Verify KeyRotationManager stores the pod_lock_manager correctly. + """ + mock_prisma = AsyncMock() + mock_lock = MagicMock() + mock_lock.redis_cache = MagicMock() + + manager = KeyRotationManager(mock_prisma, pod_lock_manager=mock_lock) + + assert manager.pod_lock_manager is mock_lock + assert manager.prisma_client is mock_prisma + + @pytest.mark.asyncio + async def test_key_rotation_manager_default_no_lock(self): + """ + When no pod_lock_manager is provided, it defaults to None. + """ + mock_prisma = AsyncMock() + manager = KeyRotationManager(mock_prisma) + + assert manager.pod_lock_manager is None + + @pytest.mark.asyncio + async def test_lock_pattern_matches_spend_log_cleanup(self): + """ + Verify the key rotation lock pattern is identical to spend_log_cleanup: + - acquire_lock with cronjob_id + - release_lock in finally + - lock_acquired flag guards release + """ + mock_prisma = AsyncMock() + mock_lock = MagicMock() + mock_lock.redis_cache = MagicMock() + mock_lock.acquire_lock = AsyncMock(return_value=True) + mock_lock.release_lock = AsyncMock() + + manager = KeyRotationManager(mock_prisma, pod_lock_manager=mock_lock) + manager._cleanup_expired_deprecated_keys = AsyncMock() + manager._find_keys_needing_rotation = AsyncMock(return_value=[]) + + await manager.process_rotations() + + # Pattern check: acquire with cronjob_id + acquire_call = mock_lock.acquire_lock.call_args + assert "cronjob_id" in acquire_call.kwargs or len(acquire_call.args) > 0 + + # Pattern check: release with same cronjob_id + release_call = mock_lock.release_lock.call_args + assert "cronjob_id" in release_call.kwargs or len(release_call.args) > 0 + + # Both should use the same job name + from litellm.constants import KEY_ROTATION_JOB_NAME + + assert acquire_call.kwargs.get("cronjob_id") == KEY_ROTATION_JOB_NAME + assert release_call.kwargs.get("cronjob_id") == KEY_ROTATION_JOB_NAME diff --git a/tests/test_litellm/proxy/common_utils/test_key_rotation_lock.py b/tests/test_litellm/proxy/common_utils/test_key_rotation_lock.py new file mode 100644 index 00000000000..c0b3611b2b4 --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_key_rotation_lock.py @@ -0,0 +1,229 @@ +""" +Test distributed lock behavior for key rotation manager. + +Verifies that PodLockManager is correctly used to prevent concurrent +key rotation across multiple pods in a distributed deployment. +""" + +import os +import sys +from unittest.mock import AsyncMock, MagicMock + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.proxy._types import LiteLLM_VerificationToken +from litellm.proxy.common_utils.key_rotation_manager import KeyRotationManager + + +class TestKeyRotationLock: + """Test distributed lock behavior in KeyRotationManager.""" + + @pytest.mark.asyncio + async def test_process_rotations_acquires_lock(self): + """ + When PodLockManager is provided and lock is acquired, + rotation logic should run normally. + """ + mock_prisma_client = AsyncMock() + mock_pod_lock_manager = MagicMock() + mock_pod_lock_manager.redis_cache = MagicMock() # Redis is available + mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True) + mock_pod_lock_manager.release_lock = AsyncMock() + + manager = KeyRotationManager( + mock_prisma_client, pod_lock_manager=mock_pod_lock_manager + ) + + # Mock _find_keys_needing_rotation to return empty list (no keys to rotate) + manager._find_keys_needing_rotation = AsyncMock(return_value=[]) + manager._cleanup_expired_deprecated_keys = AsyncMock() + + await manager.process_rotations() + + # Verify lock was acquired with custom TTL + mock_pod_lock_manager.acquire_lock.assert_called_once() + call_kwargs = mock_pod_lock_manager.acquire_lock.call_args + assert call_kwargs.kwargs["cronjob_id"] == "litellm_key_rotation_job" + assert call_kwargs.kwargs["ttl"] >= 300 # At least 5 minutes + + # Verify rotation logic ran (cleanup + find keys called) + manager._cleanup_expired_deprecated_keys.assert_called_once() + manager._find_keys_needing_rotation.assert_called_once() + + # Verify lock was released + mock_pod_lock_manager.release_lock.assert_called_once_with( + cronjob_id="litellm_key_rotation_job", + ) + + @pytest.mark.asyncio + async def test_process_rotations_skips_when_lock_held(self): + """ + When lock is held by another pod, process_rotations() should + return early without performing any rotation. + """ + mock_prisma_client = AsyncMock() + mock_pod_lock_manager = MagicMock() + mock_pod_lock_manager.redis_cache = MagicMock() + mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=False) + mock_pod_lock_manager.release_lock = AsyncMock() + + manager = KeyRotationManager( + mock_prisma_client, pod_lock_manager=mock_pod_lock_manager + ) + + manager._find_keys_needing_rotation = AsyncMock() + manager._cleanup_expired_deprecated_keys = AsyncMock() + + await manager.process_rotations() + + # Verify lock was attempted + mock_pod_lock_manager.acquire_lock.assert_called_once() + + # Verify rotation logic was NOT executed + manager._cleanup_expired_deprecated_keys.assert_not_called() + manager._find_keys_needing_rotation.assert_not_called() + + # Verify lock was NOT released (since it was never acquired) + mock_pod_lock_manager.release_lock.assert_not_called() + + @pytest.mark.asyncio + async def test_process_rotations_releases_lock_on_success(self): + """ + Lock should be released in the finally block after successful rotation. + """ + mock_prisma_client = AsyncMock() + mock_pod_lock_manager = MagicMock() + mock_pod_lock_manager.redis_cache = MagicMock() + mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True) + mock_pod_lock_manager.release_lock = AsyncMock() + + manager = KeyRotationManager( + mock_prisma_client, pod_lock_manager=mock_pod_lock_manager + ) + + # Simulate finding and rotating a key successfully + mock_key = LiteLLM_VerificationToken( + token="test-token", + auto_rotate=True, + rotation_interval="30s", + key_rotation_at=None, + rotation_count=0, + key_name="test-key", + ) + manager._find_keys_needing_rotation = AsyncMock(return_value=[mock_key]) + manager._cleanup_expired_deprecated_keys = AsyncMock() + manager._rotate_key = AsyncMock() + + await manager.process_rotations() + + # Verify rotation was performed + manager._rotate_key.assert_called_once_with(mock_key) + + # Verify lock was released after success + mock_pod_lock_manager.release_lock.assert_called_once_with( + cronjob_id="litellm_key_rotation_job", + ) + + @pytest.mark.asyncio + async def test_process_rotations_releases_lock_on_error(self): + """ + Lock should be released in the finally block even if rotation + throws an exception. + """ + mock_prisma_client = AsyncMock() + mock_pod_lock_manager = MagicMock() + mock_pod_lock_manager.redis_cache = MagicMock() + mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True) + mock_pod_lock_manager.release_lock = AsyncMock() + + manager = KeyRotationManager( + mock_prisma_client, pod_lock_manager=mock_pod_lock_manager + ) + + # Simulate an error during cleanup + manager._cleanup_expired_deprecated_keys = AsyncMock( + side_effect=Exception("Database connection failed") + ) + + await manager.process_rotations() + + # Verify lock was still released despite the error + mock_pod_lock_manager.release_lock.assert_called_once_with( + cronjob_id="litellm_key_rotation_job", + ) + + @pytest.mark.asyncio + async def test_process_rotations_works_without_lock_manager(self): + """ + When pod_lock_manager=None, rotation should run normally + without any lock logic (backward compat / single-pod mode). + """ + mock_prisma_client = AsyncMock() + + # No pod_lock_manager provided (default None) + manager = KeyRotationManager(mock_prisma_client) + + manager._find_keys_needing_rotation = AsyncMock(return_value=[]) + manager._cleanup_expired_deprecated_keys = AsyncMock() + + await manager.process_rotations() + + # Verify rotation logic ran normally + manager._cleanup_expired_deprecated_keys.assert_called_once() + manager._find_keys_needing_rotation.assert_called_once() + + @pytest.mark.asyncio + async def test_process_rotations_works_without_redis_cache(self): + """ + When pod_lock_manager exists but redis_cache is None (no Redis configured), + rotation should run normally without locking. + """ + mock_prisma_client = AsyncMock() + mock_pod_lock_manager = MagicMock() + mock_pod_lock_manager.redis_cache = None # No Redis available + + manager = KeyRotationManager( + mock_prisma_client, pod_lock_manager=mock_pod_lock_manager + ) + + manager._find_keys_needing_rotation = AsyncMock(return_value=[]) + manager._cleanup_expired_deprecated_keys = AsyncMock() + + await manager.process_rotations() + + # Verify lock was NOT attempted (no Redis) + mock_pod_lock_manager.acquire_lock.assert_not_called() + + # Verify rotation logic still ran + manager._cleanup_expired_deprecated_keys.assert_called_once() + manager._find_keys_needing_rotation.assert_called_once() + + @pytest.mark.asyncio + async def test_process_rotations_handles_none_lock_result(self): + """ + When acquire_lock returns None (edge case), it should be treated + as lock NOT acquired, and rotation should be skipped. + """ + mock_prisma_client = AsyncMock() + mock_pod_lock_manager = MagicMock() + mock_pod_lock_manager.redis_cache = MagicMock() + mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=None) + mock_pod_lock_manager.release_lock = AsyncMock() + + manager = KeyRotationManager( + mock_prisma_client, pod_lock_manager=mock_pod_lock_manager + ) + + manager._find_keys_needing_rotation = AsyncMock() + manager._cleanup_expired_deprecated_keys = AsyncMock() + + await manager.process_rotations() + + # Verify rotation logic was NOT executed (None treated as False via `or False`) + manager._cleanup_expired_deprecated_keys.assert_not_called() + manager._find_keys_needing_rotation.assert_not_called() + + # Verify lock was NOT released (lock_acquired is False) + mock_pod_lock_manager.release_lock.assert_not_called() diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 4bb9685f5e6..8b151944779 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -5,7 +5,6 @@ import sys from unittest.mock import AsyncMock, MagicMock, patch import pytest -from fastapi.testclient import TestClient sys.path.insert( 0, os.path.abspath("../../..") @@ -13,7 +12,6 @@ sys.path.insert( import litellm -from litellm.router_utils.fallback_event_handlers import run_async_fallback def test_update_kwargs_does_not_mutate_defaults_and_merges_metadata(): @@ -127,7 +125,7 @@ async def test_async_router_acreate_file(): """ Write to all deployments of a model """ - from unittest.mock import MagicMock, call, patch + from unittest.mock import MagicMock, patch router = litellm.Router( model_list=[ @@ -734,7 +732,7 @@ async def test_router_ageneric_api_call_with_fallbacks_helper(): """ Test the _ageneric_api_call_with_fallbacks_helper method with various scenarios """ - from unittest.mock import AsyncMock, MagicMock, patch + from unittest.mock import patch router = litellm.Router( model_list=[ @@ -1121,10 +1119,9 @@ def test_get_model_access_groups_cache_invalidation_upsert_deployment(): @pytest.mark.asyncio async def test_acompletion_streaming_iterator(): """Test _acompletion_streaming_iterator for normal streaming and fallback behavior.""" - from unittest.mock import AsyncMock, MagicMock + from unittest.mock import MagicMock from litellm.exceptions import MidStreamFallbackError - from litellm.types.utils import ModelResponseStream # Helper class for creating async iterators class AsyncIterator: @@ -1967,7 +1964,10 @@ def test_get_deployment_credentials_with_provider_aws_bedrock_runtime_endpoint() ) 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" @@ -1989,11 +1989,11 @@ def test_get_deployment_credentials_with_provider_resolves_credential_name(): credential_values={ "api_key": "resolved-api-key", "api_base": "https://resolved.openai.azure.com", - "api_version": "2024-02-01" - } + "api_version": "2024-02-01", + }, ) ] - + router = litellm.Router( model_list=[ { @@ -2017,7 +2017,7 @@ def test_get_deployment_credentials_with_provider_resolves_credential_name(): assert credentials["custom_llm_provider"] == "azure" # Ensure credential name is removed after resolution assert "litellm_credential_name" not in credentials - + # Cleanup litellm.credential_list = [] @@ -2122,7 +2122,10 @@ async def test_aguardrail_helper(): # Mock the original function async def mock_original_function(**kwargs): - return {"result": "success", "selected_guardrail": kwargs.get("selected_guardrail")} + return { + "result": "success", + "selected_guardrail": kwargs.get("selected_guardrail"), + } result = await router._aguardrail_helper( model="content-filter", @@ -2156,7 +2159,10 @@ async def test_aguardrail(): # Mock the original function async def mock_original_function(**kwargs): - return {"result": "success", "selected_guardrail": kwargs.get("selected_guardrail")} + return { + "result": "success", + "selected_guardrail": kwargs.get("selected_guardrail"), + } result = await router.aguardrail( guardrail_name="content-filter", @@ -2166,6 +2172,7 @@ async def test_aguardrail(): assert result["result"] == "success" assert result["selected_guardrail"]["id"] == "guardrail-1" + @pytest.mark.asyncio async def test_anthropic_messages_call_type_is_cached(): """ @@ -2237,36 +2244,33 @@ async def test_anthropic_messages_call_type_is_cached(): additional_headers=None, ), ) - + cache = DualCache() deployment_check = PromptCachingDeploymentCheck(cache=cache) prompt_cache = PromptCachingCache(cache=cache) - + # Create messages with enough tokens to pass the caching threshold test_messages = [ { - "role": "user", + "role": "user", "content": [ { - "type": "text", + "type": "text", "text": "test long message here" * 1024, - "cache_control": { - "type": "ephemeral", - "ttl": "5m" - } + "cache_control": {"type": "ephemeral", "ttl": "5m"}, } - ] + ], } ] test_model_id = "test-model-id-123" - + # Create a payload with anthropic_messages call type payload = create_standard_logging_payload() payload["call_type"] = CallTypes.anthropic_messages.value payload["messages"] = test_messages payload["model"] = "anthropic/claude-3-5-sonnet-20240620" payload["model_id"] = test_model_id - + # Log the success event (should cache the model_id) await deployment_check.async_log_success_event( kwargs={"standard_logging_object": payload}, @@ -2274,19 +2278,23 @@ async def test_anthropic_messages_call_type_is_cached(): start_time=1234567890.0, end_time=1234567891.0, ) - + # Small delay to ensure cache write completes await asyncio.sleep(0.1) - + # Verify that the model_id was actually cached cached_result = await prompt_cache.async_get_model_id( messages=test_messages, tools=None, ) - + # 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(): @@ -2502,9 +2510,7 @@ def test_credential_name_injected_as_tag(): ) kwargs: dict = {"metadata": {"tags": ["A.101"]}} - deployment = router.get_deployment_by_model_group_name( - model_group_name="xai-model" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="xai-model") router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) assert "Credential: xAI" in kwargs["metadata"]["tags"] @@ -2529,9 +2535,7 @@ def test_credential_name_not_duplicated_in_tags(): ) kwargs: dict = {"metadata": {"tags": ["Credential: xAI", "A.101"]}} - deployment = router.get_deployment_by_model_group_name( - model_group_name="xai-model" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="xai-model") router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) assert kwargs["metadata"]["tags"].count("Credential: xAI") == 1 @@ -2553,9 +2557,311 @@ def test_credential_name_not_injected_when_absent(): ) kwargs: dict = {"metadata": {"tags": ["A.101"]}} - deployment = router.get_deployment_by_model_group_name( - model_group_name="gpt-model" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-model") router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) assert kwargs["metadata"]["tags"] == ["A.101"] + + +def test_combine_fallback_usage(): + """Test that _combine_fallback_usage merges partial and fallback usage.""" + from litellm.router import Router + from litellm.types.utils import Usage + + # Create a stream chunk with usage + chunk = litellm.ModelResponseStream( + id="test", + model="gpt-4o", + choices=[], + usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), + ) + + # Call _combine_fallback_usage with no extra usage + Router._combine_fallback_usage(chunk, None) + assert chunk.usage is not None + assert chunk.usage.prompt_tokens == 10 + assert chunk.usage.completion_tokens == 5 + assert chunk.usage.total_tokens == 15 + + +@pytest.mark.asyncio +async def test_team_scoped_model_fallback(): + """ + Test that fallback works correctly for team-scoped models. + + When a team-scoped model fails and the fallback model is also team-scoped, + the router should find the fallback deployment by matching team_public_model_name. + """ + router = litellm.Router( + model_list=[ + { + "model_name": "team-a-primary-internal", + "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "fake"}, + "model_info": { + "team_id": "team-a", + "team_public_model_name": "primary-model", + }, + }, + { + "model_name": "team-a-fallback-internal", + "litellm_params": { + "model": "gpt-4", + "api_key": "fake", + "mock_response": "fallback success from team-a", + }, + "model_info": { + "team_id": "team-a", + "team_public_model_name": "fallback-model", + }, + }, + ], + fallbacks=[{"primary-model": ["fallback-model"]}], + ) + + response = await router.acompletion( + model="primary-model", + messages=[{"role": "user", "content": "Hello"}], + metadata={"user_api_key_team_id": "team-a"}, + mock_testing_fallbacks=True, + ) + assert response is not None + assert response.choices[0].message.content == "fallback success from team-a" + + +@pytest.mark.asyncio +async def test_team_scoped_model_fallback_to_global(): + """ + Test that a team-scoped model can fall back to a global (non-team) model. + + Global models (no team_id on deployment) should be accessible as fallback + targets for team-scoped requests. + """ + router = litellm.Router( + model_list=[ + { + "model_name": "team-a-primary-internal", + "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "fake"}, + "model_info": { + "team_id": "team-a", + "team_public_model_name": "primary-model", + }, + }, + { + "model_name": "global-fallback", + "litellm_params": { + "model": "gpt-4", + "api_key": "fake", + "mock_response": "global fallback success", + }, + }, + ], + fallbacks=[{"primary-model": ["global-fallback"]}], + ) + + response = await router.acompletion( + model="primary-model", + messages=[{"role": "user", "content": "Hello"}], + metadata={"user_api_key_team_id": "team-a"}, + mock_testing_fallbacks=True, + ) + assert response is not None + assert response.choices[0].message.content == "global fallback success" + + +@pytest.mark.asyncio +async def test_team_scoped_model_fallback_cross_team_blocked(): + """ + Test that cross-team fallback is correctly blocked. + + When team-a's model fails and the fallback target is scoped to team-b, + the router should NOT use it (team isolation). + """ + router = litellm.Router( + model_list=[ + { + "model_name": "team-a-primary-internal", + "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "fake"}, + "model_info": { + "team_id": "team-a", + "team_public_model_name": "primary-model", + }, + }, + { + "model_name": "team-b-fallback-internal", + "litellm_params": { + "model": "gpt-4", + "api_key": "fake", + "mock_response": "team-b response - should not reach here", + }, + "model_info": { + "team_id": "team-b", + "team_public_model_name": "fallback-model", + }, + }, + ], + fallbacks=[{"primary-model": ["fallback-model"]}], + ) + + with pytest.raises(Exception): + await router.acompletion( + model="primary-model", + messages=[{"role": "user", "content": "Hello"}], + metadata={"user_api_key_team_id": "team-a"}, + mock_testing_fallbacks=True, + ) + + +def test_get_all_deployments_with_team_id(): + """ + Test that _get_all_deployments with team_id can find deployments + by team_public_model_name when the model_name is not in the index. + """ + router = litellm.Router( + model_list=[ + { + "model_name": "internal-team-deployment", + "litellm_params": {"model": "gpt-4", "api_key": "fake"}, + "model_info": { + "team_id": "team-x", + "team_public_model_name": "gpt-4", + }, + }, + ], + ) + + # Without team_id: "gpt-4" is not in the model_name index (internal name is different) + deployments = router._get_all_deployments(model_name="gpt-4") + assert len(deployments) == 0 + + # With correct team_id: should find via O(n) scan matching team_public_model_name + deployments = router._get_all_deployments(model_name="gpt-4", team_id="team-x") + assert len(deployments) == 1 + assert deployments[0]["model_name"] == "internal-team-deployment" + + # With wrong team_id: should find nothing + deployments = router._get_all_deployments(model_name="gpt-4", team_id="team-y") + assert len(deployments) == 0 + + +def test_multiregion_team_deployments_unique_model_names(): + """ + Simulates athenahealth's exact setup: unique model_names per deployment, + same team_public_model_name, multiple regions. + + Verifies that _get_all_deployments returns ALL regional deployments + for a team when queried by team_public_model_name. + """ + router = litellm.Router( + model_list=[ + { + "model_name": "metis-claude-us-east-1", + "litellm_params": { + "model": "bedrock/anthropic.claude-3-sonnet", + "aws_region_name": "us-east-1", + "api_key": "fake", + }, + "model_info": { + "team_id": "metis-team", + "team_public_model_name": "claude-sonnet", + }, + }, + { + "model_name": "metis-claude-us-west-2", + "litellm_params": { + "model": "bedrock/anthropic.claude-3-sonnet", + "aws_region_name": "us-west-2", + "api_key": "fake", + }, + "model_info": { + "team_id": "metis-team", + "team_public_model_name": "claude-sonnet", + }, + }, + ], + ) + + # "claude-sonnet" is NOT in the model_name index + assert "claude-sonnet" not in router.model_names + + # Without team_id: returns nothing (no model_name="claude-sonnet" in index, no O(n) scan) + deployments = router._get_all_deployments(model_name="claude-sonnet") + 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" + ) + 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" + + # Wrong team: returns nothing + deployments = router._get_all_deployments( + model_name="claude-sonnet", team_id="other-team" + ) + assert len(deployments) == 0 + + +@pytest.mark.asyncio +async def test_multiregion_team_failover_between_regions(): + """ + Simulates athenahealth's multiregion failover scenario: + - Two Bedrock deployments (us-east-1 and us-west-2) with unique model_names + - Same team_public_model_name ("claude-sonnet") + - Primary region fails → router should failover to second region + + This is the exact scenario Sean Glover from athenahealth will demonstrate. + """ + router = litellm.Router( + model_list=[ + { + "model_name": "metis-claude-us-east-1", + "litellm_params": { + "model": "bedrock/anthropic.claude-3-sonnet", + "api_key": "fake", + "mock_response": "response from us-east-1", + }, + "model_info": { + "team_id": "metis-team", + "team_public_model_name": "claude-sonnet", + }, + }, + { + "model_name": "metis-claude-us-west-2", + "litellm_params": { + "model": "bedrock/anthropic.claude-3-sonnet", + "api_key": "fake", + "mock_response": "response from us-west-2", + }, + "model_info": { + "team_id": "metis-team", + "team_public_model_name": "claude-sonnet", + }, + }, + ], + num_retries=1, + ) + + # 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" + ) + + # Make a normal request — should succeed from one of the regions + response = await router.acompletion( + model="claude-sonnet", + messages=[{"role": "user", "content": "Hello"}], + metadata={"user_api_key_team_id": "metis-team"}, + ) + assert response is not None + assert response.choices[0].message.content in [ + "response from us-east-1", + "response from us-west-2", + ]