Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 32 additions & 3 deletions litellm/proxy/management_endpoints/common_daily_activity.py
Original file line number Diff line number Diff line change
Expand Up @@ -327,14 +327,43 @@ 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:
verbose_proxy_logger.debug(
"Failed to fetch deleted key metadata for missing keys"
)
Comment on lines +346 to +365

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Broad exception hides bugs

get_api_key_metadata() swallows any exception when querying litellm_deletedverificationtoken and just returns partial metadata. This can silently mask real issues (e.g., query/schema errors) and lead to key_alias/team_id unexpectedly staying null with no actionable signal. At minimum, log the exception object (or re-raise non-“table missing” errors) so production failures don’t get silently ignored.


return result


def _adjust_dates_for_timezone(
start_date: str,
Expand Down
16 changes: 15 additions & 1 deletion litellm/proxy/management_endpoints/key_management_endpoints.py
Original file line number Diff line number Diff line change
Expand Up @@ -3313,6 +3313,20 @@ async def regenerate_key_fn(

verbose_proxy_logger.debug("key_in_db: %s", _key_in_db)

# Save the old key record to deleted table before regeneration
# This preserves key_alias and team_id metadata for historical spend records
try:
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,
)
except Exception:
verbose_proxy_logger.debug(
"Failed to persist old key record to deleted table during regeneration"
)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Regeneration can lose history

In regenerate_key_fn(), persisting the old key into the deleted table is wrapped in except Exception and failures are ignored. If _persist_deleted_verification_tokens() fails (DB outage, constraint violation, etc.), the function proceeds to overwrite the token hash anyway, permanently losing the old hash→metadata mapping that this PR is trying to preserve. This needs to fail the regeneration (or otherwise guarantee persistence) when the persistence step can’t be completed.


new_token = get_new_token(data=data)

new_token_hash = hash_token(new_token)
Expand Down Expand Up @@ -3749,7 +3763,7 @@ async def list_keys(
else:
admin_team_ids = None

if not user_id and user_api_key_dict.user_role not in [
if user_id is None and user_api_key_dict.user_role not in [

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Unrelated to PR scope

This fix (user_id is None vs not user_id) is correct but unrelated to preserving key metadata after deletion/regeneration.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

does not change the flow.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Where did this change come from? looks like you're about to cause a regression: #20623

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

fixed it

LitellmUserRoles.PROXY_ADMIN.value,
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value,
]:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
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,
)
Expand Down Expand Up @@ -208,3 +209,256 @@ def __init__(self, date, endpoint, api_key, model, spend, prompt_tokens, complet
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
Loading