diff --git a/.circleci/config.yml b/.circleci/config.yml index ef6445ca0c..ba42ccbc36 100644 --- a/.circleci/config.yml +++ b/.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 diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index f70701886b..ad3afd59a0 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/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 diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.2.20-py3-none-any.whl b/litellm-proxy-extras/dist/litellm_proxy_extras-0.2.20-py3-none-any.whl new file mode 100644 index 0000000000..0a94ef6ff6 Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.2.20-py3-none-any.whl differ diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.2.20.tar.gz b/litellm-proxy-extras/dist/litellm_proxy_extras-0.2.20.tar.gz new file mode 100644 index 0000000000..1562aacc22 Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.2.20.tar.gz differ diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 2b1e20820f..9d3726d29d 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/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]) diff --git a/litellm-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml index 1c368d5807..dd0b2138dc 100644 --- a/litellm-proxy-extras/pyproject.toml +++ b/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==", diff --git a/litellm/constants.py b/litellm/constants.py index 6e70ae0671..1ed9f237a2 100644 --- a/litellm/constants.py +++ b/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" diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index e5cc96ea52..e04c75655d 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/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, diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 6f2732af37..345322a002 100644 --- a/litellm/proxy/_types.py +++ b/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): diff --git a/litellm/proxy/common_utils/key_rotation_manager.py b/litellm/proxy/common_utils/key_rotation_manager.py new file mode 100644 index 0000000000..319d5ed523 --- /dev/null +++ b/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 + ) + \ No newline at end of file diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index aa56cc68a8..c230018ece 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/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 diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index 9897e17d29..b35c7e95a1 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/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 \ No newline at end of file + turn_off_message_logging: true +# proxy_config.yaml diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 23d9acfc1e..f51b65c915 100644 --- a/litellm/proxy/proxy_server.py +++ b/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, diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 2b1e20820f..9d3726d29d 100644 --- a/litellm/proxy/schema.prisma +++ b/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]) diff --git a/poetry.lock b/poetry.lock index b6ed7801be..c05400e139 100644 --- a/poetry.lock +++ b/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" diff --git a/pyproject.toml b/pyproject.toml index 77d80732b4..f8704342b1 100644 --- a/pyproject.toml +++ b/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} diff --git a/requirements.txt b/requirements.txt index 3c55f682a0..10929ab819 100644 --- a/requirements.txt +++ b/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 diff --git a/schema.prisma b/schema.prisma index 2b1e20820f..9d3726d29d 100644 --- a/schema.prisma +++ b/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]) diff --git a/tests/test_litellm/proxy/common_utils/test_key_rotation_manager.py b/tests/test_litellm/proxy/common_utils/test_key_rotation_manager.py new file mode 100644 index 0000000000..8aab67bf58 --- /dev/null +++ b/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) diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index da64a866d5..14b5833f14 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/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