scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 109

fix req change

Browse source
This commit is contained in:
Harshit28j 2026-02-28 16:34:23 +05:30
parent 539f2eeba4
commit b39218059a
2 changed files with 41 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -126,7 +126,7 @@ def _calculate_key_rotation_time(rotation_interval: str) -> datetime:
def _set_key_rotation_fields(
data: dict, auto_rotate: bool, rotation_interval: Optional[str]
data: dict, auto_rotate: bool, rotation_interval: Optional[str], existing_key_alias: Optional[str] = None
) -> None:
"""
Helper function to set rotation fields in key data if auto_rotate is enabled.
@ -135,12 +135,14 @@ def _set_key_rotation_fields(
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)
existing_key_alias: The existing key alias from the database (if any)
"""
if auto_rotate and rotation_interval:
if (
litellm._key_management_settings is not None
and litellm._key_management_settings.store_virtual_keys is True
and data.get("key_alias") is None
and existing_key_alias is None
):
raise ProxyException(
message="key_alias is required when auto_rotate=True and store_virtual_keys is enabled. This ensures stable secret naming during rotation.",
@ -1958,6 +1960,7 @@ async def update_key_fn(
non_default_values,
non_default_values.get("auto_rotate", False),
non_default_values.get("rotation_interval"),
existing_key_alias=existing_key_row.key_alias,
)
_data = {**non_default_values, "token": key}

37
tests/test_litellm/proxy/common_utils/test_key_rotation_integration.py
View file

@ -258,3 +258,40 @@ class TestKeyRotationSecretNamingStability:
_set_key_rotation_fields(data, auto_rotate=True, rotation_interval="30d")
assert data["auto_rotate"] is True
assert "key_rotation_at" in data
@pytest.mark.asyncio
async def test_set_key_rotation_fields_with_existing_alias(self):
"""
Tests that _set_key_rotation_fields allows enabling rotation
if the key already has an alias in the database (even if not in current request).
"""
from litellm.proxy.management_endpoints.key_management_endpoints import _set_key_rotation_fields
from unittest.mock import MagicMock, patch
mock_settings = MagicMock()
mock_settings.store_virtual_keys = True
with patch("litellm._key_management_settings", mock_settings):
# 1. No alias in request, but HAS existing_key_alias
data = {"auto_rotate": True}
_set_key_rotation_fields(
data,
auto_rotate=True,
rotation_interval="30d",
existing_key_alias="already-exists-in-db"
)
# Should NOT raise, and field should be set
assert data["auto_rotate"] is True
assert "key_rotation_at" in data
# 2. Verify it still fails if NO alias AND NO existing_key_alias
from litellm.proxy._types import ProxyException
data_fail = {"auto_rotate": True}
with pytest.raises(ProxyException) as exc:
_set_key_rotation_fields(
data_fail,
auto_rotate=True,
rotation_interval="30d",
existing_key_alias=None
)
assert str(exc.value.code) == "400"