412 lines
14 KiB
Python
412 lines
14 KiB
Python
#!/usr/bin/env python3
|
|
"""Read Claude subscription quotas without exposing account credentials."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import stat
|
|
import time
|
|
import uuid
|
|
from pathlib import Path
|
|
from typing import Any
|
|
from urllib.error import HTTPError, URLError
|
|
from urllib.request import Request, urlopen
|
|
|
|
|
|
OAUTH_TOKEN_FILE = Path(
|
|
os.environ.get(
|
|
"ATLAS_AI_CLAUDE_OAUTH_TOKEN_FILE",
|
|
"/claude-oauth-access/token",
|
|
)
|
|
)
|
|
CREDENTIALS_FILE = Path(
|
|
os.environ.get(
|
|
"ATLAS_AI_CLAUDE_CREDENTIALS",
|
|
"/runtime-access/claude/.credentials.json",
|
|
)
|
|
)
|
|
USAGE_URL = os.environ.get(
|
|
"ATLAS_AI_CLAUDE_USAGE_URL",
|
|
"https://api.anthropic.com/api/oauth/usage",
|
|
)
|
|
TOKEN_URL = os.environ.get(
|
|
"ATLAS_AI_CLAUDE_TOKEN_URL",
|
|
"https://platform.claude.com/v1/oauth/token",
|
|
)
|
|
MESSAGES_URL = os.environ.get(
|
|
"ATLAS_AI_CLAUDE_MESSAGES_URL",
|
|
"https://api.anthropic.com/v1/messages",
|
|
)
|
|
OAUTH_CLIENT_ID = "9d1c250a-e61b-44d9-88ed-5944d1962f5e"
|
|
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",
|
|
}
|
|
SCOPED_WINDOWS = (
|
|
"five_hour",
|
|
"seven_day",
|
|
"seven_day_opus",
|
|
"seven_day_sonnet",
|
|
)
|
|
FABLE_DISPLAY_NAMES = frozenset(("fable", "fable 5"))
|
|
DEFAULT_SCOPES = (
|
|
"user:profile",
|
|
"user:inference",
|
|
"user:sessions:claude_code",
|
|
"user:mcp_servers",
|
|
"user:file_upload",
|
|
)
|
|
MAX_DOCUMENT_BYTES = 1 << 20
|
|
REFRESH_SKEW_MILLISECONDS = 5 * 60 * 1000
|
|
|
|
|
|
class QuotaNotExposed(RuntimeError):
|
|
"""The active provider credential cannot expose account quota."""
|
|
|
|
|
|
def _setup_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 _used_percentage(value: object) -> float | None:
|
|
"""Validate one percentage returned by the scoped usage endpoint."""
|
|
if isinstance(value, bool) or not isinstance(value, (int, float)):
|
|
return None
|
|
percentage = float(value)
|
|
if not 0 <= percentage <= 100:
|
|
return None
|
|
return percentage
|
|
|
|
|
|
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 _header_payload(headers: Any) -> dict[str, Any]:
|
|
"""Translate bounded unified-rate-limit headers into a 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 _bounded_window(value: object) -> dict[str, Any] | None:
|
|
"""Copy only the numeric utilization and bounded reset field."""
|
|
if not isinstance(value, dict):
|
|
return None
|
|
utilization = _used_percentage(value.get("utilization"))
|
|
if utilization is None:
|
|
return None
|
|
result: dict[str, Any] = {"utilization": utilization}
|
|
reset = value.get("resets_at")
|
|
if isinstance(reset, (int, float, str)) and not isinstance(reset, bool):
|
|
result["resets_at"] = reset
|
|
return result
|
|
|
|
|
|
def _scoped_payload(document: object) -> dict[str, Any]:
|
|
"""Normalize overall and exact Fable weekly limits from Claude Code usage."""
|
|
if not isinstance(document, dict):
|
|
raise QuotaNotExposed("Claude usage response is not an object")
|
|
payload: dict[str, Any] = {}
|
|
for window in SCOPED_WINDOWS:
|
|
bounded = _bounded_window(document.get(window))
|
|
if bounded is not None:
|
|
payload[window] = bounded
|
|
limits = document.get("limits")
|
|
if isinstance(limits, list):
|
|
for limit in limits[:100]:
|
|
if not isinstance(limit, dict) or limit.get("kind") != "weekly_scoped":
|
|
continue
|
|
scope = limit.get("scope")
|
|
model = scope.get("model") if isinstance(scope, dict) else None
|
|
display_name = model.get("display_name") if isinstance(model, dict) else None
|
|
if not isinstance(display_name, str):
|
|
continue
|
|
if display_name.strip().casefold() not in FABLE_DISPLAY_NAMES:
|
|
continue
|
|
utilization = _used_percentage(limit.get("percent"))
|
|
if utilization is None:
|
|
continue
|
|
fable: dict[str, Any] = {"utilization": utilization}
|
|
reset = limit.get("resets_at")
|
|
if isinstance(reset, (int, float, str)) and not isinstance(reset, bool):
|
|
fable["resets_at"] = reset
|
|
payload["seven_day_fable"] = fable
|
|
break
|
|
if not payload:
|
|
raise QuotaNotExposed("Claude usage response omitted supported quota windows")
|
|
return payload
|
|
|
|
|
|
def _read_credentials() -> tuple[dict[str, Any], dict[str, Any]]:
|
|
"""Read one private Claude OAuth document without following a final symlink."""
|
|
try:
|
|
descriptor = os.open(
|
|
CREDENTIALS_FILE,
|
|
os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0),
|
|
)
|
|
except OSError as error:
|
|
raise QuotaNotExposed("Claude scoped quota credential is unavailable") from error
|
|
try:
|
|
info = os.fstat(descriptor)
|
|
if (
|
|
not stat.S_ISREG(info.st_mode)
|
|
or info.st_mode & 0o077
|
|
or not 0 < info.st_size <= MAX_DOCUMENT_BYTES
|
|
):
|
|
raise QuotaNotExposed("Claude scoped quota credential is unsafe")
|
|
raw = os.read(descriptor, info.st_size + 1)
|
|
finally:
|
|
os.close(descriptor)
|
|
try:
|
|
document = json.loads(raw)
|
|
except (UnicodeError, json.JSONDecodeError) as error:
|
|
raise QuotaNotExposed("Claude scoped quota credential is invalid") from error
|
|
credentials = document.get("claudeAiOauth") if isinstance(document, dict) else None
|
|
if not isinstance(credentials, dict):
|
|
raise QuotaNotExposed("Claude scoped quota credential is invalid")
|
|
return document, credentials
|
|
|
|
|
|
def credential_refresh_expiry_timestamp() -> float | None:
|
|
"""Return only the safe refresh-grant expiry from the private document."""
|
|
try:
|
|
_document, credentials = _read_credentials()
|
|
except QuotaNotExposed:
|
|
return None
|
|
expires_at = credentials.get("refreshTokenExpiresAt")
|
|
if (
|
|
isinstance(expires_at, bool)
|
|
or not isinstance(expires_at, (int, float))
|
|
or expires_at <= 0
|
|
):
|
|
return None
|
|
return float(expires_at) / 1000
|
|
|
|
|
|
def _write_credentials(document: dict[str, Any]) -> None:
|
|
"""Atomically persist a provider-rotated OAuth document for Vault sync."""
|
|
encoded = (json.dumps(document, separators=(",", ":")) + "\n").encode("utf-8")
|
|
if len(encoded) > MAX_DOCUMENT_BYTES:
|
|
raise QuotaNotExposed("Claude scoped quota credential is too large")
|
|
temporary = CREDENTIALS_FILE.with_name(
|
|
f".{CREDENTIALS_FILE.name}.{uuid.uuid4().hex}.tmp"
|
|
)
|
|
flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0)
|
|
descriptor = os.open(temporary, flags, 0o600)
|
|
try:
|
|
remaining = memoryview(encoded)
|
|
while remaining:
|
|
remaining = remaining[os.write(descriptor, remaining) :]
|
|
os.fsync(descriptor)
|
|
os.fchmod(descriptor, 0o600)
|
|
finally:
|
|
os.close(descriptor)
|
|
try:
|
|
os.replace(temporary, CREDENTIALS_FILE)
|
|
finally:
|
|
temporary.unlink(missing_ok=True)
|
|
|
|
|
|
def _refresh_access_token(
|
|
document: dict[str, Any], credentials: dict[str, Any]
|
|
) -> str:
|
|
"""Refresh the scoped access token and retain any refresh-token rotation."""
|
|
refresh_token = credentials.get("refreshToken")
|
|
if not isinstance(refresh_token, str) or not refresh_token:
|
|
raise QuotaNotExposed("Claude scoped quota refresh credential is unavailable")
|
|
scopes = credentials.get("scopes")
|
|
if not (
|
|
isinstance(scopes, list)
|
|
and scopes
|
|
and all(isinstance(scope, str) and scope in DEFAULT_SCOPES for scope in scopes)
|
|
):
|
|
scopes = list(DEFAULT_SCOPES)
|
|
body = json.dumps(
|
|
{
|
|
"grant_type": "refresh_token",
|
|
"refresh_token": refresh_token,
|
|
"client_id": OAUTH_CLIENT_ID,
|
|
"scope": " ".join(scopes),
|
|
},
|
|
separators=(",", ":"),
|
|
).encode("utf-8")
|
|
request = Request(
|
|
TOKEN_URL,
|
|
data=body,
|
|
method="POST",
|
|
headers={
|
|
"Content-Type": "application/json",
|
|
"Accept": "application/json",
|
|
"User-Agent": "atlas-ai-usage-exporter/3.0",
|
|
},
|
|
)
|
|
with urlopen(request, timeout=REQUEST_TIMEOUT_SECONDS) as response:
|
|
response_document = json.loads(response.read(MAX_DOCUMENT_BYTES))
|
|
if not isinstance(response_document, dict):
|
|
raise QuotaNotExposed("Claude token refresh response is invalid")
|
|
access_token = response_document.get("access_token")
|
|
expires_in = response_document.get("expires_in")
|
|
rotated_refresh = response_document.get("refresh_token", refresh_token)
|
|
if (
|
|
not isinstance(access_token, str)
|
|
or not access_token
|
|
or isinstance(expires_in, bool)
|
|
or not isinstance(expires_in, (int, float))
|
|
or expires_in <= 0
|
|
or not isinstance(rotated_refresh, str)
|
|
or not rotated_refresh
|
|
):
|
|
raise QuotaNotExposed("Claude token refresh response is invalid")
|
|
credentials["accessToken"] = access_token
|
|
credentials["refreshToken"] = rotated_refresh
|
|
credentials["expiresAt"] = int(time.time() * 1000 + expires_in * 1000)
|
|
response_scopes = response_document.get("scope")
|
|
if isinstance(response_scopes, str) and response_scopes:
|
|
credentials["scopes"] = response_scopes.split()
|
|
refresh_expires_in = response_document.get("refresh_token_expires_in")
|
|
if (
|
|
not isinstance(refresh_expires_in, bool)
|
|
and isinstance(refresh_expires_in, (int, float))
|
|
and refresh_expires_in > 0
|
|
):
|
|
credentials["refreshTokenExpiresAt"] = int(
|
|
time.time() * 1000 + refresh_expires_in * 1000
|
|
)
|
|
_write_credentials(document)
|
|
return access_token
|
|
|
|
|
|
def _access_token(*, force_refresh: bool = False) -> str:
|
|
"""Return a current scoped token, refreshing it before provider expiry."""
|
|
document, credentials = _read_credentials()
|
|
access_token = credentials.get("accessToken")
|
|
expires_at = credentials.get("expiresAt")
|
|
current = time.time() * 1000
|
|
if (
|
|
not force_refresh
|
|
and isinstance(access_token, str)
|
|
and access_token
|
|
and not isinstance(expires_at, bool)
|
|
and isinstance(expires_at, (int, float))
|
|
and expires_at - current > REFRESH_SKEW_MILLISECONDS
|
|
):
|
|
return access_token
|
|
return _refresh_access_token(document, credentials)
|
|
|
|
|
|
def _query_scoped_usage() -> dict[str, Any]:
|
|
"""Fetch Claude Code's authenticated overall and model-scoped quotas."""
|
|
token = _access_token()
|
|
for attempt in range(2):
|
|
request = Request(
|
|
USAGE_URL,
|
|
headers={
|
|
"Authorization": f"Bearer {token}",
|
|
"Accept": "application/json",
|
|
"anthropic-beta": "oauth-2025-04-20",
|
|
"User-Agent": "atlas-ai-usage-exporter/3.0",
|
|
},
|
|
)
|
|
try:
|
|
with urlopen(request, timeout=REQUEST_TIMEOUT_SECONDS) as response:
|
|
return _scoped_payload(json.loads(response.read(MAX_DOCUMENT_BYTES)))
|
|
except HTTPError as error:
|
|
if error.code != 401 or attempt:
|
|
raise
|
|
token = _access_token(force_refresh=True)
|
|
raise QuotaNotExposed("Claude scoped quota request failed")
|
|
|
|
|
|
def _query_header_usage() -> dict[str, Any]:
|
|
"""Fetch unified quotas through the inference-scoped setup-token fallback."""
|
|
token = _setup_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/3.0",
|
|
},
|
|
)
|
|
with urlopen(request, timeout=REQUEST_TIMEOUT_SECONDS) as response:
|
|
payload = _header_payload(response.headers)
|
|
response.read(MAX_DOCUMENT_BYTES)
|
|
return payload
|
|
|
|
|
|
def query_claude() -> dict[str, Any]:
|
|
"""Prefer exact scoped quotas and preserve unified metrics as a fallback."""
|
|
try:
|
|
return _query_scoped_usage()
|
|
except (
|
|
HTTPError,
|
|
URLError,
|
|
TimeoutError,
|
|
OSError,
|
|
UnicodeError,
|
|
json.JSONDecodeError,
|
|
QuotaNotExposed,
|
|
TypeError,
|
|
ValueError,
|
|
):
|
|
return _query_header_usage()
|