v3 M2+: Codex live usage endpoint (fresh) with rollout fallback

Adopt the live ChatGPT usage API (found via github.com/rygel/AIUsageTracker, MIT) so Codex % is fresh every poll instead of only as fresh as the last local rollout write: GET https://chatgpt.com/backend-api/wham/usage Authorization: Bearer <access_token from $CODEX_HOME/auth.json tokens> ChatGPT-Account-Id: <account_id> -> rate_limit.primary_window / secondary_window {used_percent, reset_at}, plan_type poll() now tries the live endpoint first and falls back to the rollout snapshot when the token is missing/expired/offline (we read the current OAuth token but don't refresh it). Live-first fixes the staleness of the rollout-only approach: verified live s=1% (5h, resets 300m), w=0% (7d, resets 10080m), st="Plus" — a full fresh window vs the old stale past-reset heuristic. +4 tests (live preferred, limit_reached->limited, 401->rollout fallback, no-auth ->rollout). 126 passed, 2 Linux-only failures (baseline). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-07-10 09:02:12 +03:00
parent 4eeb6e242b
commit d1cc4670b1
2 changed files with 166 additions and 12 deletions
@@ -1,18 +1,21 @@
 """OpenAI Codex provider (v3 M2).

-Reads the ChatGPT/Codex rate-limit utilization the Codex CLI already records
-locally, so there's no second auth to maintain: every turn, Codex writes a
-``token_count`` event into its session rollout (``$CODEX_HOME/sessions/**/rollout-*.jsonl``)
-whose ``info.rate_limits`` snapshot mirrors Claude's model — a primary window
-(5h) and a secondary window (weekly), each with a used-percent and a reset time::
+Primary source is the live ChatGPT usage endpoint (fresh every poll)::

-    "rate_limits": {"primary":   {"used_percent": 1.0,  "window_minutes": 300,   "resets_at": 178...},
-                    "secondary": {"used_percent": 0.0,  "window_minutes": 10080, "resets_at": 178...},
-                    "plan_type": "plus", "rate_limit_reached_type": null}
+    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": {...}}}

-We surface the freshest such snapshot. It's a local read (no network, no token),
-so it's current while Codex is in use and goes stale between sessions — when a
-window's ``resets_at`` has already passed we report that window as a fresh 0%.
+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
@@ -23,8 +26,12 @@ 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.
@@ -40,11 +47,18 @@ class OpenAICodexProvider(Provider):
    label = "OpenAI Codex"
    accent = "10a37f"        # brand green

-    def __init__(self, codex_home: str | os.PathLike | None = None) -> None:
+    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)
@@ -55,6 +69,71 @@ class OpenAICodexProvider(Provider):
            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: