diff --git a/litellm/repositories/usage_alerts.py b/litellm/repositories/usage_alerts.py index 3237ca3a06c..6f38af2a862 100644 --- a/litellm/repositories/usage_alerts.py +++ b/litellm/repositories/usage_alerts.py @@ -447,12 +447,18 @@ class UsageAlertsRepository: SELECT DISTINCT scope_type,scope_id,key_id FROM jsonb_to_recordset($1::jsonb) AS x(scope_type text,scope_id text,key_id text) ) SELECT s.scope_type,s.scope_id,LEFT(CASE s.scope_type - WHEN 'key' THEN COALESCE(NULLIF(k.key_alias,''),NULLIF(k.key_name,'')) + WHEN 'key' THEN COALESCE(NULLIF(k.key_alias,''),NULLIF(k.key_name,''), + NULLIF(d.key_alias,''),NULLIF(d.key_name,'')) WHEN 'team' THEN NULLIF(t.team_alias,'') WHEN 'user' THEN COALESCE(NULLIF(u.user_alias,''),NULLIF(u.user_email,'')) WHEN 'model' THEN s.scope_id WHEN 'provider' THEN s.scope_id WHEN 'deployment' THEN 'Gateway' END,256) AS label FROM scopes s LEFT JOIN "LiteLLM_VerificationToken" k ON s.scope_type='key' AND k.token=s.key_id + LEFT JOIN LATERAL ( + SELECT key_alias,key_name FROM "LiteLLM_DeletedVerificationToken" + WHERE s.scope_type='key' AND token=s.key_id AND k.token IS NULL + ORDER BY deleted_at DESC,id DESC LIMIT 1 + ) d ON true LEFT JOIN "LiteLLM_TeamTable" t ON s.scope_type='team' AND t.team_id=s.scope_id LEFT JOIN "LiteLLM_UserTable" u ON s.scope_type='user' AND u.user_id=s.scope_id""", json.dumps( diff --git a/tests/integration/database/test_usage_alerts.py b/tests/integration/database/test_usage_alerts.py index 5339e451e51..29dc17b859e 100644 --- a/tests/integration/database/test_usage_alerts.py +++ b/tests/integration/database/test_usage_alerts.py @@ -39,6 +39,30 @@ async def test_alert_resource_labels_bound_legacy_display_values(db: ManagementD await repository.database.query_raw('DELETE FROM "LiteLLM_UserTable" WHERE user_id=$1', identity) +@pytest.mark.asyncio(loop_scope="session") +async def test_deleted_key_labels_use_the_latest_alias_and_prefer_live_keys(db: ManagementDatabase) -> None: + repository: Final = UsageAlertsRepository.from_prisma(db.prisma) + identity: Final = sha256(uuid4().bytes).hexdigest() + scope: Final = UsageAlertRuleInput( + name="Deleted key", scope_type="key", scope_id=identity, metric="spend", direction="increase" + ) + try: + await repository.database.query_raw( + """INSERT INTO "LiteLLM_DeletedVerificationToken" (id,token,key_alias,deleted_at) + VALUES ($1||'-old',$1,'Old alias',TIMESTAMP '2000-01-01'), + ($1||'-new',$1,repeat('x',512),TIMESTAMP '2000-01-02')""", + identity, + ) + assert await repository.resource_labels((scope,)) == {("key", identity): "x" * 256} + await repository.database.query_raw( + 'INSERT INTO "LiteLLM_VerificationToken" (token,key_alias) VALUES ($1,$2)', identity, "Live alias" + ) + assert await repository.resource_labels((scope,)) == {("key", identity): "Live alias"} + finally: + await repository.database.query_raw('DELETE FROM "LiteLLM_VerificationToken" WHERE token=$1', identity) + await repository.database.query_raw('DELETE FROM "LiteLLM_DeletedVerificationToken" WHERE token=$1', identity) + + @pytest.mark.parametrize("grant", ["management_routes", "/usage", "/usage/alert-rules"]) def test_scoped_rules_use_real_management_auth_and_preserve_read_only_roles(gateway: Gateway, grant: str) -> None: with gateway.scenario() as scenario: