120 lines
3.6 KiB
Python
120 lines
3.6 KiB
Python
#!/usr/bin/env python3
|
|
"""Read Claude subscription quotas from first-party response headers."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
from typing import Any
|
|
from urllib.request import Request, urlopen
|
|
|
|
|
|
OAUTH_TOKEN_FILE = Path(
|
|
os.environ.get(
|
|
"ATLAS_AI_CLAUDE_OAUTH_TOKEN_FILE",
|
|
"/claude-oauth-access/token",
|
|
)
|
|
)
|
|
MESSAGES_URL = os.environ.get(
|
|
"ATLAS_AI_CLAUDE_MESSAGES_URL",
|
|
"https://api.anthropic.com/v1/messages",
|
|
)
|
|
QUOTA_MODEL = os.environ.get(
|
|
"ATLAS_AI_CLAUDE_QUOTA_MODEL",
|
|
"claude-haiku-4-5-20251001",
|
|
)
|
|
REQUEST_TIMEOUT_SECONDS = max(
|
|
5.0,
|
|
float(os.environ.get("ATLAS_AI_CLAUDE_QUERY_TIMEOUT_SECONDS", "30")),
|
|
)
|
|
WINDOW_HEADERS = {
|
|
"five_hour": "5h",
|
|
"seven_day": "7d",
|
|
}
|
|
|
|
|
|
class QuotaNotExposed(RuntimeError):
|
|
"""The active provider credential cannot expose account quota."""
|
|
|
|
|
|
def _token() -> str:
|
|
"""Read the Vault-rendered setup token from its memory-backed file."""
|
|
try:
|
|
token = OAUTH_TOKEN_FILE.read_text(encoding="utf-8").strip()
|
|
except (OSError, UnicodeError) as error:
|
|
raise QuotaNotExposed("Claude quota token is unavailable") from error
|
|
if not token:
|
|
raise QuotaNotExposed("Claude quota token is unavailable")
|
|
return token
|
|
|
|
|
|
def _percentage(value: str | None) -> float | None:
|
|
"""Convert Anthropic's fractional utilization header to a percentage."""
|
|
try:
|
|
utilization = float(value) if value is not None else None
|
|
except ValueError:
|
|
return None
|
|
if utilization is None or not 0 <= utilization <= 1:
|
|
return None
|
|
return utilization * 100
|
|
|
|
|
|
def _timestamp(value: str | None) -> float | None:
|
|
"""Validate one Unix reset timestamp from a response header."""
|
|
try:
|
|
timestamp = float(value) if value is not None else None
|
|
except ValueError:
|
|
return None
|
|
if timestamp is None or timestamp <= 0:
|
|
return None
|
|
return timestamp
|
|
|
|
|
|
def _payload(headers: Any) -> dict[str, Any]:
|
|
"""Translate bounded unified-rate-limit headers into the usage document."""
|
|
payload: dict[str, Any] = {}
|
|
for window, abbreviation in WINDOW_HEADERS.items():
|
|
prefix = f"anthropic-ratelimit-unified-{abbreviation}"
|
|
utilization = _percentage(headers.get(f"{prefix}-utilization"))
|
|
if utilization is None:
|
|
continue
|
|
result: dict[str, Any] = {"utilization": utilization}
|
|
reset = _timestamp(headers.get(f"{prefix}-reset"))
|
|
if reset is not None:
|
|
result["resets_at"] = reset
|
|
payload[window] = result
|
|
if not payload:
|
|
raise QuotaNotExposed("Claude response omitted subscription quota headers")
|
|
return payload
|
|
|
|
|
|
def query_claude() -> dict[str, Any]:
|
|
"""Make one minimal subscription request and return exact quota windows."""
|
|
token = _token()
|
|
body = json.dumps(
|
|
{
|
|
"model": QUOTA_MODEL,
|
|
"max_tokens": 1,
|
|
"messages": [{"role": "user", "content": "quota"}],
|
|
},
|
|
separators=(",", ":"),
|
|
).encode("utf-8")
|
|
request = Request(
|
|
MESSAGES_URL,
|
|
data=body,
|
|
method="POST",
|
|
headers={
|
|
"Authorization": f"Bearer {token}",
|
|
"Content-Type": "application/json",
|
|
"Accept": "application/json",
|
|
"anthropic-version": "2023-06-01",
|
|
"anthropic-beta": "oauth-2025-04-20,claude-code-20250219",
|
|
"User-Agent": "atlas-ai-usage-exporter/2.0",
|
|
},
|
|
)
|
|
with urlopen(request, timeout=REQUEST_TIMEOUT_SECONDS) as response:
|
|
payload = _payload(response.headers)
|
|
response.read(1 << 20)
|
|
return payload
|