scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-19 00:01:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 122

fix(proxy): sync AWS Secrets Manager on body-less key regenerate

Browse source 
POST /key/{key}/regenerate with no request body reaches async_key_rotated_hook with data=None, and the secret manager sync was gated on data being present, so the rotated key never reached AWS Secrets Manager and the revoked key stayed stored. Gate on response.token_id only and read the requested alias null-safely

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin 2026-09-16 17:56:39 +00:00
parent 4e996400e2
commit ed18626291
2 changed files with 43 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
litellm/proxy/hooks/key_management_event_hooks.py
View file

@ -153,10 +153,11 @@ class KeyManagementEventHooks:
from litellm.proxy.proxy_server import litellm_proxy_admin_name
# Store the generated key in the secret manager - non-blocking, independent operation
if data is not None and response.token_id is not None:
if response.token_id is not None:
try:
initial_secret_name: Final = existing_key_row.key_alias or f"virtual-key-{existing_key_row.token}"
new_secret_name: Final = response.key_alias or data.key_alias or initial_secret_name
requested_alias: Final = data.key_alias if data is not None else None
new_secret_name: Final = response.key_alias or requested_alias or initial_secret_name
verbose_proxy_logger.info(
"Updating secret in secret manager: secret_name=%s",
new_secret_name,

40
tests/test_litellm/proxy/hooks/test_key_management_event_hooks.py
View file

@ -416,6 +416,46 @@ class TestRotateVirtualKeyInSecretManager:
assert call_kwargs["new_secret_name"] == "test-key-alias-new"
assert call_kwargs["new_secret_value"] == "sk-new-key"
@pytest.mark.parametrize("key_alias", ["test-key-alias", None])
@pytest.mark.asyncio
async def test_rotated_hook_without_request_body_syncs_secret_manager(
self, monkeypatch: pytest.MonkeyPatch, key_alias: str | None
):
"""POST /key/{key}/regenerate with no body (data=None) must still write the new key to the secret manager."""
import litellm
from litellm.proxy._types import GenerateKeyResponse, LiteLLM_VerificationToken
from litellm.secret_managers.base_secret_manager import BaseSecretManager
from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem
mock_secret_manager: Final = MagicMock(spec=BaseSecretManager)
mock_secret_manager.async_rotate_secret = AsyncMock(return_value={"status": "success"})
monkeypatch.setattr(litellm, "secret_manager_client", mock_secret_manager)
monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.AWS_SECRET_MANAGER)
monkeypatch.setattr(
litellm,
"_key_management_settings",
KeyManagementSettings(store_virtual_keys=True, prefix_for_stored_virtual_keys="litellm/"),
)
monkeypatch.setattr(litellm, "store_audit_logs", False)
existing_key_row: Final = LiteLLM_VerificationToken(token="hashed-old-token", key_alias=key_alias)
response: Final = GenerateKeyResponse(token_id="hashed-new-token", key="sk-new-key", key_alias=key_alias)
await KeyManagementEventHooks.async_key_rotated_hook(
data=None,
existing_key_row=existing_key_row,
response=response,
user_api_key_dict=MagicMock(),
)
expected_secret_name: Final = f"litellm/{key_alias or 'virtual-key-hashed-old-token'}"
mock_secret_manager.async_rotate_secret.assert_awaited_once_with(
current_secret_name=expected_secret_name,
new_secret_name=expected_secret_name,
new_secret_value="sk-new-key",
optional_params=None,
)
@pytest.mark.asyncio
async def test_rotate_virtual_key_when_store_virtual_keys_disabled(self):
"""Test that rotation is skipped when store_virtual_keys is False."""