atlas-iac/services/hermes/scripts/cli_lane_quota.py

#!/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,),
    )