mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-13 01:04:54 +00:00
Merge pull request #20684 from BerriAI/litellm_preserver_team_key_alias_after_key_regeneration_and_deletion
[Fix] Preserve key_alias and team_id metadata in /user/daily/activity/aggregated after key deletion or regeneration
This commit is contained in:
Shivam Rawat
@@ -327,14 +327,45 @@ async def get_api_key_metadata(
|
||||
prisma_client: PrismaClient,
|
||||
api_keys: Set[str],
|
||||
) -> Dict[str, Dict[str, Any]]:
|
||||
"""Update api key metadata for a single record."""
|
||||
"""Get api key metadata, falling back to deleted keys table for keys not found in active table.
|
||||
|
||||
This ensures that key_alias and team_id are preserved in historical activity logs
|
||||
even after a key is deleted or regenerated.
|
||||
"""
|
||||
key_records = await prisma_client.db.litellm_verificationtoken.find_many(
|
||||
where={"token": {"in": list(api_keys)}}
|
||||
)
|
||||
return {
|
||||
k.token: {"key_alias": k.key_alias, "team_id": k.team_id} for k in key_records
|
||||
result = {
|
||||
k.token: {"key_alias": k.key_alias, "team_id": k.team_id}
|
||||
for k in key_records
|
||||
}
|
||||
|
||||
# For any keys not found in the active table, check the deleted keys table
|
||||
missing_keys = api_keys - set(result.keys())
|
||||
if missing_keys:
|
||||
try:
|
||||
deleted_key_records = (
|
||||
await prisma_client.db.litellm_deletedverificationtoken.find_many(
|
||||
where={"token": {"in": list(missing_keys)}},
|
||||
order={"deleted_at": "desc"},
|
||||
)
|
||||
)
|
||||
# Use the most recent deleted record for each token (ordered by deleted_at desc)
|
||||
for k in deleted_key_records:
|
||||
if k.token not in result:
|
||||
result[k.token] = {
|
||||
"key_alias": k.key_alias,
|
||||
"team_id": k.team_id,
|
||||
}
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.warning(
|
||||
"Failed to fetch deleted key metadata for %d missing keys: %s",
|
||||
len(missing_keys),
|
||||
e,
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def _adjust_dates_for_timezone(
|
||||
start_date: str,
|
||||
|
||||
@@ -68,6 +68,7 @@ from litellm.proxy.management_helpers.utils import management_endpoint_wrapper
|
||||
from litellm.proxy.spend_tracking.spend_tracking_utils import _is_master_key
|
||||
from litellm.proxy.utils import (
|
||||
PrismaClient,
|
||||
ProxyLogging,
|
||||
_hash_token_if_needed,
|
||||
handle_exception_on_proxy,
|
||||
is_valid_api_key,
|
||||
@@ -3180,6 +3181,63 @@ def get_new_token(data: Optional[RegenerateKeyRequest]) -> str:
|
||||
return new_token
|
||||
|
||||
|
||||
async def _execute_virtual_key_regeneration(
|
||||
*,
|
||||
prisma_client: PrismaClient,
|
||||
key_in_db: LiteLLM_VerificationToken,
|
||||
hashed_api_key: str,
|
||||
key: str,
|
||||
data: Optional[RegenerateKeyRequest],
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
litellm_changed_by: Optional[str],
|
||||
user_api_key_cache: DualCache,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
) -> GenerateKeyResponse:
|
||||
"""Generate new token, update DB, invalidate cache, and return response."""
|
||||
from litellm.proxy.proxy_server import hash_token
|
||||
|
||||
new_token = get_new_token(data=data)
|
||||
new_token_hash = hash_token(new_token)
|
||||
new_token_key_name = f"sk-...{new_token[-4:]}"
|
||||
update_data = {"token": new_token_hash, "key_name": new_token_key_name}
|
||||
|
||||
non_default_values = {}
|
||||
if data is not None:
|
||||
non_default_values = await prepare_key_update_data(
|
||||
data=data, existing_key_row=key_in_db
|
||||
)
|
||||
verbose_proxy_logger.debug("non_default_values: %s", non_default_values)
|
||||
update_data.update(non_default_values)
|
||||
update_data = prisma_client.jsonify_object(data=update_data)
|
||||
|
||||
updated_token = await prisma_client.db.litellm_verificationtoken.update(
|
||||
where={"token": hashed_api_key},
|
||||
data=update_data, # type: ignore
|
||||
)
|
||||
updated_token_dict = dict(updated_token) if updated_token is not None else {}
|
||||
updated_token_dict["key"] = new_token
|
||||
updated_token_dict["token_id"] = updated_token_dict.pop("token")
|
||||
|
||||
if hashed_api_key or key:
|
||||
await _delete_cache_key_object(
|
||||
hashed_token=hash_token(key),
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
response = GenerateKeyResponse(**updated_token_dict)
|
||||
asyncio.create_task(
|
||||
KeyManagementEventHooks.async_key_rotated_hook(
|
||||
data=data,
|
||||
existing_key_row=key_in_db,
|
||||
response=response,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_changed_by=litellm_changed_by,
|
||||
)
|
||||
)
|
||||
return response
|
||||
|
||||
|
||||
@router.post(
|
||||
"/key/{key:path}/regenerate",
|
||||
tags=["key management"],
|
||||
@@ -3348,69 +3406,32 @@ async def regenerate_key_fn( # noqa: PLR0915
|
||||
)
|
||||
verbose_proxy_logger.debug("key_in_db: %s", _key_in_db)
|
||||
|
||||
new_token = get_new_token(data=data)
|
||||
# Normalize litellm_changed_by: if it's a Header object or not a string, convert to None
|
||||
if litellm_changed_by is not None and not isinstance(litellm_changed_by, str):
|
||||
litellm_changed_by = None
|
||||
|
||||
new_token_hash = hash_token(new_token)
|
||||
new_token_key_name = f"sk-...{new_token[-4:]}"
|
||||
|
||||
# Prepare the update data
|
||||
update_data = {
|
||||
"token": new_token_hash,
|
||||
"key_name": new_token_key_name,
|
||||
}
|
||||
|
||||
non_default_values = {}
|
||||
if data is not None:
|
||||
# Update with any provided parameters from GenerateKeyRequest
|
||||
non_default_values = await prepare_key_update_data(
|
||||
data=data, existing_key_row=_key_in_db
|
||||
)
|
||||
verbose_proxy_logger.debug("non_default_values: %s", non_default_values)
|
||||
|
||||
update_data.update(non_default_values)
|
||||
update_data = prisma_client.jsonify_object(data=update_data)
|
||||
# Update the token in the database
|
||||
updated_token = await prisma_client.db.litellm_verificationtoken.update(
|
||||
where={"token": hashed_api_key},
|
||||
data=update_data, # type: ignore
|
||||
# Save the old key record to deleted table before regeneration.
|
||||
# This preserves key_alias and team_id metadata for historical spend records.
|
||||
# If this fails, abort the regeneration to avoid permanently losing the
|
||||
# old hash→metadata mapping.
|
||||
await _persist_deleted_verification_tokens(
|
||||
keys=[_key_in_db],
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_changed_by=litellm_changed_by,
|
||||
)
|
||||
|
||||
updated_token_dict = {}
|
||||
if updated_token is not None:
|
||||
updated_token_dict = dict(updated_token)
|
||||
|
||||
updated_token_dict["key"] = new_token
|
||||
updated_token_dict["token_id"] = updated_token_dict.pop("token")
|
||||
|
||||
### 3. remove existing key entry from cache
|
||||
######################################################################
|
||||
|
||||
if hashed_api_key or key:
|
||||
await _delete_cache_key_object(
|
||||
hashed_token=hash_token(key),
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
response = GenerateKeyResponse(
|
||||
**updated_token_dict,
|
||||
return await _execute_virtual_key_regeneration(
|
||||
prisma_client=prisma_client,
|
||||
key_in_db=_key_in_db,
|
||||
hashed_api_key=hashed_api_key,
|
||||
key=key,
|
||||
data=data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_changed_by=litellm_changed_by,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
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,
|
||||
existing_key_row=_key_in_db,
|
||||
response=response,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_changed_by=litellm_changed_by,
|
||||
)
|
||||
)
|
||||
|
||||
return response
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception("Error regenerating key: %s", e)
|
||||
raise handle_exception_on_proxy(e)
|
||||
|
||||
@@ -11,6 +11,7 @@ sys.path.insert(
|
||||
from litellm.proxy.management_endpoints.common_daily_activity import (
|
||||
_is_user_agent_tag,
|
||||
compute_tag_metadata_totals,
|
||||
get_api_key_metadata,
|
||||
get_daily_activity,
|
||||
get_daily_activity_aggregated,
|
||||
)
|
||||
@@ -208,3 +209,256 @@ async def test_get_daily_activity_aggregated_with_endpoint_breakdown():
|
||||
assert chat_endpoint.api_key_breakdown["key-1"].metrics.spend == 15.0
|
||||
assert "key-2" in embeddings_endpoint.api_key_breakdown
|
||||
assert embeddings_endpoint.api_key_breakdown["key-2"].metrics.spend == 3.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_api_key_metadata_returns_active_key_metadata():
|
||||
"""Test that get_api_key_metadata should return metadata for active keys."""
|
||||
mock_prisma = MagicMock()
|
||||
|
||||
# Mock active key record
|
||||
mock_active_key = MagicMock()
|
||||
mock_active_key.token = "active-key-hash-123"
|
||||
mock_active_key.key_alias = "my-active-key"
|
||||
mock_active_key.team_id = "team-abc"
|
||||
|
||||
mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(
|
||||
return_value=[mock_active_key]
|
||||
)
|
||||
|
||||
result = await get_api_key_metadata(
|
||||
prisma_client=mock_prisma,
|
||||
api_keys={"active-key-hash-123"},
|
||||
)
|
||||
|
||||
assert "active-key-hash-123" in result
|
||||
assert result["active-key-hash-123"]["key_alias"] == "my-active-key"
|
||||
assert result["active-key-hash-123"]["team_id"] == "team-abc"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_api_key_metadata_falls_back_to_deleted_keys():
|
||||
"""Test that get_api_key_metadata should fall back to deleted keys table for missing keys."""
|
||||
mock_prisma = MagicMock()
|
||||
|
||||
# No active keys found
|
||||
mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
|
||||
|
||||
# Deleted key record exists
|
||||
mock_deleted_key = MagicMock()
|
||||
mock_deleted_key.token = "deleted-key-hash-456"
|
||||
mock_deleted_key.key_alias = "toto-test-2"
|
||||
mock_deleted_key.team_id = "team-xyz"
|
||||
|
||||
mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(
|
||||
return_value=[mock_deleted_key]
|
||||
)
|
||||
|
||||
result = await get_api_key_metadata(
|
||||
prisma_client=mock_prisma,
|
||||
api_keys={"deleted-key-hash-456"},
|
||||
)
|
||||
|
||||
assert "deleted-key-hash-456" in result
|
||||
assert result["deleted-key-hash-456"]["key_alias"] == "toto-test-2"
|
||||
assert result["deleted-key-hash-456"]["team_id"] == "team-xyz"
|
||||
|
||||
# Verify deleted table was queried with the missing key
|
||||
mock_prisma.db.litellm_deletedverificationtoken.find_many.assert_called_once_with(
|
||||
where={"token": {"in": ["deleted-key-hash-456"]}},
|
||||
order={"deleted_at": "desc"},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_api_key_metadata_mixed_active_and_deleted_keys():
|
||||
"""Test that get_api_key_metadata should return metadata for both active and deleted keys."""
|
||||
mock_prisma = MagicMock()
|
||||
|
||||
# One active key found
|
||||
mock_active_key = MagicMock()
|
||||
mock_active_key.token = "active-key-hash"
|
||||
mock_active_key.key_alias = "active-alias"
|
||||
mock_active_key.team_id = "team-active"
|
||||
|
||||
mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(
|
||||
return_value=[mock_active_key]
|
||||
)
|
||||
|
||||
# One deleted key found
|
||||
mock_deleted_key = MagicMock()
|
||||
mock_deleted_key.token = "deleted-key-hash"
|
||||
mock_deleted_key.key_alias = "deleted-alias"
|
||||
mock_deleted_key.team_id = "team-deleted"
|
||||
|
||||
mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(
|
||||
return_value=[mock_deleted_key]
|
||||
)
|
||||
|
||||
result = await get_api_key_metadata(
|
||||
prisma_client=mock_prisma,
|
||||
api_keys={"active-key-hash", "deleted-key-hash"},
|
||||
)
|
||||
|
||||
# Both keys should have metadata
|
||||
assert len(result) == 2
|
||||
assert result["active-key-hash"]["key_alias"] == "active-alias"
|
||||
assert result["active-key-hash"]["team_id"] == "team-active"
|
||||
assert result["deleted-key-hash"]["key_alias"] == "deleted-alias"
|
||||
assert result["deleted-key-hash"]["team_id"] == "team-deleted"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_api_key_metadata_deleted_table_not_queried_when_all_keys_found():
|
||||
"""Test that get_api_key_metadata should not query deleted table when all keys are active."""
|
||||
mock_prisma = MagicMock()
|
||||
|
||||
mock_active_key = MagicMock()
|
||||
mock_active_key.token = "key-hash-1"
|
||||
mock_active_key.key_alias = "alias-1"
|
||||
mock_active_key.team_id = "team-1"
|
||||
|
||||
mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(
|
||||
return_value=[mock_active_key]
|
||||
)
|
||||
mock_prisma.db.litellm_deletedverificationtoken = MagicMock()
|
||||
mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(
|
||||
return_value=[]
|
||||
)
|
||||
|
||||
result = await get_api_key_metadata(
|
||||
prisma_client=mock_prisma,
|
||||
api_keys={"key-hash-1"},
|
||||
)
|
||||
|
||||
assert len(result) == 1
|
||||
assert result["key-hash-1"]["key_alias"] == "alias-1"
|
||||
# Deleted table should NOT have been queried
|
||||
mock_prisma.db.litellm_deletedverificationtoken.find_many.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_api_key_metadata_deleted_table_error_handled_gracefully():
|
||||
"""Test that get_api_key_metadata should handle errors from deleted table gracefully."""
|
||||
mock_prisma = MagicMock()
|
||||
|
||||
# No active keys found
|
||||
mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
|
||||
|
||||
# Deleted table raises an error (e.g., table doesn't exist in older schema)
|
||||
mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(
|
||||
side_effect=Exception("Table not found")
|
||||
)
|
||||
|
||||
result = await get_api_key_metadata(
|
||||
prisma_client=mock_prisma,
|
||||
api_keys={"missing-key-hash"},
|
||||
)
|
||||
|
||||
# Should return empty dict without raising
|
||||
assert result == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_api_key_metadata_regenerated_key_uses_most_recent_deleted_record():
|
||||
"""Test that get_api_key_metadata should use the most recent deleted record for regenerated keys."""
|
||||
mock_prisma = MagicMock()
|
||||
|
||||
# No active keys found (old hash no longer in active table after regeneration)
|
||||
mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
|
||||
|
||||
# Multiple deleted records for same token (e.g., regenerated multiple times)
|
||||
mock_deleted_1 = MagicMock()
|
||||
mock_deleted_1.token = "old-key-hash"
|
||||
mock_deleted_1.key_alias = "latest-alias"
|
||||
mock_deleted_1.team_id = "latest-team"
|
||||
|
||||
mock_deleted_2 = MagicMock()
|
||||
mock_deleted_2.token = "old-key-hash"
|
||||
mock_deleted_2.key_alias = "older-alias"
|
||||
mock_deleted_2.team_id = "older-team"
|
||||
|
||||
# Ordered by deleted_at desc, so first record is the most recent
|
||||
mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(
|
||||
return_value=[mock_deleted_1, mock_deleted_2]
|
||||
)
|
||||
|
||||
result = await get_api_key_metadata(
|
||||
prisma_client=mock_prisma,
|
||||
api_keys={"old-key-hash"},
|
||||
)
|
||||
|
||||
# Should use the first (most recent) record
|
||||
assert result["old-key-hash"]["key_alias"] == "latest-alias"
|
||||
assert result["old-key-hash"]["team_id"] == "latest-team"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aggregated_activity_preserves_metadata_for_deleted_keys():
|
||||
"""Test that the full aggregation pipeline should preserve metadata for deleted keys."""
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db = MagicMock()
|
||||
|
||||
class MockRecord:
|
||||
def __init__(self, date, endpoint, api_key, model, spend, prompt_tokens, completion_tokens):
|
||||
self.date = date
|
||||
self.endpoint = endpoint
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.model_group = None
|
||||
self.custom_llm_provider = "openai"
|
||||
self.mcp_namespaced_tool_name = None
|
||||
self.spend = spend
|
||||
self.prompt_tokens = prompt_tokens
|
||||
self.completion_tokens = completion_tokens
|
||||
self.total_tokens = prompt_tokens + completion_tokens
|
||||
self.cache_read_input_tokens = 0
|
||||
self.cache_creation_input_tokens = 0
|
||||
self.api_requests = 1
|
||||
self.successful_requests = 1
|
||||
self.failed_requests = 0
|
||||
|
||||
# Records reference a deleted key
|
||||
mock_records = [
|
||||
MockRecord("2024-01-01", "/v1/chat/completions", "deleted-key-hash", "gpt-4", 10.0, 100, 50),
|
||||
]
|
||||
|
||||
mock_table = MagicMock()
|
||||
mock_table.find_many = AsyncMock(return_value=mock_records)
|
||||
mock_prisma.db.litellm_dailyuserspend = mock_table
|
||||
|
||||
# Active table returns nothing for this key
|
||||
mock_prisma.db.litellm_verificationtoken = MagicMock()
|
||||
mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
|
||||
|
||||
# Deleted table returns the metadata
|
||||
mock_deleted_key = MagicMock()
|
||||
mock_deleted_key.token = "deleted-key-hash"
|
||||
mock_deleted_key.key_alias = "toto-test-2"
|
||||
mock_deleted_key.team_id = "69cd4b77-b095-4489-8c46-4f2f31d840a2"
|
||||
|
||||
mock_prisma.db.litellm_deletedverificationtoken = MagicMock()
|
||||
mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(
|
||||
return_value=[mock_deleted_key]
|
||||
)
|
||||
|
||||
result = await get_daily_activity_aggregated(
|
||||
prisma_client=mock_prisma,
|
||||
table_name="litellm_dailyuserspend",
|
||||
entity_id_field="user_id",
|
||||
entity_id=None,
|
||||
entity_metadata_field=None,
|
||||
start_date="2024-01-01",
|
||||
end_date="2024-01-01",
|
||||
model=None,
|
||||
api_key=None,
|
||||
)
|
||||
|
||||
# Verify the deleted key's metadata is preserved
|
||||
daily_data = result.results[0]
|
||||
chat_endpoint = daily_data.breakdown.endpoints["/v1/chat/completions"]
|
||||
assert "deleted-key-hash" in chat_endpoint.api_key_breakdown
|
||||
key_data = chat_endpoint.api_key_breakdown["deleted-key-hash"]
|
||||
assert key_data.metadata.key_alias == "toto-test-2"
|
||||
assert key_data.metadata.team_id == "69cd4b77-b095-4489-8c46-4f2f31d840a2"
|
||||
assert key_data.metrics.spend == 10.0
|
||||
|
||||
Reference in New Issue
Block a user