diff --git a/litellm/repositories/usage_alerts.py b/litellm/repositories/usage_alerts.py index 130f9fb0549..e9dbf977a8a 100644 --- a/litellm/repositories/usage_alerts.py +++ b/litellm/repositories/usage_alerts.py @@ -3,7 +3,7 @@ import json from collections.abc import Mapping from datetime import date, datetime, timedelta, timezone -from typing import ( # PrismaWrapper forwards generated client methods dynamically +from typing import ( TYPE_CHECKING, Final, Protocol, @@ -92,7 +92,20 @@ def activity_path( def _observation_payload(rule: UsageAlertRule, item: UsageAlertObservation, now: datetime) -> Mapping[str, object]: alert_id: Final = str( - uuid5(NAMESPACE_URL, f"litellm:usage-alert:{rule.id}:{item.entity_id or ''}:{item.day.isoformat()}") + uuid5( + NAMESPACE_URL, + json.dumps( + [ + "litellm:usage-alert", + rule.id, + rule.scope_type, + item.entity_id, + rule.metric, + rule.direction, + item.day.isoformat(), + ] + ), + ) ) incident: Final = UsageAlert( id=alert_id, @@ -329,7 +342,7 @@ class UsageAlertsRepository: async def list_breached_entity_ids(self, rule_id: str, limit: int = 500, offset: int = 0) -> tuple[str | None, ...]: rows: Final = TypeAdapter(tuple[Mapping[str, str | None], ...]).validate_python( await self.database.query_raw( - 'SELECT DISTINCT scope_id FROM "LiteLLM_UsageAlertIncident" WHERE rule_id=$1 AND breached ORDER BY scope_id NULLS FIRST LIMIT $2 OFFSET $3', + """SELECT DISTINCT i.scope_id FROM "LiteLLM_UsageAlertIncident" i JOIN "LiteLLM_UsageAlertRule" r ON r.id=i.rule_id WHERE i.rule_id=$1 AND i.breached AND i.scope_type=r.scope_type AND i.metric=r.record->>'metric' AND i.record->>'direction'=r.record->>'direction' AND (r.scope_id IS NULL OR i.scope_id=r.scope_id) ORDER BY i.scope_id NULLS FIRST LIMIT $2 OFFSET $3""", rule_id, limit, offset, @@ -366,9 +379,12 @@ class UsageAlertsRepository: return changed[0] rows: Final = TypeAdapter(tuple[_Record, ...]).validate_python( await self.database.query_raw( - 'SELECT record FROM "LiteLLM_UsageAlertIncident" WHERE rule_id=$1 AND scope_key=$2 ORDER BY evaluation_day DESC,id DESC LIMIT 1', + """SELECT record FROM "LiteLLM_UsageAlertIncident" WHERE rule_id=$1 AND scope_key=$2 AND scope_type=$3 AND metric=$4 AND record->>'direction'=$5 ORDER BY evaluation_day DESC,id DESC LIMIT 1""", rule.id, entity_id or "", + rule.scope_type, + rule.metric, + rule.direction, ) ) return UsageAlert.model_validate(rows[0].record) if rows else None @@ -389,7 +405,10 @@ class UsageAlertsRepository: prior AS ( SELECT x.*,p.id AS prior_id,p.evaluation_day,p.breached AS prior_breached,p.record AS prior_record FROM input x LEFT JOIN LATERAL ( - SELECT * FROM "LiteLLM_UsageAlertIncident" WHERE rule_id=$1 AND scope_key=x.scope_key ORDER BY evaluation_day DESC,id DESC LIMIT 1 + SELECT * FROM "LiteLLM_UsageAlertIncident" WHERE rule_id=$1 AND scope_key=x.scope_key + AND scope_type=x.record->>'scope_type' AND scope_id IS NOT DISTINCT FROM x.record->>'scope_id' + AND metric=x.record->>'metric' AND record->>'direction'=x.record->>'direction' + ORDER BY evaluation_day DESC,id DESC LIMIT 1 ) p ON true WHERE EXISTS (SELECT 1 FROM "LiteLLM_UsageAlertRule" WHERE id=$1 AND enabled AND NOT deleted AND (record->>'updated_at')::timestamptz=$3::timestamptz) ), updated AS ( diff --git a/tests/integration/database/test_usage_alerts.py b/tests/integration/database/test_usage_alerts.py index 88c7634544a..9fd94214dd1 100644 --- a/tests/integration/database/test_usage_alerts.py +++ b/tests/integration/database/test_usage_alerts.py @@ -1,8 +1,14 @@ +from datetime import date from hashlib import sha256 from typing import Final +import pytest +from integration._support.actors import ManagementDatabase from integration._support.client import Gateway, string_value +from litellm.repositories.usage_alerts import UsageAlertsRepository +from litellm.types.proxy.usage_alerts import UsageAlertObservation, UsageAlertRuleInput + def test_scoped_rules_use_real_management_auth_and_preserve_read_only_roles(gateway: Gateway) -> None: with gateway.scenario() as scenario: @@ -78,3 +84,55 @@ def test_team_admin_can_manage_team_alerts_but_member_can_only_read(gateway: Gat assert gateway.request("DELETE", path, key=member_key).status_code == 403 finally: assert gateway.request("DELETE", path, key=admin_key).status_code == 204 + + +@pytest.mark.asyncio(loop_scope="session") +@pytest.mark.parametrize( + "changed_field,new_value", + [("metric", "latency"), ("direction", "decrease"), ("scope_type", "user"), ("scope_id", "semantic-other")], +) +async def test_rule_identity_edits_preserve_old_incident_and_filter_consistency( + db: ManagementDatabase, changed_field: str, new_value: str +) -> None: + repository: Final = UsageAlertsRepository.from_prisma(db.prisma) + definition: Final = UsageAlertRuleInput( + name="Semantic identity", + scope_type="team", + scope_id="semantic-example", + metric="spend", + direction="increase", + notification_channels=("slack",), + ) + rule: Final = await repository.save_rule(definition, "qa-admin") + assert rule is not None + initial: Final = UsageAlertObservation( + entity_id=rule.scope_id, day=date(2026, 11, 4), observed=40, baseline=5, threshold=15, breached=True + ) + try: + original: Final = (await repository.record_observations(rule, (initial,)))[0] + edited: Final = await repository.save_rule( + UsageAlertRuleInput.model_validate({**definition.model_dump(), changed_field: new_value}), "qa-admin", rule + ) + assert edited is not None + observation: Final = initial.model_copy(update={"entity_id": edited.scope_id}) + replacement: Final = (await repository.record_observations(edited, (observation,)))[0] + assert replacement.id != original.id + assert await repository.get_alert(original.id) == original + before_page: Final = await repository.list_alerts( + scope_type=original.scope_type, scope_id=original.scope_id, metric=original.metric + ) + after_page: Final = await repository.list_alerts( + scope_type=replacement.scope_type, scope_id=replacement.scope_id, metric=replacement.metric + ) + assert original.id in {item.id for item in before_page.data} + assert replacement.id in {item.id for item in after_page.data} + assert all( + item.scope_type == original.scope_type and item.metric == original.metric for item in before_page.data + ) + healthy: Final = observation.model_copy(update={"day": date(2026, 11, 5), "breached": False}) + await repository.record_observations(edited, (healthy,)) + assert (await repository.get_alert(original.id)) == original + resolved: Final = await repository.get_alert(replacement.id) + assert resolved is not None and resolved.status == "resolved" + finally: + await repository.delete_rule(rule.id)