diff --git a/litellm/integrations/SlackAlerting/usage_alerts.py b/litellm/integrations/SlackAlerting/usage_alerts.py index 0d69406ca1a..dd10b5aea6d 100644 --- a/litellm/integrations/SlackAlerting/usage_alerts.py +++ b/litellm/integrations/SlackAlerting/usage_alerts.py @@ -5,15 +5,18 @@ import hashlib import json import os from html import escape +from itertools import islice from typing import TYPE_CHECKING, ClassVar, Final from urllib.parse import urlencode, urlsplit +from uuid import NAMESPACE_URL, uuid5 from httpx import Response from pydantic import TypeAdapter +from litellm.models.issues import IssueStore from litellm.types.integrations.slack_alerting import AlertType -from litellm.types.proxy.issues import IssuePendingDelivery -from litellm.types.proxy.usage_alerts import UsageAlert, UsageAlertChannel +from litellm.types.proxy.issues import IssueDeliveryBinding, IssueDeliveryInput, IssuePendingDelivery +from litellm.types.proxy.usage_alerts import UsageAlert if TYPE_CHECKING: from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting @@ -56,19 +59,43 @@ class UsageAlertDelivery: # One retained send per worker bounds SMTP threads even when a job is cancelled. _email_task: ClassVar[asyncio.Task[None] | None] = None _email_key: ClassVar[str | None] = None + _email_store: ClassVar[IssueStore | None] = None - def __init__(self, slack: "SlackAlerting") -> None: + def __init__(self, slack: "SlackAlerting", store: IssueStore | None = None) -> None: self._slack = slack + self._store = store async def send_issue(self, pending: IssuePendingDelivery) -> bool: + urls: Final = self._webhooks() if pending.delivery.channel == "slack" else () + if len(urls) > 1 and pending.delivery.slack_destination is None: + if self._store is None: + raise RuntimeError("Multiple alert destinations require a delivery store") + bindings: Final = ( + IssueDeliveryBinding( + delivery=IssueDeliveryInput( + id=str(uuid5(NAMESPACE_URL, f"{pending.delivery.id}:{destination}")), + issue_id=pending.issue.id, + channel="slack", + slack_destination=destination, + parent_delivery_id=pending.delivery.id, + ), + rule_id=pending.rule.id, + rule_revision=pending.rule.updated_at.isoformat(), + source_id=pending.alert.id, + ) + for destination in (hashlib.sha256(url.encode()).hexdigest() for url in urls) + ) + while batch := tuple(islice(bindings, 10)): + await self._store.queue_deliveries(batch) + return True path: Final = "/alerts?" + urlencode({"issue": pending.issue.id}) issue_url: Final = activity_url( pending.alert.model_copy(update={"activity_path": path}), os.getenv("PROXY_BASE_URL") ) source_url: Final = activity_url(pending.alert, os.getenv("PROXY_BASE_URL")) return await self._send( - pending.delivery.channel, - pending.delivery.email_recipients, + pending, + urls, incident_message(pending.alert, None), f"LiteLLM alert: {pending.issue.title}", (("View alert", issue_url), ("View source activity", source_url)), @@ -76,12 +103,14 @@ class UsageAlertDelivery: async def _send( self, - channel: UsageAlertChannel, - recipients: tuple[str, ...], + pending: IssuePendingDelivery, + urls: tuple[str, ...], message: str, subject: str, links: tuple[tuple[str, str], ...], ) -> bool: + channel: Final = pending.delivery.channel + recipients: Final = pending.delivery.email_recipients if channel not in configured_destinations(self._slack): return False if channel == "email": @@ -89,8 +118,19 @@ class UsageAlertDelivery: f'

{escape(label)}

' for label, url in links ) for recipient in recipients: - await self._email(recipient, subject, html) + await self._email(pending.delivery.id, recipient, subject, html) return bool(recipients) + destination: Final = pending.delivery.slack_destination + webhook: Final = next( + (url for url in urls if destination is None or hashlib.sha256(url.encode()).hexdigest() == destination), + None, + ) + if webhook is None: + return False + await self._post(webhook, message + "\n" + "\n".join(f"{label}: {url}" for label, url in links)) + return True + + def _webhooks(self) -> tuple[str, ...]: overrides: Final = self._slack.alert_to_webhook_url or {} configured_url: Final = ( overrides.get(AlertType.usage_alerts) @@ -99,38 +139,51 @@ class UsageAlertDelivery: or os.getenv("ALERTING_WEBHOOK_URL") ) if not configured_url: - return False - urls: Final = (configured_url,) if isinstance(configured_url, str) else _URLS.validate_python(configured_url) - for webhook in urls: - await self._post(webhook, message + "\n" + "\n".join(f"{label}: {url}" for label, url in links)) - return True + return () + return (configured_url,) if isinstance(configured_url, str) else _URLS.validate_python(configured_url) @staticmethod def _email_finished(task: asyncio.Task[None]) -> None: if not task.cancelled(): task.exception() # consume a late failure after the calling job was cancelled - @classmethod - async def _email(cls, recipient: str, subject: str, html: str) -> None: - key: Final = hashlib.sha256(json.dumps((recipient, subject, html)).encode()).hexdigest() + async def _email(self, key: str, recipient: str, subject: str, html: str) -> None: + cls: Final = type(self) previous: Final = cls._email_task if previous is not None: if not previous.done(): raise RuntimeError("An alert email is still in flight; retry later") - if cls._email_key == key and not previous.cancelled() and previous.exception() is None: - cls._email_task = None - cls._email_key = None - return previous.result() + if not previous.cancelled() and previous.exception() is None: + if cls._email_store is not None and cls._email_key is not None: + await cls._email_store.mark_delivery_sent(cls._email_key) + if cls._email_task is not previous: + raise RuntimeError("Another alert email is still in flight; retry later") + elif cls._email_key != key: + raise RuntimeError("A completed alert email is awaiting its receipt") + if cls._email_key == key: + cls._clear_email(previous) + return previous.result() task: Final = asyncio.create_task(cls._send_email(recipient, subject, html)) cls._email_task = task cls._email_key = key + cls._email_store = self._store task.add_done_callback(cls._email_finished) try: await asyncio.shield(task) - finally: - if task.done() and cls._email_task is task: - cls._email_task = None - cls._email_key = None + except asyncio.CancelledError: + raise + except Exception: + cls._clear_email(task) + raise + else: + cls._clear_email(task) + + @classmethod + def _clear_email(cls, task: asyncio.Task[None]) -> None: + if cls._email_task is task: + cls._email_task = None + cls._email_key = None + cls._email_store = None @staticmethod async def _send_email(recipient: str, subject: str, html: str) -> None: diff --git a/litellm/proxy/alerting_endpoints/usage_alerts.py b/litellm/proxy/alerting_endpoints/usage_alerts.py index c954c40cc48..3950efdc1b4 100644 --- a/litellm/proxy/alerting_endpoints/usage_alerts.py +++ b/litellm/proxy/alerting_endpoints/usage_alerts.py @@ -49,6 +49,8 @@ ScopeId: TypeAlias = Annotated[str | None, Query(min_length=1, max_length=1024)] def _allows_write(grant: str, route: str) -> bool: if grant == "info_routes": return False + if route == grant or route.startswith(grant + "/"): + return True paths: Final = ( TypeAdapter(tuple[str, ...]).validate_python(LiteLLMRoutes[grant].value) if grant in LiteLLMRoutes.__members__ diff --git a/litellm/proxy/spend_tracking/usage_alert_job.py b/litellm/proxy/spend_tracking/usage_alert_job.py index d26b463ea70..374d9ee433a 100644 --- a/litellm/proxy/spend_tracking/usage_alert_job.py +++ b/litellm/proxy/spend_tracking/usage_alert_job.py @@ -145,7 +145,8 @@ class UsageAlertJob: background=True, ) metrics: Final = UsageAlertMetricsRepository(database) - delivery: Final = UsageAlertDelivery(slack) + issues: Final = IssuesRepository(database) + delivery: Final = UsageAlertDelivery(slack, issues) return cls( UsageAlertsRepository(database), metrics, @@ -158,7 +159,7 @@ class UsageAlertJob: ), deprecation_projection=deprecation_projection, performance_projection=PerformanceIssueProjector( - IssuesRepository(database), + issues, metrics, ProviderLatencyRepository(database), delivery, diff --git a/litellm/repositories/issues.py b/litellm/repositories/issues.py index c70f16a85d3..62f3b8fa9de 100644 --- a/litellm/repositories/issues.py +++ b/litellm/repositories/issues.py @@ -617,7 +617,7 @@ class IssuesRepository: SELECT r.record AS rule,a.record AS alert FROM "LiteLLM_IssueDeliveryRule" b JOIN "LiteLLM_UsageAlertRule" r ON r.id=b.rule_id JOIN "LiteLLM_UsageAlertIncident" a ON a.id=b.source_id AND a.rule_id=r.id - WHERE b.delivery_id=d.id AND r.enabled AND NOT r.deleted AND a.breached AND a.status='open' + WHERE b.delivery_id=COALESCE(d.record->>'parent_delivery_id',d.id) AND r.enabled AND NOT r.deleted AND a.breached AND a.status='open' AND (r.record->'notification_channels') ? (d.record->>'channel') AND (d.record->>'channel'<>'email' OR (jsonb_array_length(d.record->'email_recipients')>0 diff --git a/litellm/repositories/usage_alert_health.py b/litellm/repositories/usage_alert_health.py index 646aa0421c0..f1ee7643672 100644 --- a/litellm/repositories/usage_alert_health.py +++ b/litellm/repositories/usage_alert_health.py @@ -1,5 +1,6 @@ """Sample successful aggregate flushes without adding request-time telemetry.""" +import asyncio from datetime import datetime, timedelta, timezone from typing import ( TYPE_CHECKING, @@ -11,6 +12,7 @@ from typing import ( from pydantic import BaseModel, ConfigDict, TypeAdapter +from litellm.repositories.analytics_budget import AnalyticsDatabase, BoundedAnalyticsDatabase, analytics_deadline from litellm.repositories.usage_alert_metrics import MetricsDatabase if TYPE_CHECKING: @@ -34,13 +36,18 @@ class UsageAlertHealthRepository: @classmethod def from_prisma(cls, prisma: "PrismaClient") -> "UsageAlertHealthRepository": - database: Final = cast(MetricsDatabase, prisma.writer_db) # cast-ok: generated Prisma delegate - return cls(database) + database: Final = cast(AnalyticsDatabase, prisma.writer_db) # cast-ok: generated Prisma transaction delegate + return cls(BoundedAnalyticsDatabase(database, background=True)) async def record(self, source: FlushSource, now: datetime) -> bool: hour: Final = now.astimezone(timezone.utc).replace(minute=0, second=0, microsecond=0, tzinfo=None) rows: Final = await self._database.query_raw( - """INSERT INTO "LiteLLM_UsageAlertFlush" (source,hour,observed_at) + """WITH removed AS ( + DELETE FROM "LiteLLM_UsageAlertFlush" WHERE (source,hour) IN ( + SELECT source,hour FROM "LiteLLM_UsageAlertFlush" + WHERE hour<$2::timestamp-INTERVAL '35 days' ORDER BY hour,source LIMIT 10 + ) + ) INSERT INTO "LiteLLM_UsageAlertFlush" (source,hour,observed_at) SELECT $1,$2::timestamp,$3::timestamp WHERE NOT EXISTS ( SELECT 1 FROM "LiteLLM_SpendLogDelivery" WHERE state='pending' AND created_at<$2::timestamp LIMIT 1 ) ON CONFLICT (source,hour) DO UPDATE SET observed_at=EXCLUDED.observed_at RETURNING source""", @@ -48,9 +55,6 @@ class UsageAlertHealthRepository: hour, now.astimezone(timezone.utc).replace(tzinfo=None), ) - await self._database.query_raw( - 'DELETE FROM "LiteLLM_UsageAlertFlush" WHERE hour<$1::timestamp RETURNING source', hour - timedelta(days=35) - ) return bool(rows) async def coverage(self, source: FlushSource, start: datetime) -> FlushCoverage: @@ -70,11 +74,11 @@ class UsageAlertHealthRepository: class UsageAlertFlushRecorder: - """At most one heartbeat attempt per source per worker per hour after success.""" + """Retry optional freshness once a minute, and stop after an hourly success.""" def __init__(self) -> None: - self._user_hour: datetime | None = None - self._team_hour: datetime | None = None + self._user_after: datetime | None = None + self._team_after: datetime | None = None async def record(self, prisma_client: "PrismaClient", source: FlushSource) -> None: from litellm._logging import verbose_proxy_logger @@ -84,17 +88,21 @@ class UsageAlertFlushRecorder: return now: Final = datetime.now(timezone.utc) hour: Final = now.replace(minute=0, second=0, microsecond=0) - prior: Final = self._user_hour if source == "user" else self._team_hour - if prior == hour: + prior: Final = self._user_after if source == "user" else self._team_after + if prior is not None and now < prior: return + self._defer(source, now + timedelta(minutes=1)) try: - database: Final = cast(MetricsDatabase, prisma_client.writer_db) # cast-ok: generated Prisma delegate - if not await UsageAlertHealthRepository(database).record(source, now): - return + with analytics_deadline(asyncio.get_running_loop().time() + 1): + if not await UsageAlertHealthRepository.from_prisma(prisma_client).record(source, now): + return except Exception: # noqa: BLE001 # observability must not fail or requeue a committed spend flush verbose_proxy_logger.warning("Could not record usage-alert aggregate freshness") return + self._defer(source, hour + timedelta(hours=1)) + + def _defer(self, source: FlushSource, until: datetime) -> None: if source == "user": - self._user_hour = hour + self._user_after = until else: - self._team_hour = hour + self._team_after = until diff --git a/litellm/repositories/usage_alerts.py b/litellm/repositories/usage_alerts.py index 30f4de63f40..87e6be5c0f1 100644 --- a/litellm/repositories/usage_alerts.py +++ b/litellm/repositories/usage_alerts.py @@ -596,12 +596,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( diff --git a/litellm/types/proxy/issues.py b/litellm/types/proxy/issues.py index fd05bea9564..3341ee8dc7f 100644 --- a/litellm/types/proxy/issues.py +++ b/litellm/types/proxy/issues.py @@ -207,6 +207,8 @@ class IssueDeliveryInput(BaseModel, frozen=True): issue_id: str = Field(min_length=1, max_length=200) channel: UsageAlertChannel email_recipients: tuple[str, ...] = Field(default=(), max_length=20) + slack_destination: str | None = Field(default=None, pattern="^[a-f0-9]{64}$") + parent_delivery_id: str | None = Field(default=None, min_length=1, max_length=200) class IssueDeliveryBinding(BaseModel, frozen=True): diff --git a/tests/integration/database/test_bounded_daily_activity.py b/tests/integration/database/test_bounded_daily_activity.py index 420731cbaaa..1ffe35c421d 100644 --- a/tests/integration/database/test_bounded_daily_activity.py +++ b/tests/integration/database/test_bounded_daily_activity.py @@ -1,15 +1,11 @@ -from collections.abc import Mapping from datetime import date, datetime, timedelta, timezone from hashlib import sha256 from typing import Final from uuid import uuid4 -import psycopg import pytest from integration._support.actors import ManagementDatabase from integration._support.client import Gateway -from psycopg import sql -from pydantic import TypeAdapter from litellm.litellm_core_utils.ptu_pricing import PTU_COST_ATTRIBUTION_ENV_VAR from litellm.repositories.bounded_daily_activity import BoundedDailyActivityRepository @@ -190,16 +186,8 @@ async def test_bounded_activity_routes_keep_exact_totals_and_scoped_permissions( ) if include_today == "true": assert bounded_today.json()["metadata"]["total_api_requests"] == 11 - schemas: Final = TypeAdapter(tuple[Mapping[str, str], ...]).validate_python( - await repository.database.query_raw("SELECT current_schema() AS name") - ) - assert gateway.database_url is not None - with psycopg.connect(gateway.database_url) as blocker: - blocker.execute( - sql.SQL("LOCK TABLE {} IN ACCESS EXCLUSIVE MODE").format( - sql.Identifier(schemas[0]["name"], "LiteLLM_DailyUserSpend") - ) - ) + async with db.prisma.db.tx(timeout=timedelta(seconds=10)) as blocker: + await blocker.execute_raw('LOCK TABLE "LiteLLM_DailyUserSpend" IN ACCESS EXCLUSIVE MODE') busy: Final = gateway.request("GET", path, params={**query, "bounded": "true"}) assert busy.status_code == 503 and busy.headers["Retry-After"] == "5" assert gateway.request("GET", path, params={**query, "bounded": "true"}).status_code == 200 diff --git a/tests/integration/database/test_issues.py b/tests/integration/database/test_issues.py index 8fc20613a4f..670396bd0ed 100644 --- a/tests/integration/database/test_issues.py +++ b/tests/integration/database/test_issues.py @@ -1,5 +1,6 @@ import asyncio import json +import os from collections.abc import AsyncGenerator, Mapping, Sequence from contextlib import asynccontextmanager from datetime import date, datetime, timedelta, timezone @@ -406,16 +407,8 @@ async def test_issue_activity_bounds_models_without_losing_totals(gateway: Gatew ) assert non_key.model_breakdown is None and non_key.metadata == bounded.metadata assert len(non_key.results) == 3 - schemas: Final = TypeAdapter(tuple[Mapping[str, str], ...]).validate_python( - await _query(db, "SELECT current_schema() AS name") - ) - assert gateway.database_url is not None - with psycopg.connect(gateway.database_url) as blocker: - blocker.execute( - sql.SQL("LOCK TABLE {} IN ACCESS EXCLUSIVE MODE").format( - sql.Identifier(schemas[0]["name"], "LiteLLM_DailyUserSpend") - ) - ) + async with db.prisma.writer_db.tx(timeout=timedelta(seconds=10)) as blocker: + await blocker.execute_raw('LOCK TABLE "LiteLLM_DailyUserSpend" IN ACCESS EXCLUSIVE MODE') busy: Final = gateway.request("GET", f"/issues/{issue.id}/activity") assert busy.status_code == 503 and busy.headers["Retry-After"] == "5" assert gateway.request("GET", f"/issues/{issue.id}/activity").status_code == 200 @@ -654,8 +647,7 @@ class _Blocked(BaseModel, frozen=True): @asynccontextmanager async def _family_database(db: ManagementDatabase) -> AsyncGenerator[Prisma]: - url: Final = db.gateway.database_url - assert url is not None + url: Final = db.gateway.database_url or os.environ["DATABASE_URL"] schema: Final = f"issues_{uuid4().hex}" parsed: Final = urlsplit(url) scoped: Final = urlunsplit(parsed._replace(query=urlencode({**dict(parse_qsl(parsed.query)), "schema": schema}))) diff --git a/tests/integration/database/test_performance_issues.py b/tests/integration/database/test_performance_issues.py index 2e2379f8a1b..269a2a83942 100644 --- a/tests/integration/database/test_performance_issues.py +++ b/tests/integration/database/test_performance_issues.py @@ -1,4 +1,5 @@ import asyncio +import os from collections.abc import AsyncIterator from contextlib import asynccontextmanager from datetime import date, datetime, timedelta, timezone @@ -7,11 +8,14 @@ from typing import Final from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit from uuid import uuid4 +import httpx import pytest from integration._support.actors import ManagementDatabase from integration._support.prisma_clients import connected_proxy_prisma from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting +from litellm.integrations.SlackAlerting.usage_alerts import UsageAlertDelivery +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy.spend_tracking.usage_alert_job import UsageAlertJob from litellm.proxy.utils import PrismaClient from litellm.repositories.daily_activity_buckets import api_key_bucket @@ -154,8 +158,7 @@ async def test_retirement_usage_omits_oversized_models_before_materialization(db @asynccontextmanager async def _isolated(db: ManagementDatabase) -> AsyncIterator[PrismaClient]: schema: Final = "issue_test_" + uuid4().hex - assert db.gateway.database_url is not None - parsed: Final = urlsplit(db.gateway.database_url) + parsed: Final = urlsplit(db.gateway.database_url or os.environ["DATABASE_URL"]) query: Final = {**dict(parse_qsl(parsed.query)), "schema": schema} url: Final = urlunsplit(parsed._replace(query=urlencode(query))) await db.prisma.writer_db.execute_raw(f'CREATE SCHEMA "{schema}"') @@ -871,3 +874,175 @@ async def test_projection_resume_keeps_spend_and_traffic_in_one_episode(db: Mana assert tuple(item.id for item in after.data) == tuple(item.id for item in before.data) links: Final = await IssuesRepository.from_prisma(prisma).list_signals(after.data[0].id) assert any(link.usage_alert is not None and link.usage_alert.metric == "requests" for link in links.data) + + +@pytest.mark.asyncio(loop_scope="session") +async def test_slack_retry_only_repeats_the_failed_destination(db: ManagementDatabase) -> None: + paths: Final[list[str]] = [] + + def respond(request: httpx.Request) -> httpx.Response: + paths.append(request.url.path) + return httpx.Response(503 if request.url.path == "/two" and paths.count("/two") == 1 else 200) + + handler: Final = AsyncHTTPHandler(transport=httpx.MockTransport(respond)) + try: + async with _isolated(db) as prisma: + await _seed(prisma) + rules: Final = UsageAlertsRepository.from_prisma(prisma) + assert ( + await rules.save_rule( + UsageAlertRuleInput( + name="Project spend", + scope_type="key", + scope_id=_KEY_A, + metric="spend", + direction="increase", + notification_channels=("slack",), + ), + "qa-admin", + ) + is not None + ) + assert ( + await rules.save_rule( + UsageAlertRuleInput( + name="Another project rule", + scope_type="key", + scope_id=_KEY_A, + metric="spend", + direction="increase", + notification_channels=("slack",), + ), + "qa-admin", + ) + is not None + ) + slack: Final = SlackAlerting( + alerting=["slack"], + alert_types=[], + async_http_handler=handler, + alert_to_webhook_url={ + "usage_alerts": [ + "https://hooks.example.com/one", + "https://hooks.example.com/two", + "https://hooks.example.com/one", + ] + }, + ) + job: Final = UsageAlertJob.from_prisma(prisma, slack) + await job._run(_NOW) + assert paths == ["/one", "/two"] or paths == ["/two", "/one"] + await prisma.writer_db.query_raw( + 'UPDATE "LiteLLM_IssueDelivery" SET next_attempt=CURRENT_TIMESTAMP WHERE NOT sent RETURNING id' + ) + retry: Final = (await IssuesRepository.from_prisma(prisma).pending_deliveries())[0] + assert ( + await rules.save_rule( + UsageAlertRuleInput.model_validate( + retry.rule.model_dump(include=set(UsageAlertRuleInput.model_fields)) + ).model_copy(update={"enabled": False}), + "qa-admin", + retry.rule, + ) + is not None + ) + await job._performance_projection.notify() + assert paths.count("/one") == 1 and paths.count("/two") == 2 + assert await IssuesRepository.from_prisma(prisma).pending_deliveries() == () + records: Final = await prisma.writer_db.query_raw( + 'SELECT record::text AS value FROM "LiteLLM_IssueDelivery"' + ) + assert all("hooks.example.com" not in row["value"] for row in records) + finally: + await handler.close() + + +@pytest.mark.asyncio(loop_scope="session") +async def test_cancelled_email_receipt_survives_a_different_delivery( + db: ManagementDatabase, monkeypatch: pytest.MonkeyPatch +) -> None: + accepted: Final = asyncio.Event() + release: Final = asyncio.Event() + messages: Final[list[bytes]] = [] + + async def smtp(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + writer.write(b"220 example.com ready\r\n") + await writer.drain() + try: + while command := await reader.readline(): + if command.upper().startswith(b"DATA"): + writer.write(b"354 End with a dot\r\n") + await writer.drain() + messages.append(await reader.readuntil(b"\r\n.\r\n")) + accepted.set() + await release.wait() + writer.write(b"250 Accepted\r\n") + elif command.upper().startswith(b"QUIT"): + writer.write(b"221 Closing\r\n") + break + else: + writer.write(b"250 OK\r\n") + await writer.drain() + finally: + writer.close() + await writer.wait_closed() + + server: Final = await asyncio.start_server(smtp, "127.0.0.1", 0) + monkeypatch.setenv("SMTP_HOST", "127.0.0.1") + monkeypatch.setenv("SMTP_PORT", str(server.sockets[0].getsockname()[1])) + monkeypatch.setenv("SMTP_SENDER_EMAIL", "alerts@example.com") + monkeypatch.setenv("SMTP_TLS", "False") + monkeypatch.setenv("SMTP_USE_SSL", "False") + for name in ("SMTP_TIMEOUT", "SMTP_USERNAME", "SMTP_PASSWORD"): + monkeypatch.delenv(name, raising=False) + try: + async with _isolated(db) as prisma: + await _seed(prisma) + rules: Final = UsageAlertsRepository.from_prisma(prisma) + assert ( + await rules.save_rule( + UsageAlertRuleInput( + name="Project spend", + scope_type="key", + scope_id=_KEY_A, + metric="spend", + direction="increase", + notification_channels=("email",), + email_recipients=("qa@example.com",), + ), + "qa-admin", + ) + is not None + ) + job: Final = UsageAlertJob.from_prisma(prisma, SlackAlerting(alerting=[], alert_types=[])) + await job._run(_NOW) + store: Final = IssuesRepository.from_prisma(prisma) + await prisma.writer_db.query_raw( + 'UPDATE "LiteLLM_IssueDelivery" SET next_attempt=CURRENT_TIMESTAMP WHERE NOT sent RETURNING id' + ) + pending: Final = (await store.pending_deliveries())[0] + delivery: Final = UsageAlertDelivery(SlackAlerting(alerting=["email"], alert_types=[]), store) + sending: Final = asyncio.create_task(delivery.send_issue(pending)) + await asyncio.wait_for(accepted.wait(), 10) + sending.cancel() + with pytest.raises(asyncio.CancelledError): + await sending + release.set() + retained: Final = UsageAlertDelivery._email_task + assert retained is not None + await asyncio.wait_for(asyncio.shield(retained), 10) + different: Final = pending.model_copy( + update={"delivery": pending.delivery.model_copy(update={"id": "another-delivery"})} + ) + assert await delivery.send_issue(different) + assert await store.pending_delivery(pending.delivery.id) is None + assert len(messages) == 2 + finally: + release.set() + server.close() + await server.wait_closed() + if UsageAlertDelivery._email_task is not None: + await asyncio.wait_for(asyncio.shield(UsageAlertDelivery._email_task), 10) + UsageAlertDelivery._email_task = None + UsageAlertDelivery._email_key = None + UsageAlertDelivery._email_store = None diff --git a/tests/integration/database/test_usage_alert_evaluation.py b/tests/integration/database/test_usage_alert_evaluation.py index 4b926174e14..799c36df03f 100644 --- a/tests/integration/database/test_usage_alert_evaluation.py +++ b/tests/integration/database/test_usage_alert_evaluation.py @@ -1,10 +1,24 @@ +import asyncio +import json from datetime import date, datetime, timedelta, timezone -from typing import Final +from typing import Final, cast # noqa: TID251 # generated Prisma transaction implements the accounting protocol +from uuid import uuid4 import pytest from integration._support.actors import ManagementDatabase +from integration.database.test_performance_issues import _isolated -from litellm.repositories.usage_alert_health import UsageAlertHealthRepository +from litellm.accounting.application import apply_deliveries +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter +from litellm.proxy.utils import ProxyLogging +from litellm.repositories.accounting_delivery_repository import ( + AccountingDeliveryRepository, + DeliveryProjection, + LedgerDatabase, + LedgerTransaction, +) +from litellm.repositories.usage_alert_health import UsageAlertFlushRecorder, UsageAlertHealthRepository from litellm.repositories.usage_alerts import UsageAlertsRepository from litellm.types.proxy.usage_alerts import UsageAlertRuleInput @@ -53,3 +67,75 @@ async def test_freshness_timestamps_use_the_real_prisma_parameter_types(db: Mana start.replace(tzinfo=None), (start + timedelta(days=1, hours=1)).replace(tzinfo=None), ) + + +@pytest.mark.asyncio(loop_scope="session") +async def test_freshness_lock_cannot_stall_committed_spend(db: ManagementDatabase) -> None: + recorder: Final = UsageAlertFlushRecorder() + async with db.prisma.writer_db.tx(timeout=timedelta(seconds=10)) as transaction: + await transaction.execute_raw('LOCK TABLE "LiteLLM_UsageAlertFlush" IN ACCESS EXCLUSIVE MODE') + await asyncio.wait_for(recorder.record(db.prisma, "user"), timeout=2) + await asyncio.wait_for(recorder.record(db.prisma, "user"), timeout=0.1) + await UsageAlertFlushRecorder().record(db.prisma, "user") + rows: Final = await db.prisma.writer_db.query_raw( + "SELECT source FROM \"LiteLLM_UsageAlertFlush\" WHERE source='user' AND hour=date_trunc('hour',CURRENT_TIMESTAMP AT TIME ZONE 'UTC')" + ) + assert rows == [{"source": "user"}] + + +@pytest.mark.asyncio(loop_scope="session") +async def test_freshness_waits_until_accounting_has_committed_the_daily_totals(db: ManagementDatabase) -> None: + async with _isolated(db) as prisma: + await prisma.writer_db.execute_raw( + 'CREATE TABLE "LiteLLM_DailyModelUsage" (LIKE public."LiteLLM_DailyModelUsage" INCLUDING ALL)' + ) + identity: Final = "coverage-" + uuid4().hex + now: Final = datetime.now(timezone.utc).replace(minute=0, second=0, microsecond=0) + timedelta(days=3) + row: Final = { + "request_id": identity, + "user": identity, + "api_key": "synthetic-key", + "startTime": (now - timedelta(hours=2)).isoformat(), + "model": "example-model", + "model_group": "example-model", + "custom_llm_provider": "example-provider", + "spend": 5, + "metadata": {}, + } + health: Final = UsageAlertHealthRepository.from_prisma(prisma) + writer: Final = DBSpendUpdateWriter() + logging: Final = ProxyLogging(UserApiKeyCache()) + + async def apply(transaction: LedgerTransaction, deliveries: tuple[DeliveryProjection, ...]) -> None: + await apply_deliveries( + transaction, deliveries, prisma_client=prisma, writer=writer, proxy_logging_obj=logging + ) + assert not await health.record("user", now) + + try: + await prisma.writer_db.query_raw( + 'INSERT INTO "LiteLLM_SpendLogDelivery" (request_id,accounting_payload,created_at,is_batch) VALUES ($1,$2::jsonb,$3::timestamp,false) RETURNING delivery_id', + identity, + json.dumps(row), + (now - timedelta(hours=1)).replace(tzinfo=None), + ) + assert not await health.record("user", now) + database: Final = cast(LedgerDatabase, prisma) # cast-ok: generated Prisma transaction delegate + result: Final = await AccountingDeliveryRepository(database).apply_pending(apply, batch_size=1) + assert result.applied == 1 + assert await health.record("user", now) + totals: Final = await prisma.writer_db.query_raw( + 'SELECT spend,api_requests FROM "LiteLLM_DailyUserSpend" WHERE user_id=$1', identity + ) + assert totals == [{"spend": 5, "api_requests": 1}] + finally: + await prisma.writer_db.query_raw( + 'DELETE FROM "LiteLLM_SpendLogDelivery" WHERE request_id=$1 RETURNING delivery_id', identity + ) + await prisma.writer_db.query_raw( + 'DELETE FROM "LiteLLM_DailyUserSpend" WHERE user_id=$1 RETURNING id', identity + ) + await prisma.writer_db.query_raw( + "DELETE FROM \"LiteLLM_UsageAlertFlush\" WHERE source='user' AND hour=$1::timestamp RETURNING source", + now.replace(tzinfo=None), + ) diff --git a/tests/integration/database/test_usage_alerts.py b/tests/integration/database/test_usage_alerts.py index 77362be35d0..29dc17b859e 100644 --- a/tests/integration/database/test_usage_alerts.py +++ b/tests/integration/database/test_usage_alerts.py @@ -39,12 +39,37 @@ 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) -def test_scoped_rules_use_real_management_auth_and_preserve_read_only_roles(gateway: Gateway) -> None: +@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: owner: Final = scenario.user(user_role="internal_user") outsider: Final = scenario.user(user_role="internal_user") viewer: Final = scenario.user(user_role="proxy_admin_viewer") - owner_key: Final = scenario.key(user_id=owner, key_type="management") + owner_key: Final = scenario.key(user_id=owner, key_type="default", allowed_routes=[grant]) outsider_key: Final = scenario.key(user_id=outsider, key_type="management") viewer_key: Final = scenario.key(user_id=viewer, key_type="read_only") inference_key: Final = scenario.key(user_id=owner, key_type="llm_api") diff --git a/tests/unit/proxy/alerting_endpoints/test_usage_alerts.py b/tests/unit/proxy/alerting_endpoints/test_usage_alerts.py index 1270b6809a9..239151f9ce1 100644 --- a/tests/unit/proxy/alerting_endpoints/test_usage_alerts.py +++ b/tests/unit/proxy/alerting_endpoints/test_usage_alerts.py @@ -6,7 +6,7 @@ import pytest from fastapi import HTTPException from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.alerting_endpoints.usage_alerts import require_scope, require_write_access +from litellm.proxy.alerting_endpoints.usage_alerts import _allows_write, require_scope, require_write_access from litellm.repositories.usage_alerts import UsageAlertsRepository, activity_path from litellm.types.proxy.auth.user_roles import LitellmUserRoles @@ -25,6 +25,19 @@ def test_info_grants_do_not_authorize_mutations(role: LitellmUserRoles, key_type assert error.value.status_code == 403 +@pytest.mark.parametrize("grant", ["management_routes", "/usage", "/usage/alert-rules", "/usage/*"]) +def test_explicit_route_grants_authorize_alert_mutations(grant: str) -> None: + assert _allows_write(grant, "/usage/alert-rules/example") is True + + +@pytest.mark.parametrize("grant", ["/usag", "/usage/alert-rule", "/usage/alerts"]) +def test_route_grants_require_a_complete_path_segment(grant: str) -> None: + user: Final = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, allowed_routes=[grant]) + with pytest.raises(HTTPException) as error: + require_write_access(user, "/usage/alert-rules/example") + assert error.value.status_code == 403 + + @pytest.mark.asyncio @pytest.mark.parametrize( "role,key_type,write",