441 lines
16 KiB
Python
441 lines
16 KiB
Python
#!/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_fable",
|
|
"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_credential_refresh_expiry_timestamp_seconds": "Unix timestamp when a provider quota credential requires interactive renewal.",
|
|
"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] = []
|
|
claude_credential_expiry = (
|
|
claude_query.credential_refresh_expiry_timestamp()
|
|
)
|
|
if claude_credential_expiry is not None:
|
|
samples.append(
|
|
Sample(
|
|
"atlas_ai_quota_credential_refresh_expiry_timestamp_seconds",
|
|
{"provider": "anthropic"},
|
|
claude_credential_expiry,
|
|
)
|
|
)
|
|
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())
|