diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 9d3726d29d6..766625145f6 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -225,6 +225,7 @@ model LiteLLM_VerificationToken { 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 + key_rotation_at DateTime? // When this key should next be 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/_types.py b/litellm/proxy/_types.py index 345322a0021..62910ad4574 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -842,6 +842,8 @@ class UpdateKeyRequest(KeyRequestBase): metadata: Optional[dict] = None temp_budget_increase: Optional[float] = None temp_budget_expiry: Optional[datetime] = None + auto_rotate: Optional[bool] = None + rotation_interval: Optional[str] = None @model_validator(mode="after") def validate_temp_budget(self) -> "UpdateKeyRequest": @@ -1814,6 +1816,7 @@ class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase): 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 + key_rotation_at: Optional[datetime] = None # When this key should next be rotated model_config = ConfigDict(protected_namespaces=()) diff --git a/litellm/proxy/common_utils/key_rotation_manager.py b/litellm/proxy/common_utils/key_rotation_manager.py index 319d5ed5233..3c367eafbc1 100644 --- a/litellm/proxy/common_utils/key_rotation_manager.py +++ b/litellm/proxy/common_utils/key_rotation_manager.py @@ -4,7 +4,7 @@ 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 datetime import datetime, timezone from typing import List from litellm._logging import verbose_proxy_logger @@ -16,6 +16,7 @@ from litellm.proxy._types import ( ) from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks from litellm.proxy.management_endpoints.key_management_endpoints import ( + _calculate_key_rotation_time, regenerate_key_fn, ) from litellm.proxy.utils import PrismaClient @@ -60,49 +61,39 @@ class KeyRotationManager: async def _find_keys_needing_rotation(self) -> List[LiteLLM_VerificationToken]: """ - Find keys that are due for rotation based on their rotation interval. + Find keys that are due for rotation based on their key_rotation_at timestamp. 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 + - key_rotation_at is null (needs initial setup) OR key_rotation_at <= now """ + now = datetime.now(timezone.utc) + 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 + "OR": [ + {"key_rotation_at": None}, # Keys that need initial rotation time setup + {"key_rotation_at": {"lte": now}} # Keys where rotation time has passed + ] } ) - # 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 + return keys_with_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. + Determine if a key should be rotated based on key_rotation_at timestamp. """ if not key.rotation_interval: return False - # If never rotated, rotate immediately - if key.last_rotation_at is None: + # If key_rotation_at is not set, rotate immediately (and set it) + if key.key_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 + # Check if the rotation time has passed + return now >= key.key_rotation_at async def _rotate_key(self, key: LiteLLM_VerificationToken): """ @@ -125,12 +116,16 @@ class KeyRotationManager: ) # Update the NEW key with rotation info (regenerate_key_fn creates a new token) - if isinstance(response, GenerateKeyResponse) and response.token_id: + if isinstance(response, GenerateKeyResponse) and response.token_id and key.rotation_interval: + # Calculate next rotation time using helper function + now = datetime.now(timezone.utc) + next_rotation_time = _calculate_key_rotation_time(key.rotation_interval) 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) + "last_rotation_at": now, + "key_rotation_at": next_rotation_time } ) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index c230018eceb..007c0164be4 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -87,6 +87,38 @@ def _get_user_in_team( return None +def _calculate_key_rotation_time(rotation_interval: str) -> datetime: + """ + Helper function to calculate the next rotation time for a key based on the rotation interval. + + Args: + rotation_interval: String representing the rotation interval (e.g., '30d', '90d', '1h') + + Returns: + datetime: The calculated next rotation time in UTC + """ + now = datetime.now(timezone.utc) + interval_seconds = duration_in_seconds(rotation_interval) + return now + timedelta(seconds=interval_seconds) + + +def _set_key_rotation_fields(data: dict, auto_rotate: bool, rotation_interval: Optional[str]) -> None: + """ + Helper function to set rotation fields in key data if auto_rotate is enabled. + + Args: + data: Dictionary to update with rotation fields + auto_rotate: Whether auto rotation is enabled + rotation_interval: The rotation interval string (required if auto_rotate is True) + """ + if auto_rotate and rotation_interval: + data.update({ + "auto_rotate": auto_rotate, + "rotation_interval": rotation_interval, + "key_rotation_at": _calculate_key_rotation_time(rotation_interval) + }) + + def _is_allowed_to_make_key_request( user_api_key_dict: UserAPIKeyAuth, user_id: Optional[str], team_id: Optional[str] ) -> bool: @@ -1071,6 +1103,8 @@ async def update_key_fn( - allowed_routes: Optional[list] - List of allowed routes for the key. Store the actual route or store a wildcard pattern for a set of routes. Example - ["/chat/completions", "/embeddings", "/keys/*"] - prompts: Optional[List[str]] - List of allowed prompts for the key. If specified, the key will only be able to use these specific prompts. - 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. + - auto_rotate: Optional[bool] - Whether this key should be automatically rotated + - rotation_interval: Optional[str] - How often to rotate this key (e.g., '30d', '90d'). Required if auto_rotate=True Example: ```bash curl --location 'http://0.0.0.0:4000/key/update' \ @@ -1162,6 +1196,13 @@ async def update_key_fn( existing_key_token=existing_key_row.token, ) + # Handle rotation fields if auto_rotate is being enabled + _set_key_rotation_fields( + non_default_values, + non_default_values.get("auto_rotate", False), + non_default_values.get("rotation_interval") + ) + _data = {**non_default_values, "token": key} response = await prisma_client.update_data(token=key, data=_data) @@ -1727,12 +1768,11 @@ async def generate_key_helper_fn( # noqa: PLR0915 } # 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 - }) + _set_key_rotation_fields( + data=key_data, + auto_rotate=auto_rotate or False, + rotation_interval=rotation_interval + ) if ( get_secret("DISABLE_KEY_NAME", False) is True diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 9d3726d29d6..766625145f6 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -225,6 +225,7 @@ model LiteLLM_VerificationToken { 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 + key_rotation_at DateTime? // When this key should next be 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/schema.prisma b/schema.prisma index 9d3726d29d6..766625145f6 100644 --- a/schema.prisma +++ b/schema.prisma @@ -225,6 +225,7 @@ model LiteLLM_VerificationToken { 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 + key_rotation_at DateTime? // When this key should next be 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 index 8aab67bf588..6b3b4c92416 100644 --- a/tests/test_litellm/proxy/common_utils/test_key_rotation_manager.py +++ b/tests/test_litellm/proxy/common_utils/test_key_rotation_manager.py @@ -26,9 +26,9 @@ class TestKeyRotationManager: 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 + - Keys with null key_rotation_at should rotate immediately + - Keys with future key_rotation_at should not rotate + - Keys with past key_rotation_at should rotate """ # Setup mock_prisma_client = AsyncMock() @@ -36,45 +36,56 @@ class TestKeyRotationManager: now = datetime.now(timezone.utc) - # Test Case 1: Never rotated (last_rotation_at = None) - should rotate - key_never_rotated = LiteLLM_VerificationToken( + # Test Case 1: No rotation time set (key_rotation_at = None) - should rotate + key_no_rotation_time = LiteLLM_VerificationToken( token="test-token-1", auto_rotate=True, rotation_interval="30s", - last_rotation_at=None, + key_rotation_at=None, rotation_count=0 ) - assert manager._should_rotate_key(key_never_rotated, now) == True + assert manager._should_rotate_key(key_no_rotation_time, now) == True - # Test Case 2: Recently rotated (10s ago, interval 30s) - should NOT rotate - key_recently_rotated = LiteLLM_VerificationToken( + # Test Case 2: Future rotation time - should NOT rotate + key_future_rotation = LiteLLM_VerificationToken( token="test-token-2", auto_rotate=True, rotation_interval="30s", - last_rotation_at=now - timedelta(seconds=10), + key_rotation_at=now + timedelta(seconds=10), rotation_count=1 ) - assert manager._should_rotate_key(key_recently_rotated, now) == False + assert manager._should_rotate_key(key_future_rotation, now) == False - # Test Case 3: Old rotation (60s ago, interval 30s) - should rotate - key_old_rotation = LiteLLM_VerificationToken( + # Test Case 3: Past rotation time - should rotate + key_past_rotation = LiteLLM_VerificationToken( token="test-token-3", auto_rotate=True, rotation_interval="30s", - last_rotation_at=now - timedelta(seconds=60), + key_rotation_at=now - timedelta(seconds=10), rotation_count=2 ) - assert manager._should_rotate_key(key_old_rotation, now) == True + assert manager._should_rotate_key(key_past_rotation, now) == True - # Test Case 4: No rotation interval - should NOT rotate - key_no_interval = LiteLLM_VerificationToken( + # Test Case 4: Exact rotation time - should rotate + key_exact_rotation = LiteLLM_VerificationToken( token="test-token-4", auto_rotate=True, + rotation_interval="30s", + key_rotation_at=now, + rotation_count=1 + ) + + assert manager._should_rotate_key(key_exact_rotation, now) == True + + # Test Case 5: No rotation interval - should NOT rotate + key_no_interval = LiteLLM_VerificationToken( + token="test-token-5", + auto_rotate=True, rotation_interval=None, - last_rotation_at=None, + key_rotation_at=None, rotation_count=0 ) @@ -86,61 +97,63 @@ class TestKeyRotationManager: 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 + - Only keys with auto_rotate=True are considered + - Database query filters by key_rotation_at properly + - Keys are returned based on key_rotation_at being null or <= now """ # Setup mock_prisma_client = AsyncMock() manager = KeyRotationManager(mock_prisma_client) - now = datetime.now(timezone.utc) + # Use a fixed timestamp to avoid timing issues in tests + now = datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc) - # Mock database response + # Mock database response - these are the keys the database query would return mock_keys = [ LiteLLM_VerificationToken( token="token-1", auto_rotate=True, rotation_interval="30s", - last_rotation_at=None, # Should rotate + key_rotation_at=None, # Should rotate (null key_rotation_at) 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) + key_rotation_at=now - timedelta(seconds=10), # Should rotate (past time) 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() + # Mock datetime.now to return our fixed timestamp + from unittest.mock import patch + with patch('litellm.proxy.common_utils.key_rotation_manager.datetime') as mock_datetime: + mock_datetime.now.return_value = now + mock_datetime.side_effect = lambda *args, **kwargs: datetime(*args, **kwargs) + + # Execute + keys_needing_rotation = await manager._find_keys_needing_rotation() - # Verify database query + # Verify database query - should use OR condition for key_rotation_at mock_prisma_client.db.litellm_verificationtoken.find_many.assert_called_once_with( where={ "auto_rotate": True, - "rotation_interval": {"not": None} + "OR": [ + {"key_rotation_at": None}, + {"key_rotation_at": {"lte": now}} + ] } ) - # Verify filtering logic - assert len(keys_needing_rotation) == 2 # token-1 and token-3 should need rotation + # Verify all keys returned by database query are included (no additional filtering) + assert len(keys_needing_rotation) == 2 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 + assert "token-1" in tokens_needing_rotation # Null key_rotation_at + assert "token-2" in tokens_needing_rotation # Past key_rotation_at @pytest.mark.asyncio async def test_rotate_key_updates_database(self): @@ -150,6 +163,7 @@ class TestKeyRotationManager: This tests: - Rotation count is incremented - last_rotation_at is set to current time + - key_rotation_at is set to next rotation time - New key token is updated (not old one) """ # Setup @@ -162,6 +176,7 @@ class TestKeyRotationManager: auto_rotate=True, rotation_interval="30s", last_rotation_at=None, + key_rotation_at=None, rotation_count=0 ) @@ -192,3 +207,11 @@ class TestKeyRotationManager: assert update_data["rotation_count"] == 1 # Incremented from 0 assert "last_rotation_at" in update_data assert isinstance(update_data["last_rotation_at"], datetime) + assert "key_rotation_at" in update_data + assert isinstance(update_data["key_rotation_at"], datetime) + + # Verify key_rotation_at is set to future time (30s from now) + now = datetime.now(timezone.utc) + next_rotation = update_data["key_rotation_at"] + time_diff = (next_rotation - now).total_seconds() + assert 25 <= time_diff <= 35 # Should be around 30 seconds, allow some tolerance 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 14b5833f14f..e3aa7d58872 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 @@ -1162,3 +1162,75 @@ def test_key_rotation_fields_helper(): # Verify rotation fields are NOT added (missing interval) assert "auto_rotate" not in key_data3 assert "rotation_interval" not in key_data3 + + +@pytest.mark.asyncio +async def test_update_key_fn_auto_rotate_enable(): + """Test that update_key_fn properly handles enabling auto rotation.""" + from litellm.proxy._types import LiteLLM_VerificationToken, UpdateKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import ( + prepare_key_update_data, + ) + + # Mock existing key + existing_key = LiteLLM_VerificationToken( + token="test-token", + key_alias="test-key", + models=["gpt-3.5-turbo"], + user_id="test-user", + team_id=None, + auto_rotate=False, + rotation_interval=None, + metadata={} + ) + + # Test enabling auto rotation + update_request = UpdateKeyRequest( + key="test-token", + auto_rotate=True, + rotation_interval="30d" + ) + + result = await prepare_key_update_data( + data=update_request, + existing_key_row=existing_key + ) + + # Verify rotation fields are included + assert result["auto_rotate"] is True + assert result["rotation_interval"] == "30d" + + +@pytest.mark.asyncio +async def test_update_key_fn_auto_rotate_disable(): + """Test that update_key_fn properly handles disabling auto rotation.""" + from litellm.proxy._types import LiteLLM_VerificationToken, UpdateKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import ( + prepare_key_update_data, + ) + + # Mock existing key with rotation enabled + existing_key = LiteLLM_VerificationToken( + token="test-token", + key_alias="test-key", + models=["gpt-3.5-turbo"], + user_id="test-user", + team_id=None, + auto_rotate=True, + rotation_interval="30d", + metadata={} + ) + + # Test disabling auto rotation + update_request = UpdateKeyRequest( + key="test-token", + auto_rotate=False + ) + + result = await prepare_key_update_data( + data=update_request, + existing_key_row=existing_key + ) + + # Verify auto_rotate is set to False + assert result["auto_rotate"] is False diff --git a/ui/litellm-dashboard/src/components/common_components/AutoRotationView.tsx b/ui/litellm-dashboard/src/components/common_components/AutoRotationView.tsx new file mode 100644 index 00000000000..07570477ed7 --- /dev/null +++ b/ui/litellm-dashboard/src/components/common_components/AutoRotationView.tsx @@ -0,0 +1,129 @@ +import React from "react"; +import { Card, Text, Badge } from "@tremor/react"; +import { RefreshIcon, ClockIcon } from "@heroicons/react/outline"; + +interface AutoRotationViewProps { + autoRotate?: boolean; + rotationInterval?: string; + lastRotationAt?: string; + keyRotationAt?: string; + nextRotationAt?: string; + variant?: "card" | "inline"; + className?: string; +} + +const AutoRotationView: React.FC = ({ + autoRotate = false, + rotationInterval, + lastRotationAt, + keyRotationAt, + nextRotationAt, + variant = "card", + className = "", +}) => { + const formatTimestamp = (timestamp: string | Date) => { + const date = new Date(timestamp); + const dateStr = date.toLocaleDateString("en-US", { + year: "numeric", + month: "short", + day: "numeric", + }); + const timeStr = date.toLocaleTimeString("en-US", { + hour: "numeric", + minute: "2-digit", + hour12: true, + }); + return `${dateStr} at ${timeStr}`; + }; + + const content = ( +
+ {/* Status Section */} +
+
+ + Auto-Rotation + + {autoRotate ? "Enabled" : "Disabled"} + + {autoRotate && rotationInterval && ( + <> + • + Every {rotationInterval} + + )} +
+
+ + {/* Rotation History - Show if there's any rotation data OR if auto-rotation is enabled */} + {(autoRotate || lastRotationAt || keyRotationAt || nextRotationAt) && ( +
+ {/* Last Rotation - Show when available */} + {lastRotationAt && ( +
+ +
+ Last Rotation + {formatTimestamp(lastRotationAt)} +
+
+ )} + + {/* Next Scheduled Rotation - Show when available */} + {(keyRotationAt || nextRotationAt) && ( +
+ +
+ Next Scheduled Rotation + + {formatTimestamp(nextRotationAt || keyRotationAt || "")} + +
+
+ )} + + {/* No rotation data message - Only show if auto-rotation is enabled but no data */} + {autoRotate && !lastRotationAt && !keyRotationAt && !nextRotationAt && ( +
+ + No rotation history available +
+ )} +
+ )} + + {/* Disabled State - Only show if auto-rotation is disabled AND there's no rotation history */} + {!autoRotate && !lastRotationAt && !keyRotationAt && !nextRotationAt && ( +
+ + Auto-rotation is not enabled for this key +
+ )} +
+ ); + + if (variant === "card") { + return ( +
+
+
+ Auto-Rotation + + Automatic key rotation settings and status for this key + +
+
+ {content} +
+ ); + } + + return ( +
+ Auto-Rotation + {content} +
+ ); +}; + +export default AutoRotationView; diff --git a/ui/litellm-dashboard/src/components/common_components/KeyLifecycleSettings.tsx b/ui/litellm-dashboard/src/components/common_components/KeyLifecycleSettings.tsx index ca0d8f67b72..8e286574da4 100644 --- a/ui/litellm-dashboard/src/components/common_components/KeyLifecycleSettings.tsx +++ b/ui/litellm-dashboard/src/components/common_components/KeyLifecycleSettings.tsx @@ -1,4 +1,4 @@ -import React from "react"; +import React, { useState } from "react"; import { Select, Tooltip, Divider, Switch } from "antd"; import { InfoCircleOutlined } from "@ant-design/icons"; import { TextInput } from "@tremor/react"; @@ -20,6 +20,31 @@ const KeyLifecycleSettings: React.FC = ({ rotationInterval, onRotationIntervalChange, }) => { + // Predefined intervals + const predefinedIntervals = ["7d", "30d", "90d", "180d", "365d"]; + + // Check if current interval is custom + const isCustomInterval = rotationInterval && !predefinedIntervals.includes(rotationInterval); + + const [showCustomInput, setShowCustomInput] = useState(isCustomInterval); + const [customInterval, setCustomInterval] = useState(isCustomInterval ? rotationInterval : ""); + + const handleIntervalChange = (value: string) => { + if (value === "custom") { + setShowCustomInput(true); + // Don't change the actual interval yet, wait for custom input + } else { + setShowCustomInput(false); + setCustomInterval(""); + onRotationIntervalChange(value); + } + }; + + const handleCustomIntervalChange = (e: React.ChangeEvent) => { + const value = e.target.value; + setCustomInterval(value); + onRotationIntervalChange(value); + }; return (
{/* Key Expiry Section */} @@ -71,18 +96,34 @@ const KeyLifecycleSettings: React.FC = ({ - +
+ + + {showCustomInput && ( +
+ +
+ Supported formats: seconds (s), minutes (m), hours (h), days (d) +
+
+ )} +
)} diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx index 9933c4ee5bb..ccc3e0cd9c9 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx +++ b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx @@ -82,6 +82,11 @@ export interface KeyResponse { mcp_access_groups?: string[]; vector_stores: string[]; }; + auto_rotate?: boolean; + rotation_interval?: string; + last_rotation_at?: string; + key_rotation_at?: string; + next_rotation_at?: string; } interface KeyListResponse { diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index c74511e6c0a..dc523da7667 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -327,13 +327,10 @@ const CreateKey: React.FC = ({ }; } - // Add auto-rotation settings to the metadata + // Add auto-rotation settings as top-level fields if (autoRotationEnabled) { - metadata = { - ...metadata, - auto_rotation_enabled: true, - rotation_interval: rotationInterval - }; + formValues.auto_rotate = true; + formValues.rotation_interval = rotationInterval; } // Handle duration field for key expiry diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx index b758fac6643..07a311da893 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx @@ -12,6 +12,7 @@ import { extractLoggingSettings, formatMetadataForDisplay } from "../key_info_ut import { fetchMCPAccessGroups } from "../networking" import { mapInternalToDisplayNames, mapDisplayToInternalNames } from "../callback_info_helpers" import GuardrailSelector from "@/components/guardrails/GuardrailSelector" +import KeyLifecycleSettings from "../common_components/KeyLifecycleSettings" interface KeyEditViewProps { keyData: KeyResponse @@ -65,6 +66,8 @@ export function KeyEditView({ ? mapInternalToDisplayNames(keyData.metadata.litellm_disabled_callbacks) : [], ) + const [autoRotationEnabled, setAutoRotationEnabled] = useState(keyData.auto_rotate || false) + const [rotationInterval, setRotationInterval] = useState(keyData.rotation_interval || "") const fetchMcpAccessGroups = async () => { if (!accessToken) return @@ -145,6 +148,8 @@ export function KeyEditView({ disabled_callbacks: Array.isArray(keyData.metadata?.litellm_disabled_callbacks) ? mapInternalToDisplayNames(keyData.metadata.litellm_disabled_callbacks) : [], + auto_rotate: keyData.auto_rotate || false, + ...(keyData.rotation_interval && { rotation_interval: keyData.rotation_interval }), } useEffect(() => { @@ -164,9 +169,22 @@ export function KeyEditView({ disabled_callbacks: Array.isArray(keyData.metadata?.litellm_disabled_callbacks) ? mapInternalToDisplayNames(keyData.metadata.litellm_disabled_callbacks) : [], + auto_rotate: keyData.auto_rotate || false, + ...(keyData.rotation_interval && { rotation_interval: keyData.rotation_interval }) }) }, [keyData, form]) + // Sync auto-rotation state with form values + useEffect(() => { + form.setFieldValue("auto_rotate", autoRotationEnabled) + }, [autoRotationEnabled, form]) + + useEffect(() => { + if (rotationInterval) { + form.setFieldValue("rotation_interval", rotationInterval) + } + }, [rotationInterval, form]) + console.log("premiumUser:", premiumUser) return ( @@ -291,6 +309,17 @@ export function KeyEditView({ + {/* Auto-Rotation Settings */} +
+ +
+ {/* Hidden form field for token */} + {/* Hidden form fields for auto-rotation */} + + +
Cancel diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx index f5e18b68046..c7fa194f971 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx @@ -30,6 +30,7 @@ import { extractLoggingSettings, formatMetadataForDisplay } from "../key_info_ut import { CopyIcon, CheckIcon } from "lucide-react" import { callback_map, mapInternalToDisplayNames, mapDisplayToInternalNames } from "../callback_info_helpers" import { parseErrorMessage } from "../shared/errorUtils" +import AutoRotationView from "../common_components/AutoRotationView" interface KeyInfoViewProps { keyId: string @@ -484,6 +485,15 @@ export default function KeyInfoView({ } variant="card" /> + + @@ -559,6 +569,16 @@ export default function KeyInfoView({ {currentKeyData.expires ? formatTimestamp(currentKeyData.expires) : "Never"}
+ +
Spend ${formatNumberWithCommas(currentKeyData.spend, 4)} USD