fix(alerts): preserve display labels for deleted keys

This commit is contained in:
moe-berri 2026-11-08 15:14:03 +08:00
parent fba137954d
commit dfa6f6dd81
2 changed files with 31 additions and 1 deletions

View file

@ -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(

View file

@ -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: