Merge pull request #23898 from shivamrawat1/litellm_1820_stable_patch

Litellm 1820 stable patch

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))

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]:
         """

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(

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)"
                     )

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:

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

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()

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",
+    ]