tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-25 06:21:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #20481 from Harshit28j/litellm_aws_rotation_fix

Browse Source 
Fix authorization issues, same alias; verified working
This commit is contained in:
Harshit Jain
2026-02-12 09:36:10 +05:30
committed by GitHub
parent af3acdda18 ac4bd34461
commit c867740d5e
4 changed files with 270 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/hooks/key_management_event_hooks.py
+14 -3
View File
@@ -150,15 +150,26 @@ class KeyManagementEventHooks:
existing_key_row.key_alias
or f"virtual-key-{existing_key_row.token}"
)
new_secret_name = (
response.key_alias
or data.key_alias
or f"virtual-key-{response.token_id}"
)
verbose_proxy_logger.info(
"Updating secret in secret manager: secret_name=%s",
new_secret_name,
)
team_id = getattr(existing_key_row, "team_id", None)
await KeyManagementEventHooks._rotate_virtual_key_in_secret_manager(
current_secret_name=initial_secret_name,
new_secret_name=response.key_alias
or data.key_alias
or f"virtual-key-{response.token_id}",
new_secret_name=new_secret_name,
new_secret_value=response.key,
team_id=team_id,
)
verbose_proxy_logger.info(
"Secret updated in secret manager: secret_name=%s",
new_secret_name,
)
except Exception as e:
verbose_proxy_logger.warning(
f"Failed to rotate virtual key in secret manager: {e}"
litellm/proxy/management_endpoints/key_management_endpoints.py
+18 -3
View File
@@ -2770,6 +2770,7 @@ async def can_modify_verification_token(
Rules:
- Proxy admin can modify any key
- Internal jobs service account can modify any key (for auto-rotation)
- For team keys: only team admin or key owner can modify
- For personal keys: only key owner can modify
@@ -2782,13 +2783,19 @@ async def can_modify_verification_token(
Returns:
True if user can modify the key, False otherwise
"""
from litellm.constants import LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME
is_team_key = _is_team_key(data=key_info)
# 1. Proxy admin can modify any key
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
return True
# 2. For team keys: only team admin or key owner can modify
# 2. Internal jobs service account can modify any key (for auto-rotation)
if user_api_key_dict.api_key == LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME:
return True
# 3. For team keys: only team admin or key owner can modify
if is_team_key and key_info.team_id is not None:
# Get team object to check if user is team admin
team_table = await get_team_object(
@@ -2818,7 +2825,7 @@ async def can_modify_verification_token(
# Not team admin and doesn't own the key
return False
# 3. For personal keys: only key owner can modify
# 4. For personal keys: only key owner can modify
if key_info.user_id is not None and key_info.user_id == user_api_key_dict.user_id:
return True
@@ -3179,7 +3186,7 @@ def get_new_token(data: Optional[RegenerateKeyRequest]) -> str:
dependencies=[Depends(user_api_key_auth)],
)
@management_endpoint_wrapper
async def regenerate_key_fn(
async def regenerate_key_fn( # noqa: PLR0915
key: Optional[str] = None,
data: Optional[RegenerateKeyRequest] = None,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
@@ -3330,6 +3337,10 @@ async def regenerate_key_fn(
detail={"error": "You are not authorized to regenerate this key"},
)
verbose_proxy_logger.info(
"Key regeneration requested: key_alias=%s",
getattr(_key_in_db, "key_alias", None),
)
verbose_proxy_logger.debug("key_in_db: %s", _key_in_db)
new_token = get_new_token(data=data)
@@ -3380,6 +3391,10 @@ async def regenerate_key_fn(
**updated_token_dict,
)
verbose_proxy_logger.info(
"Key regeneration completed: key_alias=%s",
getattr(_key_in_db, "key_alias", None),
)
asyncio.create_task(
KeyManagementEventHooks.async_key_rotated_hook(
data=data,
litellm/secret_managers/aws_secret_manager_v2.py
+129 -32
View File
@@ -3,7 +3,8 @@ This is a file for the AWS Secret Manager Integration
Handles Async Operations for:
- Read Secret
- Write Secret
- Write Secret (CreateSecret)
- Update Secret (PutSecretValue) - for in-place rotation when alias is preserved
- Delete Secret
Relevant issue: https://github.com/BerriAI/litellm/issues/1883
@@ -42,11 +43,11 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager):
aws_profile_name: Optional[str] = None,
aws_web_identity_token: Optional[str] = None,
aws_sts_endpoint: Optional[str] = None,
**kwargs
**kwargs,
):
BaseSecretManager.__init__(self, **kwargs)
BaseAWSLLM.__init__(self, **kwargs)
# Store AWS authentication settings
self.aws_region_name = aws_region_name
self.aws_role_name = aws_role_name
@@ -61,7 +62,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager):
# AWS_REGION_NAME is only strictly required if not using a profile or role
# When using IAM roles, the region can come from multiple sources
if (
"AWS_REGION_NAME" not in os.environ
"AWS_REGION_NAME" not in os.environ
and "AWS_REGION" not in os.environ
and "AWS_DEFAULT_REGION" not in os.environ
):
@@ -83,22 +84,36 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager):
return
try:
cls.validate_environment()
# Extract AWS settings from key_management_settings if provided
aws_kwargs = {}
if key_management_settings is not None:
aws_kwargs = {
"aws_region_name": getattr(key_management_settings, "aws_region_name", None),
"aws_role_name": getattr(key_management_settings, "aws_role_name", None),
"aws_session_name": getattr(key_management_settings, "aws_session_name", None),
"aws_external_id": getattr(key_management_settings, "aws_external_id", None),
"aws_profile_name": getattr(key_management_settings, "aws_profile_name", None),
"aws_web_identity_token": getattr(key_management_settings, "aws_web_identity_token", None),
"aws_sts_endpoint": getattr(key_management_settings, "aws_sts_endpoint", None),
"aws_region_name": getattr(
key_management_settings, "aws_region_name", None
),
"aws_role_name": getattr(
key_management_settings, "aws_role_name", None
),
"aws_session_name": getattr(
key_management_settings, "aws_session_name", None
),
"aws_external_id": getattr(
key_management_settings, "aws_external_id", None
),
"aws_profile_name": getattr(
key_management_settings, "aws_profile_name", None
),
"aws_web_identity_token": getattr(
key_management_settings, "aws_web_identity_token", None
),
"aws_sts_endpoint": getattr(
key_management_settings, "aws_sts_endpoint", None
),
}
# Remove None values
aws_kwargs = {k: v for k, v in aws_kwargs.items() if v is not None}
litellm.secret_manager_client = cls(**aws_kwargs)
litellm._key_management_system = KeyManagementSystem.AWS_SECRET_MANAGER
@@ -246,13 +261,13 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager):
return primary_secret_kv_pairs.get(secret_name)
async def async_write_secret(
self,
secret_name: str,
secret_value: str,
description: Optional[str] = None,
optional_params: Optional[dict] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
tags: Optional[Union[dict, list]] = None
self,
secret_name: str,
secret_value: str,
description: Optional[str] = None,
optional_params: Optional[dict] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
tags: Optional[Union[dict, list]] = None,
) -> dict:
"""
Async function to write a secret to AWS Secrets Manager
@@ -312,6 +327,94 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager):
except httpx.TimeoutException:
raise ValueError("Timeout error occurred")
async def async_put_secret_value(
self,
secret_name: str,
secret_value: str,
optional_params: Optional[dict] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
) -> dict:
"""
Async function to update an existing secret's value in AWS Secrets Manager.
Uses PutSecretValue to update in place. Use this when rotating a secret
that keeps the same name (current_secret_name == new_secret_name).
Args:
secret_name: Name of the existing secret to update
secret_value: New value to store
optional_params: Additional AWS parameters
timeout: Request timeout
Returns:
dict: Response from AWS Secrets Manager containing update details
"""
from litellm._uuid import uuid
data: Dict[str, Any] = {
"SecretId": secret_name,
"SecretString": secret_value,
"ClientRequestToken": str(uuid.uuid4()),
}
endpoint_url, headers, body = self._prepare_request(
action="PutSecretValue",
secret_name=secret_name,
secret_value=secret_value,
optional_params=optional_params,
request_data=data,
)
async_client = get_async_httpx_client(
llm_provider=httpxSpecialProvider.SecretManager,
params={"timeout": timeout},
)
try:
response = await async_client.post(
url=endpoint_url, headers=headers, data=body.decode("utf-8")
)
response.raise_for_status()
return response.json()
except httpx.HTTPStatusError as err:
raise ValueError(f"HTTP error occurred: {err.response.text}")
except httpx.TimeoutException:
raise ValueError("Timeout error occurred")
async def async_rotate_secret(
self,
current_secret_name: str,
new_secret_name: str,
new_secret_value: str,
optional_params: Optional[dict] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
) -> dict:
"""
Rotate a secret. When current_secret_name == new_secret_name (in-place
update), uses PutSecretValue instead of create+delete to avoid
ResourceExistsException.
"""
if current_secret_name == new_secret_name:
# Same alias: update in place via PutSecretValue
verbose_logger.info(
"Secret rotated in-place (PutSecretValue): secret_name=%s",
current_secret_name,
)
return await self.async_put_secret_value(
secret_name=current_secret_name,
secret_value=new_secret_value,
optional_params=optional_params,
timeout=timeout,
)
# Different names: create new, delete old (base class logic)
return await super().async_rotate_secret(
current_secret_name=current_secret_name,
new_secret_name=new_secret_name,
new_secret_value=new_secret_value,
optional_params=optional_params,
timeout=timeout,
)
async def async_delete_secret(
self,
secret_name: str,
@@ -375,7 +478,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager):
except ImportError:
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
optional_params = optional_params or {}
# Build optional_params from instance settings if not provided
# This allows the IAM role settings to be used for Secret Manager calls
if not optional_params.get("aws_role_name") and self.aws_role_name:
@@ -388,11 +491,14 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager):
optional_params["aws_external_id"] = self.aws_external_id
if not optional_params.get("aws_profile_name") and self.aws_profile_name:
optional_params["aws_profile_name"] = self.aws_profile_name
if not optional_params.get("aws_web_identity_token") and self.aws_web_identity_token:
if (
not optional_params.get("aws_web_identity_token")
and self.aws_web_identity_token
):
optional_params["aws_web_identity_token"] = self.aws_web_identity_token
if not optional_params.get("aws_sts_endpoint") and self.aws_sts_endpoint:
optional_params["aws_sts_endpoint"] = self.aws_sts_endpoint
boto3_credentials_info = self._get_boto_credentials_from_optional_params(
optional_params
)
@@ -431,12 +537,3 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager):
prepped = request.prepare()
return endpoint_url, prepped.headers, body
# if __name__ == "__main__":
# print("loading aws secret manager v2")
# aws_secret_manager_v2 = AWSSecretsManagerV2()
# import asyncio
# print("writing secret to aws secret manager v2")
# asyncio.run(aws_secret_manager_v2.async_write_secret(secret_name="test_secret_3", secret_value="test_value_2"))
# print("reading secret from aws secret manager v2")
tests/test_litellm/secret_managers/test_aws_secret_manager_rotation.py
+109
View File
@@ -0,0 +1,109 @@
"""
Regression tests for AWS Secrets Manager same-name in-place rotation fix.
When current_secret_name == new_secret_name (e.g. key alias preserved during
rotation), AWS must use PutSecretValue to update in place instead of
create+delete, which would fail with ResourceExistsException.
"""
from unittest.mock import AsyncMock, patch
import pytest
from litellm.secret_managers.aws_secret_manager_v2 import AWSSecretsManagerV2
@pytest.mark.asyncio
async def test_rotate_secret_same_name_uses_put_secret_value():
"""
When current_secret_name == new_secret_name, async_rotate_secret should
call PutSecretValue (async_put_secret_value) instead of create+delete.
"""
secret_name = "litellm/tenant/litellm-metis-key"
new_value = "sk-new-rotated-key-value"
with patch.object(
AWSSecretsManagerV2,
"async_put_secret_value",
new_callable=AsyncMock,
return_value={"ARN": "arn:aws:secretsmanager:us-east-1:123:secret:test"},
) as mock_put:
with patch.object(
AWSSecretsManagerV2,
"async_write_secret",
new_callable=AsyncMock,
) as mock_write:
with patch.object(
AWSSecretsManagerV2,
"async_delete_secret",
new_callable=AsyncMock,
) as mock_delete:
manager = AWSSecretsManagerV2()
result = await manager.async_rotate_secret(
current_secret_name=secret_name,
new_secret_name=secret_name,
new_secret_value=new_value,
)
# PutSecretValue (in-place update) should be called
mock_put.assert_called_once_with(
secret_name=secret_name,
secret_value=new_value,
optional_params=None,
timeout=None,
)
# Create + delete should NOT be called
mock_write.assert_not_called()
mock_delete.assert_not_called()
assert result["ARN"] == "arn:aws:secretsmanager:us-east-1:123:secret:test"
@pytest.mark.asyncio
async def test_rotate_secret_different_names_uses_create_delete():
"""
When current_secret_name != new_secret_name, async_rotate_secret should
use base class logic (create new, delete old).
"""
current_name = "litellm/old-key-alias"
new_name = "litellm/virtual-key-new-token-id"
new_value = "sk-new-key-value"
with patch.object(
AWSSecretsManagerV2,
"async_read_secret",
new_callable=AsyncMock,
side_effect=["sk-old-value", new_value], # read old, then read new
):
with patch.object(
AWSSecretsManagerV2,
"async_write_secret",
new_callable=AsyncMock,
return_value={"ARN": "arn:new"},
) as mock_write:
with patch.object(
AWSSecretsManagerV2,
"async_delete_secret",
new_callable=AsyncMock,
return_value={},
) as mock_delete:
with patch.object(
AWSSecretsManagerV2,
"async_put_secret_value",
new_callable=AsyncMock,
) as mock_put:
manager = AWSSecretsManagerV2()
await manager.async_rotate_secret(
current_secret_name=current_name,
new_secret_name=new_name,
new_secret_value=new_value,
)
# PutSecretValue should NOT be called (different names)
mock_put.assert_not_called()
# Create + delete should be called
mock_write.assert_called_once()
mock_delete.assert_called_once_with(
secret_name=current_name,
recovery_window_in_days=7,
optional_params=None,
timeout=None,
)