feat(accounts): add per-account usage limits

app/core/balancer/__init__.py

@@ -1,4 +1,5 @@
 from app.core.balancer.logic import (
+    ACCOUNT_USAGE_LIMIT_REACHED_ERROR_CODE,
     ERROR_BACKOFF_THRESHOLD,
     HEALTH_TIER_DRAINING,
     HEALTH_TIER_HEALTHY,
@@ -23,6 +24,7 @@
     TrafficClass,
     UsageWeightedOrder,
     account_status_for_permanent_failure,
+    account_usage_limit_blocks_selection,
     configure_replica_salt,
     evaluate_health_tier,
     failover_decision,
@@ -31,6 +33,7 @@
     handle_rate_limit,
     plausible_rate_limit_reset_at,
     pool_usage_exhaustion,
+    routing_eligible_states,
     select_account,
 )
 
@@ -39,6 +42,7 @@
     "HEALTH_TIER_DRAINING",
     "HEALTH_TIER_HEALTHY",
     "HEALTH_TIER_PROBING",
+    "ACCOUNT_USAGE_LIMIT_REACHED_ERROR_CODE",
     "ERROR_BACKOFF_THRESHOLD",
     "REAUTH_REQUIRED_FAILURE_CODES",
     "AccountState",
@@ -59,6 +63,7 @@
     "USAGE_LIMIT_REACHED",
     "UsageWeightedOrder",
     "account_status_for_permanent_failure",
+    "account_usage_limit_blocks_selection",
     "configure_replica_salt",
     "evaluate_health_tier",
     "failover_decision",
@@ -67,5 +72,6 @@
     "handle_rate_limit",
     "plausible_rate_limit_reset_at",
     "pool_usage_exhaustion",
+    "routing_eligible_states",
     "select_account",
 ]

app/core/balancer/logic.py

@@ -6,11 +6,12 @@
 import random
 import socket
 import time
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from typing import Collection, Iterable, Literal
 
 from app.core.balancer.types import FailureClass, UpstreamError
 from app.core.usage import PLAN_CAPACITY_CREDITS_SECONDARY
+from app.core.usage.account_limits import AccountUsageLimitState
 from app.core.utils.retry import backoff_seconds, parse_retry_after
 from app.db.models import AccountStatus
 
@@ -103,6 +104,10 @@
 PRESERVE_MIN_SHORT_WINDOW_FLOOR_PCT = 10.0
 NORMAL_LAST_ACCOUNT_EMERGENCY_FLOOR_PCT = 5.0
 RECENT_FOREGROUND_ACTIVITY_SECONDS = 30 * 60
+ACCOUNT_USAGE_LIMIT_REACHED_ERROR_CODE = "account_usage_limit_reached"
+ACCOUNT_USAGE_LIMIT_REACHED_ERROR_MESSAGE = (
+    "All otherwise available accounts have reached their usage limit or lack current usage data"
+)
 
 logger = logging.getLogger(__name__)
 
@@ -140,6 +145,8 @@ class AccountState:
     leased_tokens: float = 0.0
     routing_policy: str = ROUTING_POLICY_NORMAL
     ignore_standard_quota: bool = False
+    usage_limit_state: AccountUsageLimitState = AccountUsageLimitState.DISABLED
+    usage_limit_percent: float | None = None
 
 
 @dataclass
@@ -444,6 +451,101 @@ def _fallback_secondary_capacity_credits(plan_type: str | None) -> float:
     )
 
 
+def _prepare_routing_candidates(
+    states: Iterable[AccountState],
+    *,
+    current: float,
+    ignore_standard_quota: bool,
+    bypass_quota_exceeded: bool,
+    bypass_account_ids: Collection[str] | None,
+) -> tuple[list[AccountState], list[AccountState], list[AccountState]]:
+    all_states = list(states)
+    available: list[AccountState] = []
+    in_error_backoff: list[AccountState] = []
+    usage_limit_blocked: list[AccountState] = []
+    bypass_ids = set(bypass_account_ids or ())
+
+    for state in all_states:
+        bypass_standard_quota = (
+            ignore_standard_quota
+            or state.ignore_standard_quota
+            or bypass_quota_exceeded
+            or state.account_id in bypass_ids
+        )
+        if state.status in (AccountStatus.REAUTH_REQUIRED, AccountStatus.DEACTIVATED):
+            continue
+        if state.status == AccountStatus.PAUSED:
+            continue
+        if state.status == AccountStatus.RATE_LIMITED:
+            if state.reset_at and current >= state.reset_at:
+                state.status = AccountStatus.ACTIVE
+                state.used_percent = 0.0
+                state.error_count = 0
+                state.reset_at = None
+            elif not bypass_standard_quota:
+                continue
+        if state.status == AccountStatus.QUOTA_EXCEEDED:
+            if state.reset_at and current >= state.reset_at:
+                state.status = AccountStatus.ACTIVE
+                state.used_percent = 0.0
+                state.secondary_used_percent = 0.0
+                state.reset_at = None
+            elif not bypass_standard_quota:
+                continue
+        if state.cooldown_until and current >= state.cooldown_until:
+            state.cooldown_until = None
+            state.last_error_at = None
+            state.error_count = 0
+        if state.cooldown_until and current < state.cooldown_until:
+            continue
+        if account_usage_limit_blocks_selection(state):
+            usage_limit_blocked.append(state)
+            continue
+        if state.error_count >= ERROR_BACKOFF_THRESHOLD:
+            backoff = min(300, 30 * (2 ** (state.error_count - ERROR_BACKOFF_THRESHOLD)))
+            if state.last_error_at and current - state.last_error_at < backoff:
+                in_error_backoff.append(state)
+                continue
+            # Error backoff expired — reset error state so recovery is
+            # not penalised by stale counts. The account has already
+            # been held back for the full backoff period; letting it
+            # re-enter the pool with a clean slate avoids the problem
+            # where a previously-high error_count causes an immediate
+            # return to maximum backoff on the very next transient error.
+            state.error_count = 0
+            state.last_error_at = None
+        available.append(state)
+
+    return available, in_error_backoff, usage_limit_blocked
+
+
+def routing_eligible_states(
+    states: Iterable[AccountState],
+    *,
+    now: float | None = None,
+    traffic_class: TrafficClass = TRAFFIC_CLASS_FOREGROUND,
+) -> list[AccountState]:
+    """Return the pool-wide states eligible for the requested traffic class."""
+    current = time.time() if now is None else now
+    state_list = list(states)
+    evaluated_states = [replace(state) for state in state_list]
+    available, _, _ = _prepare_routing_candidates(
+        evaluated_states,
+        current=current,
+        ignore_standard_quota=False,
+        bypass_quota_exceeded=False,
+        bypass_account_ids=None,
+    )
+    if traffic_class == TRAFFIC_CLASS_OPPORTUNISTIC and available:
+        available, _ = _filter_opportunistic_candidates(available, current)
+    eligible_evaluations = {id(state) for state in available}
+    return [
+        state
+        for state, evaluated in zip(state_list, evaluated_states, strict=True)
+        if id(evaluated) in eligible_evaluations
+    ]
+
+
 def select_account(
     states: Iterable[AccountState],
     now: float | None = None,
@@ -530,59 +632,16 @@ def select_account(
         human-readable error message when no account is eligible.
     """
     current = now or time.time()
-    available: list[AccountState] = []
-    in_error_backoff: list[AccountState] = []
+    bypass_account_ids = None if bypass_quota_exceeded_account_ids is None else set(bypass_quota_exceeded_account_ids)
     all_states = list(states)
+    available, in_error_backoff, usage_limit_blocked = _prepare_routing_candidates(
+        all_states,
+        current=current,
+        ignore_standard_quota=ignore_standard_quota,
+        bypass_quota_exceeded=bypass_quota_exceeded,
+        bypass_account_ids=bypass_account_ids,
+    )
     usage_exhaustion_state_list = list(usage_exhaustion_states) if usage_exhaustion_states is not None else all_states
-    bypass_account_ids = None if bypass_quota_exceeded_account_ids is None else set(bypass_quota_exceeded_account_ids)
-
-    for state in all_states:
-        bypass_standard_quota = (
-            ignore_standard_quota
-            or state.ignore_standard_quota
-            or bypass_quota_exceeded
-            or (bypass_account_ids is not None and state.account_id in bypass_account_ids)
-        )
-        if state.status in (AccountStatus.REAUTH_REQUIRED, AccountStatus.DEACTIVATED):
-            continue
-        if state.status == AccountStatus.PAUSED:
-            continue
-        if state.status == AccountStatus.RATE_LIMITED:
-            if state.reset_at and current >= state.reset_at:
-                state.status = AccountStatus.ACTIVE
-                state.used_percent = 0.0
-                state.error_count = 0
-                state.reset_at = None
-            elif not bypass_standard_quota:
-                continue
-        if state.status == AccountStatus.QUOTA_EXCEEDED:
-            if state.reset_at and current >= state.reset_at:
-                state.status = AccountStatus.ACTIVE
-                state.used_percent = 0.0
-                state.secondary_used_percent = 0.0
-                state.reset_at = None
-            elif not bypass_standard_quota:
-                continue
-        if state.cooldown_until and current >= state.cooldown_until:
-            state.cooldown_until = None
-            state.last_error_at = None
-            state.error_count = 0
-        if state.cooldown_until and current < state.cooldown_until:
-            continue
-        if state.error_count >= ERROR_BACKOFF_THRESHOLD:
-            backoff = min(300, 30 * (2 ** (state.error_count - ERROR_BACKOFF_THRESHOLD)))
-            if state.last_error_at and current - state.last_error_at < backoff:
-                in_error_backoff.append(state)
-                continue
-            # Error backoff expired — reset error state so recovery is
-            # not penalised by stale counts. The account has already
-            # been held back for the full backoff period; letting it
-            # re-enter the pool with a clean slate avoids the problem
-            # where a previously-high error_count causes an immediate
-            # return to maximum backoff on the very next transient error.
-            state.error_count = 0
-            state.last_error_at = None
-        available.append(state)
 
     if traffic_class == TRAFFIC_CLASS_OPPORTUNISTIC and available:
         opportunistic_available, reason = _filter_opportunistic_candidates(available, current)
@@ -592,7 +651,7 @@ def select_account(
 
     if not available:
         in_error_backoff_ids = {state.account_id for state in in_error_backoff}
-        hard_blocked_exists = any(
+        hard_blocked_exists = bool(_routing_relevant_usage_limit_blocks(usage_limit_blocked)) or any(
             state.status
             in (
                 AccountStatus.PAUSED,
@@ -617,6 +676,12 @@ def _backoff_expires_at(s: AccountState) -> float:
                     return SelectionResult(None, f"opportunistic burn window closed: {reason}")
                 available = opportunistic_available
         else:
+            if _routing_relevant_usage_limit_blocks(usage_limit_blocked):
+                return SelectionResult(
+                    None,
+                    ACCOUNT_USAGE_LIMIT_REACHED_ERROR_MESSAGE,
+                    ACCOUNT_USAGE_LIMIT_REACHED_ERROR_CODE,
+                )
             if allow_usage_exhaustion_error:
                 usage_exhaustion = pool_usage_exhaustion(
                     usage_exhaustion_state_list,
@@ -802,6 +867,21 @@ def _oldest_due_probing_account(
     )
 
 
+def account_usage_limit_blocks_selection(state: AccountState) -> bool:
+    return state.usage_limit_state in {
+        AccountUsageLimitState.REACHED,
+        AccountUsageLimitState.DATA_UNAVAILABLE,
+    }
+
+
+def _routing_relevant_usage_limit_blocks(states: Iterable[AccountState]) -> list[AccountState]:
+    return [
+        state
+        for state in states
+        if state.status not in {AccountStatus.PAUSED, AccountStatus.REAUTH_REQUIRED, AccountStatus.DEACTIVATED}
+    ]
+
+
 def _remaining_secondary_credits(state: AccountState) -> float:
     """Return remaining absolute credits for the secondary (7-day) window."""
     capacity = (

app/core/usage/__init__.py

@@ -297,7 +297,7 @@ def _has_real_quota_metadata(row: UsageWindowRow) -> bool:
     return row.window_minutes is not None and row.window_minutes > 0 and row.reset_at is not None
 
 
-def _is_no_data_placeholder(row: UsageWindowRow) -> bool:
+def is_no_data_placeholder(row: UsageWindowRow) -> bool:
     """A no-data placeholder is the absence of a measurement, not 0% used.
 
     Such rows (no positive window duration AND no reset deadline) are written
@@ -339,9 +339,9 @@ def _should_prefer_primary_row(primary_row: UsageWindowRow, secondary_row: Usage
     # 0% used, so it must never displace a real weekly sample (otherwise the
     # dashboard jumps to 100% remaining every refresh).
     primary_has_real = _has_real_quota_metadata(primary_row)
-    if primary_has_real and _is_no_data_placeholder(secondary_row):
+    if primary_has_real and is_no_data_placeholder(secondary_row):
         return True
-    if _has_real_quota_metadata(secondary_row) and _is_no_data_placeholder(primary_row):
+    if _has_real_quota_metadata(secondary_row) and is_no_data_placeholder(primary_row):
         return False
 
     # Both real or both placeholder (same fetch / indeterminate ordering):