fix: re-add scheduled rotations

.circleci/config.yml

@@ -148,7 +148,8 @@ jobs:
             python -m pip install types-requests types-setuptools types-redis types-PyYAML
             if ! python -m mypy . \
                 --config-file mypy.ini \
-                --ignore-missing-imports; then
+                --ignore-missing-imports \
+                --no-incremental; then
                 echo "mypy detected errors"
                 exit 1
             fi

docs/my-website/docs/proxy/config_settings.md

@@ -614,6 +614,8 @@ router_settings:
 | LITELLM_MIGRATION_DIR | Custom migrations directory for prisma migrations, used for baselining db in read-only file systems.
 | LITELLM_HOSTED_UI | URL of the hosted UI for LiteLLM
 | LITELM_ENVIRONMENT | Environment of LiteLLM Instance, used by logging services. Currently only used by DeepEval.
+| LITELLM_KEY_ROTATION_ENABLED | Enable auto-key rotation for LiteLLM (boolean). Default is false.
+| LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS | Interval in seconds for how often to run job that auto-rotates keys. Default is 86400 (24 hours).
 | LITELLM_LICENSE | License key for LiteLLM usage
 | LITELLM_LOCAL_MODEL_COST_MAP | Local configuration for model cost mapping in LiteLLM
 | LITELLM_LOG | Enable detailed logging for LiteLLM

litellm-proxy-extras/litellm_proxy_extras/schema.prisma

@@ -221,6 +221,10 @@ model LiteLLM_VerificationToken {
     created_by String?
     updated_at      DateTime?               @default(now()) @updatedAt @map("updated_at")
     updated_by String?
+    rotation_count   Int?       @default(0) // Number of times key has been rotated
+    auto_rotate      Boolean?   @default(false) // Whether this key should be auto-rotated
+    rotation_interval String?   // How often to rotate (e.g., "30d", "90d")
+    last_rotation_at DateTime?  // When this key was last rotated
     litellm_budget_table LiteLLM_BudgetTable?   @relation(fields: [budget_id], references: [budget_id])
     litellm_organization_table LiteLLM_OrganizationTable?   @relation(fields: [organization_id], references: [organization_id])
     object_permission LiteLLM_ObjectPermissionTable?   @relation(fields: [object_permission_id], references: [object_permission_id])

litellm-proxy-extras/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "litellm-proxy-extras"
-version = "0.2.19"
+version = "0.2.20"
 description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
 authors = ["BerriAI"]
 readme = "README.md"
@@ -22,7 +22,7 @@ requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"
 
 [tool.commitizen]
-version = "0.2.19"
+version = "0.2.20"
 version_files = [
     "pyproject.toml:version",
     "../requirements.txt:litellm-proxy-extras==",

litellm/constants.py

@@ -989,7 +989,11 @@ HEALTH_CHECK_TIMEOUT_SECONDS = int(
 )  # 60 seconds
 LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME = "litellm-internal-health-check"
 LITTELM_CLI_SERVICE_ACCOUNT_NAME = "litellm-cli"
+LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME = "litellm_internal_jobs"
 
+# Key Rotation Constants
+LITELLM_KEY_ROTATION_ENABLED = os.getenv("LITELLM_KEY_ROTATION_ENABLED", "false")
+LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS = int(os.getenv("LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS", 86400))  # 24 hours default
 UI_SESSION_TOKEN_TEAM_ID = "litellm-dashboard"
 LITELLM_PROXY_ADMIN_NAME = "default_user_id"
 

litellm/model_prices_and_context_window_backup.json

@@ -16660,8 +16660,8 @@
         "output_cost_per_token_above_200k_tokens": 2.25e-05,
         "litellm_provider": "openrouter",
         "max_input_tokens": 1000000,
-        "max_output_tokens": 1000000,
-        "max_tokens": 1000000,
+        "max_output_tokens": 64000,
+        "max_tokens": 64000,
         "mode": "chat",
         "output_cost_per_token": 1.5e-05,
         "supports_assistant_prefill": true,

litellm/proxy/_types.py

@@ -1,6 +1,5 @@
 import enum
 import json
-from litellm._uuid import uuid
 from datetime import datetime
 from typing import (
     TYPE_CHECKING,
@@ -24,6 +23,7 @@ from pydantic import (
 )
 from typing_extensions import Required, TypedDict
 
+from litellm._uuid import uuid
 from litellm.types.integrations.slack_alerting import AlertType
 from litellm.types.llms.openai import AllMessageValues, OpenAIFileObject
 from litellm.types.mcp import (
@@ -787,6 +787,14 @@ class GenerateKeyRequest(KeyRequestBase):
         default=LiteLLMKeyType.DEFAULT,
         description="Type of key that determines default allowed routes.",
     )
+    auto_rotate: Optional[bool] = Field(
+        default=False,
+        description="Whether this key should be automatically rotated"
+    )
+    rotation_interval: Optional[str] = Field(
+        default=None,
+        description="How often to rotate this key (e.g., '30d', '90d'). Required if auto_rotate=True"
+    )
 
 
 class GenerateKeyResponse(KeyRequestBase):
@@ -1802,6 +1810,10 @@ class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase):
     updated_by: Optional[str] = None
     object_permission_id: Optional[str] = None
     object_permission: Optional[LiteLLM_ObjectPermissionTable] = None
+    rotation_count: Optional[int] = 0  # Number of times key has been rotated
+    auto_rotate: Optional[bool] = False  # Whether this key should be auto-rotated
+    rotation_interval: Optional[str] = None  # How often to rotate (e.g., "30d", "90d")
+    last_rotation_at: Optional[datetime] = None  # When this key was last rotated
 
     model_config = ConfigDict(protected_namespaces=())
 
@@ -1932,6 +1944,23 @@ class UserAPIKeyAuth(
             key_alias=LITTELM_CLI_SERVICE_ACCOUNT_NAME,
             team_alias=LITTELM_CLI_SERVICE_ACCOUNT_NAME,
         )
+    
+    @classmethod
+    def get_litellm_internal_jobs_user_api_key_auth(cls) -> "UserAPIKeyAuth":
+        """
+        Returns a `UserAPIKeyAuth` object for internal LiteLLM jobs like key rotation.
+
+        This is used to track actions performed by automated system jobs.
+        """
+        from litellm.constants import LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME
+
+        return cls(
+            api_key=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME,
+            team_id="system",
+            key_alias=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME,
+            team_alias="system",
+            user_id="system",
+        )
 
 
 class UserInfoResponse(LiteLLMPydanticObjectBase):

litellm/proxy/common_utils/key_rotation_manager.py

@@ -0,0 +1,146 @@
+"""
+Key Rotation Manager - Automated key rotation based on rotation schedules
+
+Handles finding keys that need rotation based on their individual schedules.
+"""
+
+from datetime import datetime, timedelta, timezone
+from typing import List
+
+from litellm._logging import verbose_proxy_logger
+from litellm.constants import LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME
+from litellm.proxy._types import (
+    GenerateKeyResponse,
+    LiteLLM_VerificationToken,
+    RegenerateKeyRequest,
+)
+from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks
+from litellm.proxy.management_endpoints.key_management_endpoints import (
+    regenerate_key_fn,
+)
+from litellm.proxy.utils import PrismaClient
+
+
+class KeyRotationManager:
+    """
+    Manages automated key rotation based on individual key rotation schedules.
+    """
+    
+    def __init__(self, prisma_client: PrismaClient):
+        self.prisma_client = prisma_client
+        
+    async def process_rotations(self):
+        """
+        Main entry point - find and rotate keys that are due for rotation
+        """
+        try:
+            verbose_proxy_logger.info("Starting scheduled key rotation check...")
+            
+            # Find keys that are due for rotation
+            keys_to_rotate = await self._find_keys_needing_rotation()
+            
+            if not keys_to_rotate:
+                verbose_proxy_logger.debug("No keys are due for rotation at this time")
+                return
+                
+            verbose_proxy_logger.info(f"Found {len(keys_to_rotate)} keys due for rotation")
+            
+            # Rotate each key
+            for key in keys_to_rotate:
+                try:
+                    await self._rotate_key(key)
+                    key_identifier = key.key_name or (key.token[:8] + "..." if key.token else "unknown")
+                    verbose_proxy_logger.info(f"Successfully rotated key: {key_identifier}")
+                except Exception as e:
+                    key_identifier = key.key_name or (key.token[:8] + "..." if key.token else "unknown")
+                    verbose_proxy_logger.error(f"Failed to rotate key {key_identifier}: {e}")
+                    
+        except Exception as e:
+            verbose_proxy_logger.error(f"Key rotation process failed: {e}")
+    
+    async def _find_keys_needing_rotation(self) -> List[LiteLLM_VerificationToken]:
+        """
+        Find keys that are due for rotation based on their rotation interval.
+        
+        Logic:
+        - Key has auto_rotate = true
+        - Key has rotation_interval set
+        - Either: never been rotated (last_rotation_at is null) OR
+        - Time since last rotation >= rotation_interval
+        """
+        keys_with_rotation = await self.prisma_client.db.litellm_verificationtoken.find_many(
+            where={
+                "auto_rotate": True,  # Only keys marked for auto rotation
+                "rotation_interval": {"not": None}  # Must have rotation interval set
+            }
+        )
+        
+        # Filter keys that need rotation based on last_rotation_at + interval
+        keys_needing_rotation = []
+        now = datetime.now(timezone.utc)
+        
+        for key in keys_with_rotation:
+            if self._should_rotate_key(key, now):
+                keys_needing_rotation.append(key)
+        
+        return keys_needing_rotation
+    
+    def _should_rotate_key(self, key: LiteLLM_VerificationToken, now: datetime) -> bool:
+        """
+        Determine if a key should be rotated based on last rotation time and interval.
+        """
+        if not key.rotation_interval:
+            return False
+        
+        # If never rotated, rotate immediately
+        if key.last_rotation_at is None:
+            return True
+        
+        # Calculate if enough time has passed since last rotation
+        from litellm.litellm_core_utils.duration_parser import duration_in_seconds
+        
+        interval_seconds = duration_in_seconds(key.rotation_interval)
+        next_rotation_time = key.last_rotation_at + timedelta(seconds=interval_seconds)
+        
+        return now >= next_rotation_time
+    
+    async def _rotate_key(self, key: LiteLLM_VerificationToken):
+        """
+        Rotate a single key using existing regenerate_key_fn and call the rotation hook
+        """
+        # Create regenerate request 
+        regenerate_request = RegenerateKeyRequest(
+            key=key.token or ""
+        )
+        
+        # Create a system user for key rotation
+        from litellm.proxy._types import UserAPIKeyAuth
+        system_user = UserAPIKeyAuth.get_litellm_internal_jobs_user_api_key_auth()
+        
+        # Use existing regenerate key function
+        response = await regenerate_key_fn(
+            data=regenerate_request,
+            user_api_key_dict=system_user,
+            litellm_changed_by=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME
+        )
+       
+        # Update the NEW key with rotation info (regenerate_key_fn creates a new token)
+        if isinstance(response, GenerateKeyResponse) and response.token_id:
+            await self.prisma_client.db.litellm_verificationtoken.update(
+                where={"token": response.token_id},
+                data={
+                    "rotation_count": (key.rotation_count or 0) + 1,
+                    "last_rotation_at": datetime.now(timezone.utc)
+                }
+            )
+        
+        # Call the existing rotation hook for notifications, audit logs, etc.
+        if isinstance(response, GenerateKeyResponse):
+            await KeyManagementEventHooks.async_key_rotated_hook(
+                data=regenerate_request,
+                existing_key_row=key,
+                response=response,
+                user_api_key_dict=system_user,
+                litellm_changed_by=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME
+            )
+    

litellm/proxy/management_endpoints/key_management_endpoints.py

@@ -14,7 +14,6 @@ import copy
 import json
 import secrets
 import traceback
-from litellm._uuid import uuid
 from datetime import datetime, timedelta, timezone
 from typing import List, Literal, Optional, Tuple, cast
 
@@ -23,6 +22,7 @@ from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request, s
 
 import litellm
 from litellm._logging import verbose_proxy_logger
+from litellm._uuid import uuid
 from litellm.caching import DualCache
 from litellm.constants import LENGTH_OF_LITELLM_GENERATED_KEY, UI_SESSION_TOKEN_TEAM_ID
 from litellm.litellm_core_utils.duration_parser import duration_in_seconds
@@ -642,6 +642,9 @@ async def generate_key_fn(
     - object_permission: Optional[LiteLLM_ObjectPermissionBase] - key-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"]}. IF null or {} then no object permission.
     - key_type: Optional[str] - Type of key that determines default allowed routes. Options: "llm_api" (can call LLM API routes), "management" (can call management routes), "read_only" (can only call info/read routes), "default" (uses default allowed routes). Defaults to "default".
     - prompts: Optional[List[str]] - List of allowed prompts for the key. If specified, the key will only be able to use these specific prompts.
+    - auto_rotate: Optional[bool] - Whether this key should be automatically rotated (regenerated)
+    - rotation_interval: Optional[str] - How often to auto-rotate this key (e.g., '30s', '30m', '30h', '30d'). Required if auto_rotate=True.
+
     Examples:
 
     1. Allow users to turn on/off pii masking
@@ -1558,6 +1561,8 @@ def _check_model_access_group(
     return True
 
 
+
+
 async def generate_key_helper_fn(  # noqa: PLR0915
     request_type: Literal[
         "user", "key"
@@ -1611,6 +1616,8 @@ async def generate_key_helper_fn(  # noqa: PLR0915
         str
     ] = None,  # object_permission_id <-> LiteLLM_ObjectPermissionTable
     object_permission: Optional[LiteLLM_ObjectPermissionBase] = None,
+    auto_rotate: Optional[bool] = None,
+    rotation_interval: Optional[str] = None,
 ):
     from litellm.proxy.proxy_server import premium_user, prisma_client
 
@@ -1718,6 +1725,14 @@ async def generate_key_helper_fn(  # noqa: PLR0915
             "allowed_routes": allowed_routes or [],
             "object_permission_id": object_permission_id,
         }
+        
+        # Add rotation fields if auto_rotate is enabled
+        if auto_rotate and rotation_interval:
+            key_data.update({
+                "auto_rotate": auto_rotate,
+                "rotation_interval": rotation_interval
+                # last_rotation_at will be null initially - rotation happens on first check
+            })
 
         if (
             get_secret("DISABLE_KEY_NAME", False) is True

litellm/proxy/proxy_config.yaml

@@ -42,4 +42,5 @@ litellm_settings:
   datadog_params:
     turn_off_message_logging: true
   datadog_llm_observability_params:
-    turn_off_message_logging: true
+    turn_off_message_logging: true
+# proxy_config.yaml

litellm/proxy/proxy_server.py

@@ -9,7 +9,6 @@ import subprocess
 import sys
 import time
 import traceback
-from litellm._uuid import uuid
 import warnings
 from datetime import datetime, timedelta
 from typing import (
@@ -27,6 +26,7 @@ from typing import (
     get_type_hints,
 )
 
+from litellm._uuid import uuid
 from litellm.constants import (
     BASE_MCP_ROUTE,
     DEFAULT_MAX_RECURSE_DEPTH,
@@ -3811,9 +3811,10 @@ class ProxyStartupEvent:
         cls, scheduler: AsyncIOScheduler
     ):
         """
-        Initialize the spend tracking background jobs
+        Initialize the spend tracking and other background jobs
         1. CloudZero Background Job
         2. Prometheus Background Job
+        3. Key Rotation Background Job
 
         Args:
             scheduler: The scheduler to add the background jobs to
@@ -3838,6 +3839,41 @@ class ProxyStartupEvent:
             except Exception:
                 PrometheusLogger = None
 
+        ########################################################
+        # Key Rotation Background Job
+        ########################################################
+        from litellm.constants import (
+            LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS,
+            LITELLM_KEY_ROTATION_ENABLED,
+        )
+        
+        key_rotation_enabled: Optional[bool] = str_to_bool(LITELLM_KEY_ROTATION_ENABLED)
+        verbose_proxy_logger.debug(f"key_rotation_enabled: {key_rotation_enabled}")
+        
+        if key_rotation_enabled is True:
+            try:
+                from litellm.proxy.common_utils.key_rotation_manager import (
+                    KeyRotationManager,
+                )
+
+                # Get prisma_client from global scope
+                global prisma_client
+                if prisma_client is not None:
+                    key_rotation_manager = KeyRotationManager(prisma_client)
+                    verbose_proxy_logger.debug(f"Key rotation background job scheduled every {LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS} seconds (LITELLM_KEY_ROTATION_ENABLED=true)")
+                    scheduler.add_job(
+                        key_rotation_manager.process_rotations,
+                        "interval",
+                        seconds=LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS,
+                        id="key_rotation_job"
+                    )
+                else:
+                    verbose_proxy_logger.warning("Key rotation enabled but prisma_client not available")
+            except Exception as e:
+                verbose_proxy_logger.warning(f"Failed to setup key rotation job: {e}")
+        else:
+            verbose_proxy_logger.debug("Key rotation disabled (set LITELLM_KEY_ROTATION_ENABLED=true to enable)")
+
     @classmethod
     async def _setup_prisma_client(
         cls,

litellm/proxy/schema.prisma

@@ -221,6 +221,10 @@ model LiteLLM_VerificationToken {
     created_by String?
     updated_at      DateTime?               @default(now()) @updatedAt @map("updated_at")
     updated_by String?
+    rotation_count   Int?       @default(0) // Number of times key has been rotated
+    auto_rotate      Boolean?   @default(false) // Whether this key should be auto-rotated
+    rotation_interval String?   // How often to rotate (e.g., "30d", "90d")
+    last_rotation_at DateTime?  // When this key was last rotated
     litellm_budget_table LiteLLM_BudgetTable?   @relation(fields: [budget_id], references: [budget_id])
     litellm_organization_table LiteLLM_OrganizationTable?   @relation(fields: [organization_id], references: [organization_id])
     object_permission LiteLLM_ObjectPermissionTable?   @relation(fields: [object_permission_id], references: [object_permission_id])

poetry.lock

@@ -1,4 +1,4 @@
-# This file is automatically @generated by Poetry 2.2.1 and should not be changed by hand.
+# This file is automatically @generated by Poetry 2.2.0 and should not be changed by hand.
 
 [[package]]
 name = "aiohappyeyeballs"
@@ -3804,15 +3804,15 @@ files = [
 
 [[package]]
 name = "litellm-proxy-extras"
-version = "0.2.19"
+version = "0.2.20"
 description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
 optional = true
 python-versions = "!=2.7.*,!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,!=3.6.*,!=3.7.*,>=3.8"
 groups = ["main"]
 markers = "extra == \"proxy\""
 files = [
-    {file = "litellm_proxy_extras-0.2.19-py3-none-any.whl", hash = "sha256:cd4cc5fc639ac24bc99af70b8b56b7cc0480b0e72a2a53638f557beb7e9f64fd"},
-    {file = "litellm_proxy_extras-0.2.19.tar.gz", hash = "sha256:e53592e39eb0b9c3b6cd32f29e6a0a53be51d7ad31dedfb69b58c24073b5b3ec"},
+    {file = "litellm_proxy_extras-0.2.20-py3-none-any.whl", hash = "sha256:80ebe03be210eed99ddab0bf8e1a14169f24a502ecdaa98f4bb5032ebe65a944"},
+    {file = "litellm_proxy_extras-0.2.20.tar.gz", hash = "sha256:8dea657f965122490be480327995ce534e75dc96778cb3e05c562c7996f85cce"},
 ]
 
 [[package]]
@@ -9598,4 +9598,4 @@ utils = ["numpydoc"]
 [metadata]
 lock-version = "2.1"
 python-versions = ">=3.8.1,<4.0, !=3.9.7"
-content-hash = "b775339256b9eb6ae6ffc72cbe7d17d379897a3d72f762cb0df58e4044cded0c"
+content-hash = "1fc7c0289412c4ea4b35448c3a2857eb4c5656796ec51df5f20d4c134c383271"

pyproject.toml

@@ -59,7 +59,7 @@ websockets = {version = "^13.1.0", optional = true}
 boto3 = {version = "1.36.0", optional = true}
 redisvl = {version = "^0.4.1", optional = true, markers = "python_version >= '3.9' and python_version < '3.14'"}
 mcp = {version = "^1.10.0", optional = true, python = ">=3.10"}
-litellm-proxy-extras = {version = "0.2.19", optional = true}
+litellm-proxy-extras = {version = "0.2.20", optional = true}
 rich = {version = "13.7.1", optional = true}
 litellm-enterprise = {version = "0.1.20", optional = true}
 diskcache = {version = "^5.6.1", optional = true}

requirements.txt

@@ -43,7 +43,7 @@ sentry_sdk==2.21.0 # for sentry error handling
 detect-secrets==1.5.0 # Enterprise - secret detection / masking in LLM requests
 cryptography==43.0.1
 tzdata==2025.1 # IANA time zone database
-litellm-proxy-extras==0.2.19 # for proxy extras - e.g. prisma migrations
+litellm-proxy-extras==0.2.20 # for proxy extras - e.g. prisma migrations
 ### LITELLM PACKAGE DEPENDENCIES
 python-dotenv==1.0.1 # for env
 tiktoken==0.8.0 # for calculating usage

schema.prisma

@@ -221,6 +221,10 @@ model LiteLLM_VerificationToken {
     created_by String?
     updated_at      DateTime?               @default(now()) @updatedAt @map("updated_at")
     updated_by String?
+    rotation_count   Int?       @default(0) // Number of times key has been rotated
+    auto_rotate      Boolean?   @default(false) // Whether this key should be auto-rotated
+    rotation_interval String?   // How often to rotate (e.g., "30d", "90d")
+    last_rotation_at DateTime?  // When this key was last rotated
     litellm_budget_table LiteLLM_BudgetTable?   @relation(fields: [budget_id], references: [budget_id])
     litellm_organization_table LiteLLM_OrganizationTable?   @relation(fields: [organization_id], references: [organization_id])
     object_permission LiteLLM_ObjectPermissionTable?   @relation(fields: [object_permission_id], references: [object_permission_id])

tests/test_litellm/proxy/common_utils/test_key_rotation_manager.py

@@ -0,0 +1,194 @@
+"""
+Test key rotation manager functionality
+"""
+import os
+import sys
+from datetime import datetime, timedelta, timezone
+from unittest.mock import AsyncMock, MagicMock
+
+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 TestKeyRotationManager:
+    """Test the KeyRotationManager class functionality."""
+
+    @pytest.mark.asyncio
+    async def test_should_rotate_key_logic(self):
+        """
+        Test the core logic for determining when a key should be rotated.
+        
+        This tests:
+        - Keys with null last_rotation_at should rotate immediately
+        - Keys with recent rotation should not rotate
+        - Keys with old rotation should rotate
+        """
+        # Setup
+        mock_prisma_client = AsyncMock()
+        manager = KeyRotationManager(mock_prisma_client)
+        
+        now = datetime.now(timezone.utc)
+        
+        # Test Case 1: Never rotated (last_rotation_at = None) - should rotate
+        key_never_rotated = LiteLLM_VerificationToken(
+            token="test-token-1",
+            auto_rotate=True,
+            rotation_interval="30s",
+            last_rotation_at=None,
+            rotation_count=0
+        )
+        
+        assert manager._should_rotate_key(key_never_rotated, now) == True
+        
+        # Test Case 2: Recently rotated (10s ago, interval 30s) - should NOT rotate
+        key_recently_rotated = LiteLLM_VerificationToken(
+            token="test-token-2",
+            auto_rotate=True,
+            rotation_interval="30s",
+            last_rotation_at=now - timedelta(seconds=10),
+            rotation_count=1
+        )
+        
+        assert manager._should_rotate_key(key_recently_rotated, now) == False
+        
+        # Test Case 3: Old rotation (60s ago, interval 30s) - should rotate
+        key_old_rotation = LiteLLM_VerificationToken(
+            token="test-token-3",
+            auto_rotate=True,
+            rotation_interval="30s",
+            last_rotation_at=now - timedelta(seconds=60),
+            rotation_count=2
+        )
+        
+        assert manager._should_rotate_key(key_old_rotation, now) == True
+        
+        # Test Case 4: No rotation interval - should NOT rotate
+        key_no_interval = LiteLLM_VerificationToken(
+            token="test-token-4",
+            auto_rotate=True,
+            rotation_interval=None,
+            last_rotation_at=None,
+            rotation_count=0
+        )
+        
+        assert manager._should_rotate_key(key_no_interval, now) == False
+
+    @pytest.mark.asyncio
+    async def test_find_keys_needing_rotation(self):
+        """
+        Test finding keys that need rotation from database.
+        
+        This tests:
+        - Only keys with auto_rotate=True and rotation_interval are considered
+        - Filtering logic works correctly
+        - Database query is constructed properly
+        """
+        # Setup
+        mock_prisma_client = AsyncMock()
+        manager = KeyRotationManager(mock_prisma_client)
+        
+        now = datetime.now(timezone.utc)
+        
+        # Mock database response
+        mock_keys = [
+            LiteLLM_VerificationToken(
+                token="token-1",
+                auto_rotate=True,
+                rotation_interval="30s",
+                last_rotation_at=None,  # Should rotate
+                rotation_count=0
+            ),
+            LiteLLM_VerificationToken(
+                token="token-2",
+                auto_rotate=True,
+                rotation_interval="60s",
+                last_rotation_at=now - timedelta(seconds=30),  # Should NOT rotate (30s < 60s)
+                rotation_count=1
+            ),
+            LiteLLM_VerificationToken(
+                token="token-3",
+                auto_rotate=True,
+                rotation_interval="30s",
+                last_rotation_at=now - timedelta(seconds=45),  # Should rotate (45s > 30s)
+                rotation_count=2
+            )
+        ]
+        
+        mock_prisma_client.db.litellm_verificationtoken.find_many.return_value = mock_keys
+        
+        # Execute
+        keys_needing_rotation = await manager._find_keys_needing_rotation()
+        
+        # Verify database query
+        mock_prisma_client.db.litellm_verificationtoken.find_many.assert_called_once_with(
+            where={
+                "auto_rotate": True,
+                "rotation_interval": {"not": None}
+            }
+        )
+        
+        # Verify filtering logic
+        assert len(keys_needing_rotation) == 2  # token-1 and token-3 should need rotation
+        
+        tokens_needing_rotation = [key.token for key in keys_needing_rotation]
+        assert "token-1" in tokens_needing_rotation  # Never rotated
+        assert "token-2" not in tokens_needing_rotation  # Recently rotated
+        assert "token-3" in tokens_needing_rotation  # Old rotation
+
+    @pytest.mark.asyncio
+    async def test_rotate_key_updates_database(self):
+        """
+        Test that key rotation properly updates the database with new rotation info.
+        
+        This tests:
+        - Rotation count is incremented
+        - last_rotation_at is set to current time
+        - New key token is updated (not old one)
+        """
+        # Setup
+        mock_prisma_client = AsyncMock()
+        manager = KeyRotationManager(mock_prisma_client)
+        
+        # Mock key to rotate
+        key_to_rotate = LiteLLM_VerificationToken(
+            token="old-token",
+            auto_rotate=True,
+            rotation_interval="30s",
+            last_rotation_at=None,
+            rotation_count=0
+        )
+        
+        # Mock regenerate_key_fn response
+        mock_response = GenerateKeyResponse(
+            key="new-api-key",
+            token_id="new-token-id",
+            user_id="test-user"
+        )
+        
+        # Mock the regenerate function
+        from unittest.mock import patch
+        with patch('litellm.proxy.common_utils.key_rotation_manager.regenerate_key_fn', return_value=mock_response):
+            with patch('litellm.proxy.common_utils.key_rotation_manager.KeyManagementEventHooks.async_key_rotated_hook'):
+                # Execute
+                await manager._rotate_key(key_to_rotate)
+        
+        # Verify database update was called with correct data
+        mock_prisma_client.db.litellm_verificationtoken.update.assert_called_once()
+        
+        call_args = mock_prisma_client.db.litellm_verificationtoken.update.call_args
+        
+        # Check the WHERE clause targets the new token
+        assert call_args[1]["where"]["token"] == "new-token-id"
+        
+        # Check the data being updated
+        update_data = call_args[1]["data"]
+        assert update_data["rotation_count"] == 1  # Incremented from 0
+        assert "last_rotation_at" in update_data
+        assert isinstance(update_data["last_rotation_at"], datetime)

tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py

@@ -23,6 +23,7 @@ from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth
 from litellm.proxy.management_endpoints.key_management_endpoints import (
     _common_key_generation_helper,
     _list_key_helper,
+    generate_key_helper_fn,
     prepare_key_update_data,
     validate_key_team_change,
 )
@@ -1091,3 +1092,73 @@ def test_validate_key_team_change_with_member_permissions():
                         team_table=mock_team,
                         route=KeyManagementRoutes.KEY_UPDATE.value
                     )
+
+
+def test_key_rotation_fields_helper():
+    """
+    Test the key data update logic for rotation fields.
+    
+    This test focuses on the core logic that adds rotation fields to key_data
+    when auto_rotate is enabled, without the complexity of full key generation.
+    """
+    # Test Case 1: With rotation enabled
+    key_data = {
+        "models": ["gpt-3.5-turbo"],
+        "user_id": "test-user"
+    }
+    
+    auto_rotate = True
+    rotation_interval = "30d"
+    
+    # Simulate the rotation logic from generate_key_helper_fn
+    if auto_rotate and rotation_interval:
+        key_data.update({
+            "auto_rotate": auto_rotate,
+            "rotation_interval": rotation_interval
+        })
+    
+    # Verify rotation fields are added
+    assert key_data["auto_rotate"] == True
+    assert key_data["rotation_interval"] == "30d"
+    assert key_data["models"] == ["gpt-3.5-turbo"]  # Original fields preserved
+    
+    # Test Case 2: Without rotation enabled
+    key_data2 = {
+        "models": ["gpt-4"],
+        "user_id": "test-user"
+    }
+    
+    auto_rotate2 = False
+    rotation_interval2 = None
+    
+    # Simulate the rotation logic
+    if auto_rotate2 and rotation_interval2:
+        key_data2.update({
+            "auto_rotate": auto_rotate2,
+            "rotation_interval": rotation_interval2
+        })
+    
+    # Verify rotation fields are NOT added
+    assert "auto_rotate" not in key_data2
+    assert "rotation_interval" not in key_data2
+    assert key_data2["models"] == ["gpt-4"]  # Original fields preserved
+    
+    # Test Case 3: auto_rotate=True but no interval
+    key_data3 = {
+        "models": ["claude-3"],
+        "user_id": "test-user"
+    }
+    
+    auto_rotate3 = True
+    rotation_interval3 = None
+    
+    # Simulate the rotation logic
+    if auto_rotate3 and rotation_interval3:
+        key_data3.update({
+            "auto_rotate": auto_rotate3,
+            "rotation_interval": rotation_interval3
+        })
+    
+    # Verify rotation fields are NOT added (missing interval)
+    assert "auto_rotate" not in key_data3
+    assert "rotation_interval" not in key_data3