diff --git a/litellm/models/issues.py b/litellm/models/issues.py index 36e7d087b6d..d250b0726b5 100644 --- a/litellm/models/issues.py +++ b/litellm/models/issues.py @@ -81,7 +81,7 @@ class IssueStore(Protocol): scope_ids: tuple[str, ...] | None = None, ) -> tuple[AlertProjectionSignal, ...]: ... - async def recover_issues(self, day: date, issue_ids: tuple[str, ...] | None = None) -> int: ... + async def recover_issues(self, day: date, issue_ids: tuple[str, ...]) -> int: ... async def active_provider_issues(self, limit: int = 500, cursor: str | None = None) -> tuple[Alert, ...]: ... diff --git a/litellm/repositories/issues.py b/litellm/repositories/issues.py index a902ffefbe3..afe80da1fbe 100644 --- a/litellm/repositories/issues.py +++ b/litellm/repositories/issues.py @@ -392,15 +392,14 @@ class IssuesRepository: ) ) - async def recover_issues(self, day: date, issue_ids: tuple[str, ...] | None = None) -> int: + async def recover_issues(self, day: date, issue_ids: tuple[str, ...]) -> int: if issue_ids == (): return 0 - if issue_ids is not None and len(issue_ids) > 10: + if len(issue_ids) > 10: raise ValueError("Recover at most 10 resource alerts") - selection: Final = "i.id=ANY($3::text[])" if issue_ids is not None else "$3::text[] IS NULL" rows: Final = TypeAdapter(tuple[_Count, ...]).validate_python( await self.database.query_raw( - f"""WITH primary_signals AS ( + """WITH primary_signals AS ( SELECT i.id,a.id AS source_id,a.breached,a.evaluation_day,a.evaluation_revision, max(a.evaluation_day) OVER (PARTITION BY i.id) AS latest_observation_day, r.id AS rule_id,r.record AS rule_record, @@ -419,7 +418,7 @@ class IssuesRepository: LEFT JOIN "LiteLLM_UsageAlertIncident" a ON a.id=s.source_id LEFT JOIN "LiteLLM_UsageAlertRule" r ON r.id=a.rule_id WHERE i.kind='performance' AND i.active AND i.scope_type<>'provider' AND s.source='performance' - AND ({selection}) + AND i.id=ANY($3::text[]) AND s.record->'performance'->>'scope_type'=i.scope_type AND s.record->'performance'->>'scope_id' IS NOT DISTINCT FROM i.scope_id AND ((i.family='spend' AND s.record->'performance'->>'metric'='spend') @@ -501,7 +500,9 @@ class IssuesRepository: """WITH changed AS ( UPDATE "LiteLLM_Issue" SET active=false, status=CASE WHEN status='open' THEN 'resolved' ELSE status END,updated_at=CURRENT_TIMESTAMP - WHERE kind='model_deprecation' AND active AND last_seen<$1::timestamp(3) RETURNING id + WHERE id IN (SELECT id FROM "LiteLLM_Issue" + WHERE kind='model_deprecation' AND active AND last_seen<$1::timestamp(3) + ORDER BY id LIMIT 10) AND active AND last_seen<$1::timestamp(3) RETURNING id ) SELECT count(*)::int AS count FROM changed""", _utc(observed_at), ) diff --git a/litellm/repositories/usage_alerts.py b/litellm/repositories/usage_alerts.py index 664f0c6466d..90c8052c8f8 100644 --- a/litellm/repositories/usage_alerts.py +++ b/litellm/repositories/usage_alerts.py @@ -526,7 +526,7 @@ class UsageAlertsRepository: async def team_role(self, team_id: str, user_id: str) -> str | None: rows: Final = TypeAdapter(tuple[Mapping[str, str], ...]).validate_python( await self.database.query_raw( - """SELECT member->>'role' AS role FROM "LiteLLM_TeamTable", jsonb_array_elements(members_with_roles) member WHERE team_id=$1 AND member->>'user_id'=$2 LIMIT 1""", + """SELECT member->>'role' AS role FROM "LiteLLM_TeamTable", jsonb_array_elements(members_with_roles) member WHERE team_id=$1 AND member->>'user_id'=$2 ORDER BY (member->>'role'='admin') DESC NULLS LAST LIMIT 1""", team_id, user_id, ) diff --git a/tests/integration/database/test_issues.py b/tests/integration/database/test_issues.py index ac766c66ff3..6956620803a 100644 --- a/tests/integration/database/test_issues.py +++ b/tests/integration/database/test_issues.py @@ -573,7 +573,7 @@ async def test_issue_delivery_guards_snapshot_freeze_and_delayed_recovery(db: Ma update={"day": observation.day + timedelta(days=1), "breached": False, "observed": 120} ) await alerts.record_observations(enabled, (healthy,)) - await store.recover_issues(observation.day + timedelta(days=2)) + await store.recover_issues(observation.day + timedelta(days=2), (initial.id,)) recovered: Final = await store.get_issue(initial.id) assert recovered is not None and not recovered.active and recovered.status == "resolved" assert await store.find_active_issues((original.cohort_key,)) == () @@ -1057,7 +1057,7 @@ async def test_resource_recovery_ignores_broader_and_traffic_context(db: Managem ), ), ) - await store.recover_issues(issue.observed_on + timedelta(days=2)) + await store.recover_issues(issue.observed_on + timedelta(days=2), (issue.id,)) recovered: Final = await store.get_issue(issue.id) assert recovered is not None and recovered.status == "resolved" and not recovered.active supporting: Final = tuple([await alerts.get_alert(source.id) for source in sources[1:]]) @@ -1100,7 +1100,7 @@ async def test_provider_recovery_pages_saved_cohorts_and_rejects_stale_observati resolved: Final = await store.get_issue(inputs[1].id) assert dismissed is not None and not dismissed.active and dismissed.status == "dismissed" assert resolved is not None and not resolved.active and resolved.status == "resolved" - await store.recover_issues(template.observed_on + timedelta(days=1)) + await store.recover_issues(template.observed_on + timedelta(days=1), (inputs[2].id,)) remaining: Final = await store.active_provider_issues(1, prefix) assert remaining[0].id == inputs[2].id finally: @@ -1400,7 +1400,7 @@ async def test_disabled_later_breach_requires_new_healthy_primary_evidence(db: M ), ) assert await alerts.save_rule(definition.model_copy(update={"enabled": False}), "qa-admin", second) - await store.recover_issues(issue.observed_on + timedelta(days=3)) + await store.recover_issues(issue.observed_on + timedelta(days=3), (issue.id,)) unresolved: Final = await store.get_issue(issue.id) assert unresolved is not None and unresolved.active await alerts.record_observations( @@ -1415,7 +1415,7 @@ async def test_disabled_later_breach_requires_new_healthy_primary_evidence(db: M ), ), ) - await store.recover_issues(issue.observed_on + timedelta(days=3)) + await store.recover_issues(issue.observed_on + timedelta(days=3), (issue.id,)) recovered: Final = await store.get_issue(issue.id) assert recovered is not None and not recovered.active finally: @@ -1442,3 +1442,37 @@ async def test_family_settings_expose_projection_backlog_without_delayed_rules(d WHERE id='performance-dirty:2026-11-04' RETURNING id""" ) assert not (await store.family_settings()).projection_pending + + +@pytest.mark.asyncio(loop_scope="session") +async def test_retirement_recovery_is_batched_and_resumable(db: ManagementDatabase) -> None: + store: Final = IssuesRepository.from_prisma(db.prisma) + template: Final = _issue(AlertCohort(scope_type="model", scope_id="example-model")) + retirement: Final = template.observed_on + timedelta(days=10) + details: Final = AlertDeprecation( + model=ModelDeprecationInfo( + model_name="example-model", deprecation_date=retirement, days_until_deprecation=10, status="imminent" + ), + usage_model="example-model", + requests=10, + last_used_date=template.observed_on, + last_day_requests=10, + state=DeprecationAlertState(model="example-model", deprecation_date=retirement), + ) + inputs: Final = tuple( + template.model_copy(update={ + "id": str(uuid4()), "cohort_key": str(uuid4()), "kind": "model_deprecation", + "family": "retirement", "deprecation": details, "grouping_reason": "model_deprecation", + "comparisons": (), + }) + for _ in range(23) + ) + try: + await store.upsert_issues(inputs) + observed_at: Final = template.observed_at + timedelta(seconds=1) + counts: Final = tuple([await store.finish_deprecation_projection(observed_at) for _ in range(4)]) + assert counts == (10, 10, 3, 0) + recovered: Final = tuple([await store.get_issue(item.id) for item in inputs]) + assert all(item is not None and not item.active and item.status == "resolved" for item in recovered) + finally: + await _remove_issues(db, tuple(item.id for item in inputs)) diff --git a/tests/integration/database/test_usage_alerts.py b/tests/integration/database/test_usage_alerts.py index b477f07d447..fbcb89200d4 100644 --- a/tests/integration/database/test_usage_alerts.py +++ b/tests/integration/database/test_usage_alerts.py @@ -536,3 +536,24 @@ async def test_reported_model_search_preserves_exact_identity_and_cursor( assert gateway.request("GET", "/alerts/models", key=key).status_code == 403 finally: await db.prisma.writer_db.execute_raw('DELETE FROM "LiteLLM_DailyUserSpend" WHERE api_key=$1', prefix) + + +@pytest.mark.asyncio(loop_scope="session") +@pytest.mark.parametrize("roles", (("user", "admin"), ("admin", "user"), ("user",))) +async def test_duplicate_team_membership_preserves_admin_access( + db: ManagementDatabase, roles: tuple[str, ...] +) -> None: + import json + + repository: Final = UsageAlertsRepository.from_prisma(db.prisma) + identity: Final = "duplicate-role-" + uuid4().hex + await repository.database.query_raw( + 'INSERT INTO "LiteLLM_TeamTable" (team_id,members_with_roles) VALUES ($1,$2::jsonb)', + identity, + json.dumps(tuple({"user_id": identity, "role": role} for role in roles)), + ) + try: + assert await repository.team_role(identity, identity) == ("admin" if "admin" in roles else "user") + assert await repository.team_role(identity, "outsider") is None + finally: + await repository.database.query_raw('DELETE FROM "LiteLLM_TeamTable" WHERE team_id=$1', identity)