fix(alerts): preserve scoped authorization and policy ownership

This commit is contained in:
moe-berri 2026-11-06 16:50:07 +08:00
parent 68a1ec4909
commit 08760d243b
10 changed files with 308 additions and 27 deletions

View file

@ -53,8 +53,11 @@ class IssueStore(Protocol):
status: UsageAlertStatus,
expected_status: UsageAlertStatus,
actor: str | None = None,
authorized_scopes: tuple[IssueCohort, ...] | None = None,
) -> Issue | None: ...
async def status_scopes(self, issue_id: str) -> tuple[IssueCohort, ...]: ...
async def upsert_issue(
self,
data: IssueInput,

View file

@ -25,6 +25,7 @@ from litellm.types.proxy.auth.user_roles import LitellmUserRoles
from litellm.types.proxy.issues import (
Issue,
IssueActivityResponse,
IssueCohort,
IssueFamily,
IssueFamilyList,
IssueFamilySetting,
@ -106,6 +107,29 @@ async def require_activity_access(issue: Issue, user: User, repository: IssuesRe
raise HTTPException(status_code=403, detail="Full team Activity permission is required for this investigation")
async def require_status_scopes(
issue: Issue, user: User, repository: IssueStore, alerts: UsageAlertsRepositoryDep
) -> tuple[IssueCohort, ...] | None:
store: Final = require_repository(alerts)
await require_scope(
user,
store,
"model" if issue.cohort.scope_type == "provider" else issue.cohort.scope_type,
issue.cohort.scope_id,
write=True,
)
if user.user_role == LitellmUserRoles.PROXY_ADMIN:
return None
scopes: Final = await repository.status_scopes(issue.id)
if len(scopes) > 10:
raise HTTPException(status_code=403, detail="An administrator must update this alert's status")
for scope in scopes:
await require_scope(
user, store, "model" if scope.scope_type == "provider" else scope.scope_type, scope.scope_id, write=True
)
return scopes
def activity_range(issue: Issue, start_date: date | None, end_date: date | None) -> tuple[date, date]:
start: Final = start_date or issue.range_start
end: Final = end_date or issue.range_end
@ -175,13 +199,7 @@ async def get_issue(
activity_allowed: Final = await can_view_activity(issue, user, require_issues(repository))
try:
require_write_access(user, f"/issues/{issue_id}")
await require_scope(
user,
require_repository(alerts),
"model" if issue.cohort.scope_type == "provider" else issue.cohort.scope_type,
issue.cohort.scope_id,
write=True,
)
await require_status_scopes(issue, user, require_issues(repository), alerts)
except HTTPException as exc:
if exc.status_code != 403:
raise
@ -198,11 +216,12 @@ async def update_issue(
alerts: UsageAlertsRepositoryDep,
) -> Issue:
require_write_access(user, f"/issues/{issue_id}")
issue: Final = await authorized_issue(issue_id, user, repository, alerts, write=True)
issue: Final = await authorized_issue(issue_id, user, repository, alerts)
scopes: Final = await require_status_scopes(issue, user, require_issues(repository), alerts)
if issue.kind == "model_deprecation" and not issue.active and data.status == "open":
raise HTTPException(status_code=409, detail="This retirement issue reopens when the model receives new traffic")
changed: Final = await require_issues(repository).set_status(
issue_id, data.status, data.expected_status, user.user_id
issue_id, data.status, data.expected_status, user.user_id, authorized_scopes=scopes
)
if changed is None:
raise HTTPException(status_code=409, detail="Issue status changed; reload and retry")

View file

@ -229,9 +229,7 @@ async def update_usage_alert_rule(
)
await require_scope(user, store, data.scope_type, data.scope_id, write=True)
require_notification_access(user, data, rule)
saved: Final = await store.save_rule(
data, user.user_id, existing=rule, overrides_defaults=user.user_role == LitellmUserRoles.PROXY_ADMIN
)
saved: Final = await store.save_rule(data, user.user_id, existing=rule, overrides_defaults=rule.overrides_defaults)
if saved is None:
raise HTTPException(status_code=409, detail="Alert rule changed while being updated. Refresh and try again")
return saved

View file

@ -132,6 +132,8 @@ async def update_family(
**({"enabled": data.enabled} if data.enabled is not None else {}),
**updates.get(row.record.direction, {}),
"updated_at": timestamp,
"evaluation_status": "pending",
"last_evaluated_day": None,
}
)
for row in rows

View file

@ -248,12 +248,20 @@ class IssuesRepository:
status: UsageAlertStatus,
expected_status: UsageAlertStatus,
actor: str | None = None,
authorized_scopes: tuple[IssueCohort, ...] | None = None,
) -> Issue | None:
if authorized_scopes is not None and len(authorized_scopes) > 10:
raise ValueError("Authorize at most 10 alert source scopes")
changed: Final = TypeAdapter(tuple[_Identifier, ...]).validate_python(
await self.database.query_raw(
f"""WITH candidate AS MATERIALIZED (
SELECT i.*,({_STATUS}) AS effective_status {_JOIN} WHERE i.id=$1
AND NOT(i.kind='model_deprecation' AND NOT i.active AND $2='open')
AND ($5::jsonb IS NULL OR NOT EXISTS (
SELECT 1 FROM "LiteLLM_IssueSignal" s WHERE s.issue_id=i.id AND NOT EXISTS (
SELECT 1 FROM jsonb_to_recordset($5::jsonb) AS x(scope_type text,scope_id text)
WHERE s.source='usage_alert' AND x.scope_type=s.record->'usage_alert'->>'scope_type'
AND x.scope_id IS NOT DISTINCT FROM s.record->'usage_alert'->>'scope_id')))
), ack AS (
INSERT INTO "LiteLLM_Config" AS prior (param_name,param_value)
SELECT ack_key,jsonb_build_object('model',record->'deprecation'->>'usage_model',
@ -266,18 +274,34 @@ class IssuesRepository:
WHERE EXISTS(SELECT 1 FROM candidate i WHERE i.ack_key=prior.param_name
AND ({_PRIOR_STATUS})=$3)
RETURNING param_name
) UPDATE "LiteLLM_Issue" i SET status=$2,updated_at=CURRENT_TIMESTAMP
FROM candidate c WHERE i.id=c.id AND c.effective_status=$3 AND
) UPDATE "LiteLLM_Issue" i SET status=$2,
updated_at=GREATEST(CURRENT_TIMESTAMP,i.updated_at+interval '1 millisecond')
FROM candidate c WHERE i.id=c.id AND (c.kind='model_deprecation' OR i.updated_at=c.updated_at)
AND c.effective_status=$3 AND
((c.kind<>'model_deprecation' AND i.status=$3) OR i.ack_key IN(SELECT param_name FROM ack))
RETURNING i.id""",
issue_id,
status,
expected_status,
actor,
json.dumps(tuple(scope.model_dump(mode="json") for scope in authorized_scopes))
if authorized_scopes is not None
else None,
)
)
return await self.get_issue(issue_id) if changed else None
async def status_scopes(self, issue_id: str) -> tuple[IssueCohort, ...]:
return TypeAdapter(tuple[IssueCohort, ...]).validate_python(
await self.database.query_raw(
"""SELECT DISTINCT CASE WHEN source='usage_alert' THEN record->'usage_alert'->>'scope_type'
ELSE 'deployment' END AS scope_type,
CASE WHEN source='usage_alert' THEN record->'usage_alert'->>'scope_id' END AS scope_id
FROM "LiteLLM_IssueSignal" WHERE issue_id=$1 ORDER BY scope_type,scope_id LIMIT 11""",
issue_id,
)
)
async def add_signals(self, issue_id: str, signals: tuple[IssueSignalInput, ...]) -> None:
await self.add_signal_batch(
tuple(IssueSignal.model_validate({**signal.model_dump(), "issue_id": issue_id}) for signal in signals)
@ -289,11 +313,30 @@ class IssuesRepository:
if len(signals) > 500:
raise ValueError("Write at most 500 issue signals per batch")
await self.database.query_raw(
"""INSERT INTO "LiteLLM_IssueSignal" (issue_id,source,source_id,record)
SELECT DISTINCT ON(x->>'issue_id',x->>'source',x->>'source_id')
x->>'issue_id',x->>'source',x->>'source_id',x FROM jsonb_array_elements($1::jsonb) x
WHERE EXISTS(SELECT 1 FROM "LiteLLM_Issue" i WHERE i.id=x->>'issue_id')
ON CONFLICT(issue_id,source,source_id) DO NOTHING RETURNING source_id""",
"""WITH incoming AS MATERIALIZED (
SELECT DISTINCT ON(x->>'issue_id',x->>'source',x->>'source_id') x AS record
FROM jsonb_array_elements($1::jsonb) x
), parents AS MATERIALIZED (
SELECT i.id FROM "LiteLLM_Issue" i WHERE i.id IN(SELECT record->>'issue_id' FROM incoming)
ORDER BY i.id FOR UPDATE
), added AS (
INSERT INTO "LiteLLM_IssueSignal" (issue_id,source,source_id,record)
SELECT x.record->>'issue_id',x.record->>'source',x.record->>'source_id',x.record
FROM incoming x JOIN parents p ON p.id=x.record->>'issue_id'
ON CONFLICT(issue_id,source,source_id) DO NOTHING RETURNING issue_id,source,record
) UPDATE "LiteLLM_Issue" i SET
updated_at=GREATEST(CURRENT_TIMESTAMP,i.updated_at+interval '1 millisecond'),
status=CASE WHEN i.kind='performance' AND i.active AND i.status<>'open' AND EXISTS (
SELECT 1 FROM added a WHERE a.issue_id=i.id AND a.source='usage_alert'
AND (a.record->'usage_alert'->>'scope_type'<>i.scope_type
OR a.record->'usage_alert'->>'scope_id' IS DISTINCT FROM i.scope_id)
AND NOT EXISTS(SELECT 1 FROM "LiteLLM_IssueSignal" prior
WHERE prior.issue_id=i.id AND prior.source='usage_alert'
AND prior.record->'usage_alert'->>'scope_type'=a.record->'usage_alert'->>'scope_type'
AND prior.record->'usage_alert'->>'scope_id'
IS NOT DISTINCT FROM a.record->'usage_alert'->>'scope_id')
) THEN 'open' ELSE i.status END
WHERE i.id IN(SELECT issue_id FROM added) RETURNING i.id""",
json.dumps(tuple(signal.model_dump(mode="json") for signal in signals)),
)

View file

@ -534,7 +534,12 @@ class UsageAlertsRepository:
async def key_ownership(self, key_id: str) -> tuple[str | None, str | None] | None:
rows: Final = TypeAdapter(tuple[_Ownership, ...]).validate_python(
await self.database.query_raw(
'SELECT user_id,team_id FROM "LiteLLM_VerificationToken" WHERE token=$1', key_id
"""SELECT user_id,team_id FROM "LiteLLM_VerificationToken" WHERE token=$1
UNION ALL
(SELECT user_id,team_id FROM "LiteLLM_DeletedVerificationToken" WHERE token=$1
AND NOT EXISTS(SELECT 1 FROM "LiteLLM_VerificationToken" WHERE token=$1)
ORDER BY deleted_at DESC,id DESC LIMIT 1)""",
key_id,
)
)
return (rows[0].user_id, rows[0].team_id) if rows else None

View file

@ -1,7 +1,8 @@
from datetime import date, datetime
from typing import Literal, Self, TypeAlias
from typing import Literal, TypeAlias
from pydantic import BaseModel, ConfigDict, Field, model_validator
from typing_extensions import Self
from litellm.types.proxy.management_endpoints.common_daily_activity import (
BoundedDailyActivityResponse,

View file

@ -1,9 +1,10 @@
"""Daily usage alert rules and incidents shared by the API and evaluator."""
from datetime import date, datetime
from typing import Literal, Self, TypeAlias
from typing import Literal, TypeAlias
from pydantic import BaseModel, ConfigDict, Field, model_validator
from typing_extensions import Self
UsageAlertScope: TypeAlias = Literal["deployment", "key", "user", "team", "model"]
UsageAlertMetric: TypeAlias = Literal["spend", "requests", "latency", "errors", "error_rate"]

View file

@ -157,7 +157,7 @@ async def test_issue_http_scope_pagination_and_source_visibility(gateway: Gatewa
path: Final = f"/issues/{first.id}"
own: Final = gateway.request("GET", path, key=owner_key)
assert own.status_code == 200 and own.json()["scope_label"] == "Example checkout"
assert own.json()["can_update_status"] is True
assert own.json()["can_update_status"] is False, "Broader source scopes require their own write permission"
assert own.json()["can_view_activity"] is True
assert gateway.request("GET", path, key=viewer_key).json()["can_update_status"] is False
assert gateway.request("GET", path, key=info_key).json()["can_update_status"] is False
@ -204,15 +204,18 @@ async def test_issue_http_scope_pagination_and_source_visibility(gateway: Gatewa
gateway.request(
"PATCH", path, {"status": "dismissed", "expected_status": "open"}, key=owner_key
).status_code
== 200
== 403
)
assert gateway.request("PATCH", path, {"status": "dismissed", "expected_status": "open"}).status_code == 200
assert gateway.request("PATCH", path, {"status": "resolved", "expected_status": "open"}).status_code == 409
assert gateway.request("GET", f"/issues/{second.id}", key=owner_key).json()["status"] == "open"
assert gateway.request("GET", f"/issues/{second.id}", key=owner_key).json()["can_update_status"] is True
assert (
gateway.request(
"PATCH", path, {"status": "resolved", "expected_status": "open"}, key=owner_key
"PATCH", f"/issues/{second.id}", {"status": "dismissed", "expected_status": "open"}, key=owner_key
).status_code
== 409
== 200
)
assert gateway.request("GET", f"/issues/{second.id}", key=owner_key).json()["status"] == "open"
finally:
await _remove_issues(db, (first.id, second.id, global_issue.id))
await alerts.delete_rule(rule.id)
@ -784,6 +787,189 @@ async def test_family_settings_preserve_custom_deleted_rules_and_abort_concurren
assert await alerts.get_rule(custom.id) == custom
@pytest.mark.asyncio(loop_scope="session")
async def test_family_edit_requires_a_new_evaluation(db: ManagementDatabase) -> None:
async with _family_database(db) as database:
borrowed: Final = cast(AlertDatabase, database) # cast-ok: real generated Prisma query delegate
alerts: Final = UsageAlertsRepository(borrowed)
store: Final = IssuesRepository(borrowed)
await alerts.ensure_default_rules(default_usage_rules())
rule: Final = await alerts.get_rule(default_usage_rule_id("key", "spend", "increase"))
assert rule is not None
await alerts.record_evaluation_status(rule, date.today() - timedelta(days=1), "complete")
result: Final = await store.update_family(
"spend", IssueFamilyUpdate(policies=(IssueFamilyPolicyUpdate(direction="increase", multiplier=5),))
)
saved: Final = await alerts.get_rule(rule.id)
assert result is not None and saved is not None
assert (result.evaluation_status_counts.complete, saved.evaluation_status, saved.last_evaluated_day) == (
0,
"pending",
None,
), "Changed family policy must reset stale evaluation status"
@pytest.mark.asyncio(loop_scope="session")
async def test_deleted_key_owner_retains_historical_issue_access(gateway: Gateway, db: ManagementDatabase) -> None:
store: Final = IssuesRepository.from_prisma(db.prisma)
with gateway.scenario() as scenario:
owner: Final = scenario.user(user_role="internal_user")
outsider: Final = scenario.user(user_role="internal_user")
owner_key: Final = scenario.key(user_id=owner, key_type="management")
outsider_key: Final = scenario.key(user_id=outsider, key_type="management")
reporting_id: Final = sha256(uuid4().bytes).hexdigest()
issue: Final = _issue(IssueCohort(scope_type="key", scope_id=reporting_id))
await _query(
db,
'INSERT INTO "LiteLLM_DeletedVerificationToken" (id,token,user_id,models) VALUES ($1,$1,$2,ARRAY[]::text[])',
reporting_id,
owner,
)
try:
await store.upsert_issue(issue)
path: Final = f"/issues/{issue.id}"
detail: Final = gateway.request("GET", path, key=owner_key)
assert detail.status_code == 200, "Deleted key owner must retain historical issue access"
assert detail.json()["can_update_status"] is True
assert gateway.request("GET", path, key=outsider_key).status_code == 403
assert (
gateway.request(
"PATCH", path, {"status": "dismissed", "expected_status": "open"}, key=owner_key
).status_code
== 200
)
await _query(
db,
'INSERT INTO "LiteLLM_VerificationToken" (token,user_id,models) VALUES ($1,$2,ARRAY[]::text[])',
reporting_id,
outsider,
)
assert gateway.request("GET", path, key=owner_key).status_code == 403
assert gateway.request("GET", path, key=outsider_key).status_code == 200
finally:
await _remove_issues(db, (issue.id,))
await _query(db, 'DELETE FROM "LiteLLM_VerificationToken" WHERE token=$1', reporting_id)
await _query(db, 'DELETE FROM "LiteLLM_DeletedVerificationToken" WHERE token=$1', reporting_id)
@pytest.mark.asyncio(loop_scope="session")
async def test_issue_status_rechecks_changed_scope_set_and_caps_sources(db: ManagementDatabase) -> None:
store: Final = IssuesRepository.from_prisma(db.prisma)
alerts: Final = UsageAlertsRepository.from_prisma(db.prisma)
issue: Final = _issue(IssueCohort(scope_type="user", scope_id=f"fixture-{uuid4()}"))
rule: Final = await alerts.save_rule(
UsageAlertRuleInput(
name="Scope guard",
scope_type="user",
scope_id=issue.cohort.scope_id,
metric="requests",
direction="increase",
),
"qa-admin",
)
assert rule is not None
try:
source: Final = (
await alerts.record_observations(
rule,
(
UsageAlertObservation(
entity_id=rule.scope_id,
day=issue.observed_on,
observed=420,
baseline=120,
threshold=360,
breached=True,
),
),
)
)[0]
await store.upsert_issue(issue)
await store.add_signals(
issue.id, (IssueSignalInput(source="usage_alert", source_id=source.id, usage_alert=source),)
)
authorized: Final = await store.status_scopes(issue.id)
assert authorized == (issue.cohort,)
assert await store.set_status(issue.id, "dismissed", "open", authorized_scopes=authorized) is not None
context: Final = source.model_copy(update={"id": str(uuid4()), "metric": "spend"})
same_scope: Final = (IssueSignalInput(source="usage_alert", source_id=context.id, usage_alert=context),)
await store.add_signals(issue.id, same_scope)
unchanged: Final = await store.get_issue(issue.id)
assert unchanged is not None and unchanged.status == "dismissed"
await store.add_signals(issue.id, same_scope)
assert await store.get_issue(issue.id) == unchanged
broader: Final = source.model_copy(update={"id": str(uuid4()), "scope_type": "deployment", "scope_id": None})
await store.add_signals(
issue.id, (IssueSignalInput(source="usage_alert", source_id=broader.id, usage_alert=broader),)
)
assert await store.set_status(issue.id, "dismissed", "open", authorized_scopes=authorized) is None
current: Final = await store.get_issue(issue.id)
assert current is not None and current.status == "open"
all_scopes: Final = await store.status_scopes(issue.id)
assert await store.set_status(issue.id, "dismissed", "open", authorized_scopes=all_scopes) is not None
broader_context: Final = broader.model_copy(update={"id": str(uuid4()), "metric": "spend"})
await store.add_signals(
issue.id,
(IssueSignalInput(source="usage_alert", source_id=broader_context.id, usage_alert=broader_context),),
)
repeated_scope: Final = await store.get_issue(issue.id)
assert repeated_scope is not None and repeated_scope.status == "dismissed"
assert await store.set_status(issue.id, "open", "dismissed") is not None
racing: Final = source.model_copy(update={"id": str(uuid4()), "scope_type": "team", "scope_id": "racing-team"})
async with asyncio.TaskGroup() as tasks:
async with db.prisma.writer_db.tx(timeout=timedelta(seconds=30)) as transaction:
await transaction.query_raw('SELECT id FROM "LiteLLM_Issue" WHERE id=$1 FOR UPDATE', issue.id)
pid: Final = (
TypeAdapter(tuple[_Pid, ...])
.validate_python(await transaction.query_raw("SELECT pg_backend_pid() AS pid"))[0]
.pid
)
update: Final = tasks.create_task(
store.set_status(issue.id, "dismissed", "open", authorized_scopes=all_scopes)
)
async with asyncio.timeout(10):
while (
not TypeAdapter(tuple[_Blocked, ...])
.validate_python(
await _query(
db,
"SELECT EXISTS(SELECT 1 FROM pg_stat_activity WHERE $1::int=ANY(pg_blocking_pids(pid))) AS blocked",
pid,
)
)[0]
.blocked
):
await asyncio.sleep(0.01)
locked: Final = IssuesRepository(
cast(AlertDatabase, transaction) # cast-ok: generated Prisma transaction query delegate
)
await locked.add_signals(
issue.id, (IssueSignalInput(source="usage_alert", source_id=racing.id, usage_alert=racing),)
)
assert update.result() is None
raced: Final = await store.get_issue(issue.id)
assert raced is not None and raced.status == "open"
await store.add_signals(
issue.id,
tuple(
IssueSignalInput(
source="usage_alert",
source_id=f"scope-{index}",
usage_alert=source.model_copy(update={"id": f"scope-{index}", "scope_id": f"other-user-{index}"}),
)
for index in range(20)
),
)
limited: Final = await store.status_scopes(issue.id)
assert len(limited) == 11
with pytest.raises(ValueError, match="at most 10"):
await store.set_status(issue.id, "open", "dismissed", authorized_scopes=limited)
finally:
await _remove_issues(db, (issue.id,))
await alerts.delete_rule(rule.id)
await _query(db, 'DELETE FROM "LiteLLM_UsageAlertIncident" WHERE rule_id=$1', rule.id)
@pytest.mark.asyncio(loop_scope="session")
async def test_resource_recovery_ignores_broader_and_traffic_context(db: ManagementDatabase) -> None:
store: Final = IssuesRepository.from_prisma(db.prisma)

View file

@ -86,6 +86,29 @@ def test_scoped_rules_use_real_management_auth_and_preserve_read_only_roles(gate
assert gateway.request("GET", path, key=owner_key).status_code == 404
def test_admin_edit_preserves_owner_rule_access(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
owner: Final = scenario.user(user_role="internal_user")
owner_key: Final = scenario.key(user_id=owner, key_type="management")
config: Final = {
"name": "Owner policy",
"scope_type": "user",
"scope_id": owner,
"metric": "spend",
"direction": "increase",
}
created: Final = gateway.request("POST", "/usage/alert-rules", config, key=owner_key)
assert created.status_code == 201, created.text
path: Final = f"/usage/alert-rules/{created.json()['id']}"
try:
edited: Final = gateway.request("PUT", path, {**config, "name": "Renamed by admin"})
assert edited.status_code == 200, edited.text
owner_edit: Final = gateway.request("PUT", path, {**config, "multiplier": 5}, key=owner_key)
assert owner_edit.status_code == 200, "Admin edit must preserve the owner's rule access"
finally:
assert gateway.request("DELETE", path).status_code == 204
def test_team_admin_can_manage_team_alerts_but_member_can_only_read(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
admin_id: Final = scenario.user(user_role="internal_user")