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

#!/usr/bin/env python3
"""Export first-party coding CLI quotas without exposing account credentials."""

from __future__ import annotations

import json
import os
import threading
import time
from dataclasses import dataclass, field
from datetime import UTC, date, datetime, timedelta
from pathlib import Path
from typing import Any

import ai_usage_claude as claude_query
import ai_usage_codex as codex_query
import ai_usage_http as http_engine
import ai_usage_polling as polling_engine


PROVIDER_HEALTH_ROOT = Path(
    os.environ.get("ATLAS_AI_PROVIDER_HEALTH_ROOT", "/provider-health")
)
AUTHENTICATION_GRACE_SECONDS = max(
    60,
    int(os.environ.get("ATLAS_AI_AUTHENTICATION_GRACE_SECONDS", "1200")),
)
CLAUDE_WINDOWS = (
    "five_hour",
    "seven_day",
    "seven_day_opus",
    "seven_day_sonnet",
)
METRIC_HELP = {
    "atlas_ai_account_tokens": "First-party account token usage for a fixed period.",
    "atlas_ai_account_usage_summary": "First-party account usage summary values.",
    "atlas_ai_extra_usage_enabled": "Whether metered extra usage is enabled for the account.",
    "atlas_ai_provider_authenticated": "Whether the first-party provider access boundary is authenticated.",
    "atlas_ai_quota_fetch_duration_seconds": "Duration of the latest provider quota fetch.",
    "atlas_ai_quota_fetch_success": "Whether the latest provider quota fetch succeeded.",
    "atlas_ai_quota_last_attempt_timestamp_seconds": "Unix timestamp of the latest quota fetch attempt.",
    "atlas_ai_quota_last_success_timestamp_seconds": "Unix timestamp of the latest successful quota fetch.",
    "atlas_ai_quota_remaining_percent": "Remaining percentage in a first-party coding CLI quota window.",
    "atlas_ai_quota_reset_timestamp_seconds": "Unix timestamp when a coding CLI quota window resets.",
    "atlas_ai_quota_used_percent": "Used percentage in a first-party coding CLI quota window.",
    "atlas_ai_quota_window_duration_seconds": "Nominal duration of a coding CLI quota window.",
}

QuotaNotExposed = claude_query.QuotaNotExposed


@dataclass(frozen=True)
class Sample:
    """One Prometheus gauge sample."""

    name: str
    labels: dict[str, str]
    value: float


@dataclass
class ProviderState:
    """Latest safe samples and fetch health for one provider."""

    samples: list[Sample] = field(default_factory=list)
    authenticated: bool = False
    fetch_success: bool = False
    last_attempt: float = 0
    last_success: float = 0
    duration: float = 0


def _number(value: Any) -> float | None:
    """Return a finite numeric value while rejecting booleans and nulls."""
    if isinstance(value, bool) or not isinstance(value, (int, float)):
        return None
    converted = float(value)
    return (
        converted if converted == converted and abs(converted) != float("inf") else None
    )


def _timestamp(value: Any) -> float | None:
    """Parse either a Unix timestamp or an ISO-8601 timestamp."""
    number = _number(value)
    if number is not None:
        return number
    if not isinstance(value, str) or not value:
        return None
    try:
        return datetime.fromisoformat(value.replace("Z", "+00:00")).timestamp()
    except ValueError:
        return None


def _window_name(window: dict[str, Any], fallback: str) -> str:
    """Give common rolling windows stable, human-readable labels."""
    minutes = _number(window.get("windowDurationMins"))
    known = {300: "five_hour", 1440: "one_day", 10080: "seven_day"}
    if minutes is not None and int(minutes) in known:
        return known[int(minutes)]
    return fallback


def _codex_limit_name(limit_id: str, snapshot: dict[str, Any]) -> str:
    """Return a stable low-cardinality name for a Codex quota bucket."""
    if limit_id == "codex":
        return "codex"
    name = snapshot.get("limitName")
    if isinstance(name, str) and name:
        normalized = "".join(char.lower() if char.isalnum() else "-" for char in name)
        return "-".join(filter(None, normalized.split("-")))[:64]
    return "additional"


def parse_codex_payloads(
    rate_response: dict[str, Any],
    usage_response: dict[str, Any],
    *,
    today: date | None = None,
) -> list[Sample]:
    """Convert structured Codex app-server responses into bounded metrics."""
    samples: list[Sample] = []
    limits = rate_response.get("rateLimitsByLimitId")
    if not isinstance(limits, dict) or not limits:
        limits = {"codex": rate_response.get("rateLimits")}
    for limit_id, snapshot in limits.items():
        if not isinstance(limit_id, str) or not isinstance(snapshot, dict):
            continue
        limit_name = _codex_limit_name(limit_id, snapshot)
        for fallback, raw_window in (
            ("primary", snapshot.get("primary")),
            ("secondary", snapshot.get("secondary")),
        ):
            if not isinstance(raw_window, dict):
                continue
            used = _number(raw_window.get("usedPercent"))
            if used is None:
                continue
            labels = {
                "provider": "openai",
                "limit": limit_name,
                "window": _window_name(raw_window, fallback),
            }
            samples.extend(
                (
                    Sample("atlas_ai_quota_used_percent", labels, used),
                    Sample(
                        "atlas_ai_quota_remaining_percent", labels, max(0, 100 - used)
                    ),
                )
            )
            reset = _timestamp(raw_window.get("resetsAt"))
            duration = _number(raw_window.get("windowDurationMins"))
            if reset is not None:
                samples.append(
                    Sample("atlas_ai_quota_reset_timestamp_seconds", labels, reset)
                )
            if duration is not None:
                samples.append(
                    Sample(
                        "atlas_ai_quota_window_duration_seconds", labels, duration * 60
                    )
                )

    summary = usage_response.get("summary")
    if isinstance(summary, dict):
        for source, metric in (
            ("lifetimeTokens", "lifetime_tokens"),
            ("peakDailyTokens", "peak_daily_tokens"),
            ("currentStreakDays", "current_streak_days"),
            ("longestStreakDays", "longest_streak_days"),
            ("longestRunningTurnSec", "longest_running_turn_seconds"),
        ):
            value = _number(summary.get(source))
            if value is not None:
                samples.append(
                    Sample(
                        "atlas_ai_account_usage_summary",
                        {"provider": "openai", "metric": metric},
                        value,
                    )
                )

    current_day = today or datetime.now(UTC).date()
    daily = usage_response.get("dailyUsageBuckets")
    parsed_daily: list[tuple[date, float]] = []
    if isinstance(daily, list):
        for bucket in daily:
            if not isinstance(bucket, dict):
                continue
            value = _number(bucket.get("tokens"))
            try:
                start = date.fromisoformat(str(bucket.get("startDate")))
            except ValueError:
                continue
            if value is not None:
                parsed_daily.append((start, value))
    # The account endpoint publishes completed daily buckets and may not include
    # the current UTC day. Anchor fixed periods to the latest reported day so a
    # delayed bucket is not mislabeled as zero usage.
    anchor_day = max((start for start, _ in parsed_daily), default=current_day)
    for period, days in (("latest_day", 1), ("seven_day", 7), ("thirty_day", 30)):
        earliest = anchor_day - timedelta(days=days - 1)
        value = sum(
            tokens for start, tokens in parsed_daily if earliest <= start <= anchor_day
        )
        samples.append(
            Sample(
                "atlas_ai_account_tokens",
                {"provider": "openai", "period": period},
                value,
            )
        )
    return samples


def parse_claude_payload(payload: dict[str, Any]) -> list[Sample]:
    """Convert Claude's first-party OAuth usage document into bounded metrics."""
    samples: list[Sample] = []
    for window_name in CLAUDE_WINDOWS:
        window = payload.get(window_name)
        if not isinstance(window, dict):
            continue
        used = _number(window.get("utilization"))
        if used is None:
            continue
        labels = {"provider": "anthropic", "limit": "claude", "window": window_name}
        samples.extend(
            (
                Sample("atlas_ai_quota_used_percent", labels, used),
                Sample("atlas_ai_quota_remaining_percent", labels, max(0, 100 - used)),
            )
        )
        reset = _timestamp(window.get("resets_at"))
        if reset is not None:
            samples.append(
                Sample("atlas_ai_quota_reset_timestamp_seconds", labels, reset)
            )
    extra = payload.get("extra_usage")
    if isinstance(extra, dict):
        enabled = extra.get("is_enabled")
        if isinstance(enabled, bool):
            samples.append(
                Sample(
                    "atlas_ai_extra_usage_enabled",
                    {"provider": "anthropic"},
                    float(enabled),
                )
            )
    return samples


query_codex = codex_query.query_codex
query_claude = claude_query.query_claude


def _provider_authenticated(provider: str) -> bool:
    """Read one fresh non-secret broker authentication snapshot."""
    path = PROVIDER_HEALTH_ROOT / ("codex.json" if provider == "openai" else "claude.json")
    try:
        if time.time() - path.stat().st_mtime > 10 * 60:
            return False
        document = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, TypeError, ValueError, json.JSONDecodeError):
        return False
    return isinstance(document, dict) and document.get("authenticated") is True


def _escape_label(value: str) -> str:
    """Escape one Prometheus label value."""
    return value.replace("\\", "\\\\").replace("\n", "\\n").replace('"', '\\"')


class Collector:
    """Poll providers independently and expose a thread-safe metrics snapshot."""

    def __init__(self) -> None:
        self._lock = threading.Lock()
        self._providers = {name: ProviderState() for name in ("openai", "anthropic")}

    def record_failure(
        self,
        provider: str,
        *,
        started: float | None = None,
        monotonic_started: float | None = None,
    ) -> None:
        """Record a failed attempt while retaining the provider's last good samples."""
        attempted_at = time.time() if started is None else started
        duration_started = (
            time.monotonic() if monotonic_started is None else monotonic_started
        )
        with self._lock:
            state = self._providers[provider]
            state.last_attempt = attempted_at
            state.duration = max(0, time.monotonic() - duration_started)
            state.fetch_success = False

    def refresh_provider(self, provider: str) -> None:
        """Refresh one provider while retaining the last good values on failure."""
        started = time.time()
        monotonic_started = time.monotonic()
        try:
            if provider not in self._providers:
                raise ValueError("unknown provider")
            authenticated = _provider_authenticated(provider)
            with self._lock:
                self._providers[provider].authenticated = authenticated
            if provider == "openai":
                samples = parse_codex_payloads(*query_codex())
            elif provider == "anthropic":
                samples = parse_claude_payload(query_claude())
            else:
                raise ValueError("unknown provider")
        except QuotaNotExposed:
            self.record_failure(
                provider,
                started=started,
                monotonic_started=monotonic_started,
            )
            return
        except Exception as error:
            print(
                f"{provider} quota collection deferred: {type(error).__name__}",
                flush=True,
            )
            self.record_failure(
                provider,
                started=started,
                monotonic_started=monotonic_started,
            )
            return
        with self._lock:
            state = self._providers[provider]
            state.last_attempt = started
            state.duration = time.monotonic() - monotonic_started
            # A successful first-party quota request is direct authentication
            # proof. Broker health files may be quiet when no routed model call
            # has occurred recently, so they must not demote this live result.
            state.authenticated = True
            state.fetch_success = True
            state.samples = samples
            state.last_success = time.time()

    def render(self) -> bytes:
        """Render the current provider states in Prometheus text format."""
        with self._lock:
            states = {
                provider: ProviderState(**vars(state))
                for provider, state in self._providers.items()
            }
        now = time.time()
        for provider, state in states.items():
            quota_proves_access = (
                state.last_success > 0
                and 0 <= now - state.last_success <= AUTHENTICATION_GRACE_SECONDS
            )
            state.authenticated = (
                _provider_authenticated(provider) or quota_proves_access
            )
        samples: list[Sample] = []
        for provider, state in states.items():
            labels = {"provider": provider}
            samples.extend(state.samples)
            samples.extend(
                (
                    Sample(
                        "atlas_ai_provider_authenticated",
                        labels,
                        float(state.authenticated),
                    ),
                    Sample(
                        "atlas_ai_quota_fetch_success",
                        labels,
                        float(state.fetch_success),
                    ),
                    Sample(
                        "atlas_ai_quota_last_attempt_timestamp_seconds",
                        labels,
                        state.last_attempt,
                    ),
                    Sample(
                        "atlas_ai_quota_last_success_timestamp_seconds",
                        labels,
                        state.last_success,
                    ),
                    Sample(
                        "atlas_ai_quota_fetch_duration_seconds", labels, state.duration
                    ),
                )
            )
        lines: list[str] = []
        for name in sorted({sample.name for sample in samples}):
            lines.extend((f"# HELP {name} {METRIC_HELP[name]}", f"# TYPE {name} gauge"))
            for sample in sorted(
                (item for item in samples if item.name == name),
                key=lambda item: sorted(item.labels.items()),
            ):
                labels = ",".join(
                    f'{key}="{_escape_label(value)}"'
                    for key, value in sorted(sample.labels.items())
                )
                lines.append(f"{name}{{{labels}}} {sample.value:.12g}")
        return ("\n".join(lines) + "\n").encode("utf-8")


make_handler = http_engine.make_handler
Server = http_engine.Server
PollingEngine = polling_engine.PollingEngine


def main() -> int:
    """Poll quota APIs and serve only sanitized metrics and health endpoints."""
    collector = Collector()
    interval = max(60, int(os.environ.get("ATLAS_AI_USAGE_INTERVAL_SECONDS", "300")))

    poller = PollingEngine(collector, interval=interval)
    poller.start()
    port = int(os.environ.get("ATLAS_AI_USAGE_PORT", "9010"))
    server = Server(("0.0.0.0", port), make_handler(collector, poller))
    server.serve_forever()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())