clawdmeter/daemon/providers/openai_codex.py

"""OpenAI Codex provider (v3 M2).

Primary source is the live ChatGPT usage endpoint (fresh every poll)::

    GET https://chatgpt.com/backend-api/wham/usage
    Authorization: Bearer <access_token from $CODEX_HOME/auth.json tokens>
    ChatGPT-Account-Id: <account_id>            # when present
    -> {"plan_type": "...", "rate_limit": {"primary_window": {...}, "secondary_window": {...}}}

with each window carrying ``used_percent`` + ``reset_at`` — Claude's short (5h) /
weekly model. (Endpoint + auth discovered from github.com/rygel/AIUsageTracker,
MIT.) The access token is Codex's own OAuth token; we read the current one and do
NOT refresh it, so when it has expired (Codex not run in a while) we fall back to
the local session-rollout snapshot Codex writes every turn
(``$CODEX_HOME/sessions/**/rollout-*.jsonl`` -> ``payload.rate_limits``). The
rollout read needs no token but goes stale between sessions; a window whose
``resets_at`` has passed is reported as a fresh 0%. Live-first + rollout-fallback
gives fresh data when possible and last-known otherwise.
"""

from __future__ import annotations

import asyncio
import json
import os
import time
from pathlib import Path

import httpx

from .base import Provider, ProviderStatus

USAGE_URL = "https://chatgpt.com/backend-api/wham/usage"

# Newest-first cap: the active session's rollout has the latest reading, so we
# rarely look past the first file — but scan a few in case the newest is a
# just-opened session with no turns (hence no rate_limits) yet.
_MAX_ROLLOUTS_SCANNED = 12


def _codex_home() -> Path:
    return Path(os.environ.get("CODEX_HOME") or (Path.home() / ".codex"))


class OpenAICodexProvider(Provider):
    id = "openai"
    label = "OpenAI Codex"
    accent = "10a37f"        # brand green

    def __init__(self, codex_home: str | os.PathLike | None = None,
                 transport: httpx.BaseTransport | None = None) -> None:
        self._home = Path(codex_home) if codex_home else _codex_home()
        self._transport = transport   # tests inject an httpx.MockTransport

    async def poll(self) -> ProviderStatus:
        from daemon.claude_usage_daemon_windows import log
        # 1) live endpoint (fresh) — None if no token / expired / offline.
        live = await self._live_usage(log)
        if live is not None:
            return live
        # 2) fall back to the local rollout snapshot (no token needed).
        try:
            snap = await asyncio.to_thread(self._latest_rate_limits)
        except Exception as e:                      # never break the loop (SC#5)
            log(f"Codex poll error: {e!r}")
            return ProviderStatus(ok=False, st="error")
        if snap is None:
            # Configured but nothing to read (Codex never run / logged out).
            return ProviderStatus(ok=False, st="no data")
        return self._to_status(snap)

    # -- live usage endpoint ---------------------------------------------------

    def _read_auth(self) -> dict | None:
        """access_token (+ account_id) from Codex's auth.json, or None."""
        try:
            data = json.loads((self._home / "auth.json").read_text(encoding="utf-8"))
        except (OSError, ValueError):
            return None
        toks = data.get("tokens") or {}
        at = toks.get("access_token")
        return {"access_token": at, "account_id": toks.get("account_id")} if at else None

    async def _live_usage(self, log) -> ProviderStatus | None:
        creds = await asyncio.to_thread(self._read_auth)
        if not creds:
            return None
        headers = {"Authorization": f"Bearer {creds['access_token']}",
                   "Content-Type": "application/json"}
        if creds.get("account_id"):
            headers["ChatGPT-Account-Id"] = creds["account_id"]
        client_kw = {"timeout": 20.0}
        if self._transport is not None:
            client_kw["transport"] = self._transport
        try:
            async with httpx.AsyncClient(**client_kw) as http:
                resp = await http.get(USAGE_URL, headers=headers)
        except httpx.HTTPError as e:
            log(f"Codex live usage failed ({e}); using local snapshot")
            return None
        if resp.status_code >= 400:
            # 401/403 => token expired (Codex refreshes it when it runs); other
            # 4xx/5xx are transient. Either way, fall back to the rollout read.
            log(f"Codex usage HTTP {resp.status_code}; using local snapshot")
            return None
        try:
            root = resp.json()
        except ValueError:
            return None
        rl = root.get("rate_limit")
        return self._status_from_live(root, rl) if isinstance(rl, dict) else None

    @staticmethod
    def _status_from_live(root: dict, rl: dict) -> ProviderStatus:
        now = time.time()

        def window(w: dict | None) -> tuple[float, int]:
            if not isinstance(w, dict):
                return 0.0, -1
            used = float(w.get("used_percent") or 0.0)
            ra = w.get("reset_at")
            if isinstance(ra, (int, float)):
                m = (ra - now) / 60.0
                return used, (int(round(m)) if m > 0 else -1)
            ras = w.get("reset_after_seconds")
            if isinstance(ras, (int, float)) and ras > 0:
                return used, int(ras // 60)
            return used, -1

        s, sr = window(rl.get("primary_window"))
        w, wr = window(rl.get("secondary_window"))
        reached = rl.get("limit_reached") or root.get("rate_limit_reached_type")
        plan = root.get("plan_type")
        st = "limited" if reached else (str(plan).capitalize() if plan else "allowed")
        return ProviderStatus(s=s, sr=sr, w=w, wr=wr, st=st, ok=True)

    # -- local rollout reading -------------------------------------------------

    def _latest_rate_limits(self) -> dict | None:
        """The freshest non-empty ``rate_limits`` object from the most recently
        written session rollout. Pure/sync — run via asyncio.to_thread."""
        sessions = self._home / "sessions"
        if not sessions.is_dir():
            return None
        try:
            files = sorted(sessions.rglob("rollout-*.jsonl"),
                           key=lambda p: p.stat().st_mtime, reverse=True)
        except OSError:
            return None
        for path in files[:_MAX_ROLLOUTS_SCANNED]:
            found = None
            try:
                with path.open(encoding="utf-8", errors="replace") as fh:
                    for line in fh:
                        if '"rate_limits"' not in line:
                            continue
                        try:
                            obj = json.loads(line)
                        except ValueError:
                            continue
                        # token_count events carry rate_limits as a sibling of
                        # info under payload; tolerate the info-nested shape too.
                        payload = obj.get("payload") or {}
                        rl = payload.get("rate_limits")
                        if not isinstance(rl, dict):
                            rl = (payload.get("info") or {}).get("rate_limits")
                        if isinstance(rl, dict) and (rl.get("primary") or rl.get("secondary")):
                            found = rl          # keep the LAST one in the file
            except OSError:
                continue
            if found is not None:
                return found                    # newest file that has a reading wins
        return None

    # -- mapping ---------------------------------------------------------------

    @staticmethod
    def _window(w: dict | None) -> tuple[float, int]:
        """(used_percent 0-100, minutes-to-reset) for one window. A reset time
        already in the past means the window rolled over since the snapshot, so
        report it as a fresh 0% with an unknown reset."""
        if not isinstance(w, dict):
            return 0.0, -1
        used = float(w.get("used_percent") or 0.0)
        resets_at = w.get("resets_at")
        if isinstance(resets_at, (int, float)):
            delta = resets_at - time.time()
            if delta <= 0:
                return 0.0, -1
            return used, int(delta // 60)
        return used, -1

    def _to_status(self, rl: dict) -> ProviderStatus:
        s, sr = self._window(rl.get("primary"))
        w, wr = self._window(rl.get("secondary"))
        reached = rl.get("rate_limit_reached_type")
        plan = rl.get("plan_type")
        if reached:
            st = "limited"
        elif plan:
            st = str(plan).capitalize()         # "Plus" / "Pro" / "Team" / ...
        else:
            st = "allowed"
        # Subscription plan → no per-token cost; the two rate-limit bars are the
        # metric (matches the locked v3 decision). tokens/cost left at 0.
        return ProviderStatus(s=s, sr=sr, w=w, wr=wr, st=st, ok=True)