From b39218059a381d4ba842078e39a4b4aa2cd5a458 Mon Sep 17 00:00:00 2001
From: Harshit28j <Harshitjain0562@gmail.com>
Date: Sat, 28 Feb 2026 16:34:23 +0530
Subject: [PATCH] fix req change

---
 .../key_management_endpoints.py               |  5 ++-
 .../test_key_rotation_integration.py          | 37 +++++++++++++++++++
 2 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py
index 01bc5704502..f027c068784 100644
--- a/litellm/proxy/management_endpoints/key_management_endpoints.py
+++ b/litellm/proxy/management_endpoints/key_management_endpoints.py
@@ -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}
diff --git a/tests/test_litellm/proxy/common_utils/test_key_rotation_integration.py b/tests/test_litellm/proxy/common_utils/test_key_rotation_integration.py
index 65c89a35f0b..234b83bcd95 100644
--- a/tests/test_litellm/proxy/common_utils/test_key_rotation_integration.py
+++ b/tests/test_litellm/proxy/common_utils/test_key_rotation_integration.py
@@ -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"