diff --git a/litellm/models/usage_alerts.py b/litellm/models/usage_alerts.py index 8858835eafc..bf90f6acb80 100644 --- a/litellm/models/usage_alerts.py +++ b/litellm/models/usage_alerts.py @@ -25,7 +25,7 @@ class UsageAlertStore(Protocol): self, data: UsageAlertRuleInput, created_by: str | None, existing: UsageAlertRule | None = None ) -> UsageAlertRule | None: ... - async def delete_rule(self, rule_id: str) -> None: ... + async def delete_rule(self, rule_id: str, *, existing: UsageAlertRule | None = None) -> bool: ... async def list_rules( self, diff --git a/litellm/proxy/alerting_endpoints/usage_alerts.py b/litellm/proxy/alerting_endpoints/usage_alerts.py index e6b1b22cc00..d571fc776d9 100644 --- a/litellm/proxy/alerting_endpoints/usage_alerts.py +++ b/litellm/proxy/alerting_endpoints/usage_alerts.py @@ -188,7 +188,7 @@ async def update_usage_alert_rule( await require_scope(user, store, data.scope_type, data.scope_id, write=True) saved: Final = await store.save_rule(data, user.user_id, existing=rule) if saved is None: - raise HTTPException(status_code=404, detail="Alert rule not found") + raise HTTPException(status_code=409, detail="Alert rule changed while being updated. Refresh and try again") return saved @@ -199,7 +199,8 @@ async def delete_usage_alert_rule(rule_id: str, user: User, repository: Reposito if rule is None: raise HTTPException(status_code=404, detail="Alert rule not found") await require_scope(user, store, rule.scope_type, rule.scope_id, write=True) - await store.delete_rule(rule_id) + if not await store.delete_rule(rule_id, existing=rule): + raise HTTPException(status_code=409, detail="Alert rule changed while being deleted. Refresh and try again") return Response(status_code=204) diff --git a/litellm/repositories/usage_alerts.py b/litellm/repositories/usage_alerts.py index 3ecfe60d679..a7f65c535cd 100644 --- a/litellm/repositories/usage_alerts.py +++ b/litellm/repositories/usage_alerts.py @@ -223,19 +223,35 @@ class UsageAlertsRepository: rows: Final = await self.database.query_raw( """INSERT INTO "LiteLLM_UsageAlertRule" (id,scope_type,scope_id,enabled,record) VALUES ($1,$2,$3,$4,$5::jsonb) ON CONFLICT (id) DO UPDATE SET scope_type=EXCLUDED.scope_type,scope_id=EXCLUDED.scope_id,enabled=EXCLUDED.enabled,record=EXCLUDED.record - WHERE NOT "LiteLLM_UsageAlertRule".deleted RETURNING record""", + WHERE NOT "LiteLLM_UsageAlertRule".deleted + AND ("LiteLLM_UsageAlertRule".record->>'updated_at')::timestamptz=$6::text::timestamptz + AND "LiteLLM_UsageAlertRule".scope_type=$7 + AND "LiteLLM_UsageAlertRule".scope_id IS NOT DISTINCT FROM $8::text + RETURNING record""", rule.id, rule.scope_type, rule.scope_id, rule.enabled, rule.model_dump_json(), + existing.updated_at.isoformat() if existing else None, + existing.scope_type if existing else None, + existing.scope_id if existing else None, ) saved: Final = TypeAdapter(tuple[_Record, ...]).validate_python(rows) return _rule(saved[0].record) if saved else None - async def delete_rule(self, rule_id: str) -> None: - await self.database.query_raw( - 'UPDATE "LiteLLM_UsageAlertRule" SET deleted=true,enabled=false WHERE id=$1 RETURNING id', rule_id + async def delete_rule(self, rule_id: str, *, existing: UsageAlertRule | None = None) -> bool: + return bool( + await self.database.query_raw( + """UPDATE "LiteLLM_UsageAlertRule" SET deleted=true,enabled=false + WHERE id=$1 AND NOT deleted AND ($2::text::timestamptz IS NULL OR ( + (record->>'updated_at')::timestamptz=$2::text::timestamptz + AND scope_type=$3 AND scope_id IS NOT DISTINCT FROM $4::text)) RETURNING id""", + rule_id, + existing.updated_at.isoformat() if existing else None, + existing.scope_type if existing else None, + existing.scope_id if existing else None, + ) ) async def list_rules( diff --git a/tests/integration/database/test_usage_alerts.py b/tests/integration/database/test_usage_alerts.py index 51978152ed2..c152d6c95f0 100644 --- a/tests/integration/database/test_usage_alerts.py +++ b/tests/integration/database/test_usage_alerts.py @@ -210,6 +210,36 @@ async def test_ongoing_episode_preserves_its_original_evidence_and_activity_link await repository.delete_rule(rule.id) +@pytest.mark.asyncio(loop_scope="session") +async def test_rule_scope_changes_reject_stale_updates_and_deletes(db: ManagementDatabase) -> None: + repository: Final = UsageAlertsRepository.from_prisma(db.prisma) + definition: Final = UsageAlertRuleInput( + name="Concurrent scope edit", scope_type="user", scope_id="original-owner", metric="spend", direction="increase" + ) + original: Final = await repository.save_rule(definition, "qa-admin") + assert original is not None + try: + moved: Final = await repository.save_rule( + definition.model_copy(update={"scope_id": "different-owner"}), "qa-admin", existing=original + ) + assert moved is not None + assert await repository.save_rule(definition, "original-owner", existing=original) is None + assert not await repository.delete_rule(original.id, existing=original) + assert await repository.get_rule(original.id) == moved + current: Final = await repository.save_rule( + definition.model_copy(update={"scope_id": moved.scope_id, "name": "Current authorized edit"}), + "qa-admin", + existing=moved, + ) + assert current is not None + assert not await repository.delete_rule(original.id, existing=moved) + assert await repository.delete_rule(original.id, existing=current) + assert await repository.save_rule(definition, "qa-admin", existing=current) is None + assert await repository.get_rule(original.id) is None + finally: + await repository.delete_rule(original.id) + + @pytest.mark.asyncio(loop_scope="session") @pytest.mark.parametrize( "changed_field,new_value",