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

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