#!/usr/bin/env python3 """Quota-aware soft routing for automatic CLI lane provider selection.""" from __future__ import annotations import re import sys import urllib.request from dataclasses import dataclass from typing import Callable from cli_lane_config import ( QUOTA_METRICS_URL, QUOTA_MIN_REMAINING_PERCENT_DEFAULT, kanban_setting, ) from cli_lane_metrics import ( record_provider_quota, record_quota_fetch_failure, record_soft_exclusion, ) from cli_lane_routing import fresh_unavailable_provider # Exporter accounts are labeled by vendor; the lane routes by CLI provider. QUOTA_PROVIDER_LABELS = {"openai": "codex", "anthropic": "claude"} # Gate only on account-wide windows: a model-specific Claude window (for # example seven_day_opus) must not veto the whole provider. GATED_QUOTA_WINDOWS = { "codex": None, "claude": ("five_hour", "seven_day"), } _SAMPLE_PATTERN = re.compile( r"^atlas_ai_quota_(remaining_percent|reset_timestamp_seconds|fetch_success)" r"\{([^}]*)\} (-?[0-9.eE+]+)$" ) _LABEL_PATTERN = re.compile(r'([a-zA-Z_][a-zA-Z0-9_]*)="([^"]*)"') @dataclass(frozen=True) class ProviderQuota: """One provider's binding remaining quota as seen by the lane.""" remaining_percent: float reset_timestamp: float | None @dataclass(frozen=True) class SelectionConstraint: """One provider exclusion decided at a routing boundary, with evidence.""" exclude_provider: str | None exclude_reason: str | None source: str | None notes: tuple[str, ...] def _gated_window(provider: str, labels: dict[str, str]) -> str | None: """Return the window key when a sample belongs to a gated account limit.""" if labels.get("limit") != provider: return None window = labels.get("window", "") allowed = GATED_QUOTA_WINDOWS[provider] if allowed is not None and window not in allowed: return None return window def parse_quota_metrics(text: str) -> dict[str, ProviderQuota]: """Reduce exporter samples to each provider's tightest account window.""" remaining: dict[tuple[str, str], float] = {} resets: dict[tuple[str, str], float] = {} fetch_success: dict[str, bool] = {} for line in text.splitlines(): match = _SAMPLE_PATTERN.match(line.strip()) if not match: continue metric, raw_labels, raw_value = match.groups() labels = dict(_LABEL_PATTERN.findall(raw_labels)) provider = QUOTA_PROVIDER_LABELS.get(labels.get("provider", "")) if provider is None: continue value = float(raw_value) if metric == "fetch_success": fetch_success[provider] = value >= 1.0 continue window = _gated_window(provider, labels) if window is None: continue if metric == "remaining_percent": remaining[(provider, window)] = value else: resets[(provider, window)] = value snapshot: dict[str, ProviderQuota] = {} for provider in QUOTA_PROVIDER_LABELS.values(): windows = { window: value for (candidate, window), value in remaining.items() if candidate == provider } if not windows or not fetch_success.get(provider, True): # Failed or stale exporter fetches leave this provider ungated. continue binding = min(windows, key=lambda window: (windows[window], window)) snapshot[provider] = ProviderQuota( remaining_percent=windows[binding], reset_timestamp=resets.get((provider, binding)), ) return snapshot def fetch_quota_snapshot( url: str = QUOTA_METRICS_URL, open_url: Callable = urllib.request.urlopen, ) -> dict[str, ProviderQuota]: """Read the in-pod usage exporter; empty means the signal is unavailable.""" try: with open_url(url, timeout=10) as response: text = response.read(1 << 20).decode("utf-8", errors="replace") snapshot = parse_quota_metrics(text) except (OSError, ValueError) as error: record_quota_fetch_failure() print( f"quota snapshot unavailable; routing fails open: {type(error).__name__}", file=sys.stderr, flush=True, ) return {} if not snapshot: record_quota_fetch_failure() print( "quota snapshot carried no gated windows; routing fails open", file=sys.stderr, flush=True, ) return snapshot def quota_soft_exclusion( snapshot: dict[str, ProviderQuota], threshold: float ) -> tuple[str | None, str | None]: """Choose at most one provider to soft-exclude from new automatic work.""" below = { provider: quota.remaining_percent for provider, quota in snapshot.items() if quota.remaining_percent < threshold } if not below: return None, None excluded = min(below, key=lambda provider: (below[provider], provider)) if len(below) == len(snapshot) and len(snapshot) > 1: preferred = next(name for name in snapshot if name != excluded) note = ( "Quota degraded: every provider is below the " f"{threshold:.6g}% remaining threshold; preferring {preferred} " f"({snapshot[preferred].remaining_percent:.6g}% left) over " f"{excluded} ({below[excluded]:.6g}% left) for new auto work." ) return excluded, note note = ( f"Quota guard: {excluded} has {below[excluded]:.6g}% remaining " f"(threshold {threshold:.6g}%); routing new auto work to the other " "provider while active work finishes." ) return excluded, note def selection_constraint( assignee: str, *, snapshot_from: Callable[[], dict[str, ProviderQuota]] | None = None, now: float | None = None, ) -> SelectionConstraint: """Combine provider health and quota gating for one selection boundary.""" if assignee != "cli-auto": # Explicit lanes are operator decisions; they are never re-routed. return SelectionConstraint(None, None, None, ()) snapshot = (snapshot_from or fetch_quota_snapshot)() resets: dict[str, float] = {} for provider, quota in snapshot.items(): record_provider_quota(provider, quota.remaining_percent, quota.reset_timestamp) if quota.reset_timestamp is not None: resets[provider] = quota.reset_timestamp health_excluded = fresh_unavailable_provider(now=now, quota_resets=resets) if health_excluded is not None: return SelectionConstraint( health_excluded, "is unavailable according to fresh native health", "health", (), ) threshold = kanban_setting( "provider_quota_min_remaining_percent", QUOTA_MIN_REMAINING_PERCENT_DEFAULT ) excluded, note = quota_soft_exclusion(snapshot, threshold) for provider in snapshot: record_soft_exclusion(provider, provider == excluded) if excluded is None or note is None: return SelectionConstraint(None, None, None, ()) return SelectionConstraint( excluded, "is below its remaining-quota routing threshold", "quota", (note,), )