diff --git a/.circleci/config.yml b/.circleci/config.yml index ba42ccbc36..ef6445ca0c 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -148,8 +148,7 @@ jobs: python -m pip install types-requests types-setuptools types-redis types-PyYAML if ! python -m mypy . \ --config-file mypy.ini \ - --ignore-missing-imports \ - --no-incremental; then + --ignore-missing-imports; 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 4cd15e8ae4..974e95a07b 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -613,8 +613,6 @@ 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 deleted file mode 100644 index 0a94ef6ff6..0000000000 Binary files a/litellm-proxy-extras/dist/litellm_proxy_extras-0.2.20-py3-none-any.whl and /dev/null 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 deleted file mode 100644 index 1562aacc22..0000000000 Binary files a/litellm-proxy-extras/dist/litellm_proxy_extras-0.2.20.tar.gz and /dev/null differ diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 9d3726d29d..2b1e20820f 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -221,10 +221,6 @@ 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 dd0b2138dc..1c368d5807 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.20" +version = "0.2.19" 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.20" +version = "0.2.19" version_files = [ "pyproject.toml:version", "../requirements.txt:litellm-proxy-extras==", diff --git a/litellm/constants.py b/litellm/constants.py index 1ed9f237a2..6e70ae0671 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -989,11 +989,7 @@ 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/proxy/_types.py b/litellm/proxy/_types.py index eea370fa22..a92b7d8ad2 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -787,14 +787,6 @@ 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): @@ -1810,10 +1802,6 @@ 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=()) @@ -1944,23 +1932,6 @@ 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 deleted file mode 100644 index 319d5ed523..0000000000 --- a/litellm/proxy/common_utils/key_rotation_manager.py +++ /dev/null @@ -1,146 +0,0 @@ -""" -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 363cbf7182..eb389450f9 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -642,9 +642,6 @@ 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 @@ -1561,8 +1558,6 @@ def _check_model_access_group( return True - - async def generate_key_helper_fn( # noqa: PLR0915 request_type: Literal[ "user", "key" @@ -1616,8 +1611,6 @@ 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 @@ -1725,14 +1718,6 @@ 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 b9cd33175e..921b564407 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -36,5 +36,3 @@ guardrails: breakdown: true payload: true dev_info: true - -# proxy_config.yaml \ No newline at end of file diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 399cadac24..7255161b8b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3802,10 +3802,9 @@ class ProxyStartupEvent: cls, scheduler: AsyncIOScheduler ): """ - Initialize the spend tracking and other background jobs + Initialize the spend tracking 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 @@ -3830,41 +3829,6 @@ 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 9d3726d29d..2b1e20820f 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -221,10 +221,6 @@ 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 9677775ad2..07822390af 100644 --- a/poetry.lock +++ b/poetry.lock @@ -2715,16 +2715,21 @@ files = [ [[package]] name = "litellm-proxy-extras" -version = "0.2.20" +version = "0.2.19" 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 = [ +<<<<<<< HEAD {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"}, <<<<<<< HEAD +======= + {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"}, +>>>>>>> parent of ded823a3a3 ([Feat] Initial support for scheduled key rotations (#14877)) ] [[package]] @@ -6742,4 +6747,4 @@ utils = ["numpydoc"] [metadata] lock-version = "2.1" python-versions = ">=3.8.1,<4.0, !=3.9.7" -content-hash = "6ea59a97b3d69dc5fcc45645131e2fc843ec533854dba19852cecf8e10da5e7a" \ No newline at end of file +content-hash = "6ea59a97b3d69dc5fcc45645131e2fc843ec533854dba19852cecf8e10da5e7a" diff --git a/pyproject.toml b/pyproject.toml index 5d4605ed5b..cc45c35cda 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.20", optional = true} +litellm-proxy-extras = {version = "0.2.19", 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 10929ab819..3c55f682a0 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.20 # for proxy extras - e.g. prisma migrations +litellm-proxy-extras==0.2.19 # 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 9d3726d29d..2b1e20820f 100644 --- a/schema.prisma +++ b/schema.prisma @@ -221,10 +221,6 @@ 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 deleted file mode 100644 index 8aab67bf58..0000000000 --- a/tests/test_litellm/proxy/common_utils/test_key_rotation_manager.py +++ /dev/null @@ -1,194 +0,0 @@ -""" -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 14b5833f14..da64a866d5 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,7 +23,6 @@ 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, ) @@ -1092,73 +1091,3 @@ 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