diff --git a/app/models/notification.py b/app/models/notification.py index 6ec4d9b..afd524c 100644 --- a/app/models/notification.py +++ b/app/models/notification.py @@ -1,12 +1,12 @@ """Notification model and notification-type catalog. -A persisted per-user notification produced through the single ``notify()`` +A persisted per-user notification produced through the single `notify()` service. Notification types are shared with the notification wiring (149) and the preference model (150) so opt-outs map 1:1 to the types the app produces. -Security: ``payload`` must never contain secrets, invitation tokens, or another +Security: `payload` must never contain secrets, invitation tokens, or another user's data. Users can only ever read their own notifications (enforced by the -notification API which always filters on ``current_user.id``). +notification API which always filters on `current_user.id`). """ from __future__ import annotations @@ -21,6 +21,7 @@ NOTIF_ROLE_CHANGE = "role_change" NOTIF_AI_EVENT = "ai_event" NOTIF_SHARE = "share" +NOTIF_BUDGET_ALERT = "budget_alert" # Notification types that are always delivered: users cannot disable security- # critical system notifications about their own access. @@ -31,6 +32,7 @@ NOTIF_ROLE_CHANGE, NOTIF_AI_EVENT, NOTIF_SHARE, + NOTIF_BUDGET_ALERT, ) diff --git a/app/services/token_usage.py b/app/services/token_usage.py index cdf1e12..1d398e7 100644 --- a/app/services/token_usage.py +++ b/app/services/token_usage.py @@ -13,7 +13,7 @@ def estimate_tokens(text) -> int: - """Estimate the token count of ``text`` (>= 1 for non-empty input).""" + """Estimate the token count of `text` (>= 1 for non-empty input).""" if text is None: return 0 value = str(text) @@ -67,7 +67,7 @@ def usage_from_text(prompt_text, completion_text) -> dict: def sum_usage(messages) -> dict: - """Cumulative usage across ``messages`` (missing counts ignored).""" + """Cumulative usage across `messages` (missing counts ignored).""" prompt = completion = total = 0 for message in messages or []: if getattr(message, "total_tokens", None) is None: @@ -76,3 +76,45 @@ def sum_usage(messages) -> dict: completion += message.completion_tokens or 0 total += message.total_tokens or 0 return _totals(prompt, completion) + + +class BudgetAlertDeduplicator: + """Per-period deduplication of budge threshold alerts. + + A given threshold alert fires at most once per budget per period. The + deduplicator records the period in which each threshold was alerted and + suppresses repeated crossings within that period. Starting a new period + resets the recorded alerts so the threshold can fire again. + """ + + def __init__(self): + self._period = None + self._alerted = set() + + def _ensure_period(self, period) -> None: + if period != self._period: + self._period = period + self._alerted = set() + + def should_alert(self, period, threshold) -> bool: + """Return True if the alert has not yet fired for this period/threshold.""" + self._ensure_period(period) + return threshold not in self._alerted + + def mark_alerted(self, period, threshold) -> None: + """Record that the alert fired for this period/threshold.""" + self._ensure_period(period) + self._alerted.add(threshold) + + def check_and_mark(self, period, threshold) -> bool: + """Atomically check and record the alert; True if it should fire.""" + self._ensure_period(period) + if threshold in self._alerted: + return False + self._alerted.add(threshold) + return True + + def reset(self, period=None) -> None: + """Clear recorded alerts, optionally switching to a new period.""" + self._period = period + self._alerted = set() diff --git a/tests/test_notifications.py b/tests/test_notifications.py index d62b309..df144a8 100644 --- a/tests/test_notifications.py +++ b/tests/test_notifications.py @@ -167,3 +167,30 @@ def test_remove_member_creates_membership_notification(self, client, make_user, created = Notification.query.filter_by(user_id=member.id, type="membership").one() assert created.is_read is False + + +class TestBudgetAlertDeduplication: + def test_repeated_crossings_alert_once_per_period(self, app, make_user): + from app.services.budget_alerts import check_budget_alerts + + user = make_user() + workspace = Workspace(user_id=user.id, name="Budget ws") + db.session.add(workspace) + db.session.commit() + + alerts = check_budget_alerts(workspace.id, spend=150, limit=100, period="2024-01") + assert len(alerts) == 1 + alerts = check_budget_alerts(workspace.id, spend=150, limit=100, period="2024-01") + assert len(alerts) == 0 + + def test_new_period_resets_deduplication(self/ app, make_user): + from app.services.budget_alerts import check_budget_alerts + + user = make_user() + workspace = Workspace(user_id=user.id, name="Budget ws") + db.session.add(workspace) + db.session.commit() + + assert len(check_budget_alerts(workspace.id, spend=150, limit=100, period="2024-01")) == 1 + assert len(check_budget_alerts(workspace.id, spend=150, limit=100, period="2024-01")) == 0 + assert len(check_budget_alerts(workspace.id, spend=150, limit=100, period="2024-02")) == 1