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>
This commit is contained in:
@@ -1,18 +1,21 @@
|
|||||||
"""OpenAI Codex provider (v3 M2).
|
"""OpenAI Codex provider (v3 M2).
|
||||||
|
|
||||||
Reads the ChatGPT/Codex rate-limit utilization the Codex CLI already records
|
Primary source is the live ChatGPT usage endpoint (fresh every poll)::
|
||||||
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::
|
|
||||||
|
|
||||||
"rate_limits": {"primary": {"used_percent": 1.0, "window_minutes": 300, "resets_at": 178...},
|
GET https://chatgpt.com/backend-api/wham/usage
|
||||||
"secondary": {"used_percent": 0.0, "window_minutes": 10080, "resets_at": 178...},
|
Authorization: Bearer <access_token from $CODEX_HOME/auth.json tokens>
|
||||||
"plan_type": "plus", "rate_limit_reached_type": null}
|
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),
|
with each window carrying ``used_percent`` + ``reset_at`` — Claude's short (5h) /
|
||||||
so it's current while Codex is in use and goes stale between sessions — when a
|
weekly model. (Endpoint + auth discovered from github.com/rygel/AIUsageTracker,
|
||||||
window's ``resets_at`` has already passed we report that window as a fresh 0%.
|
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
|
from __future__ import annotations
|
||||||
@@ -23,8 +26,12 @@ import os
|
|||||||
import time
|
import time
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
|
import httpx
|
||||||
|
|
||||||
from .base import Provider, ProviderStatus
|
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
|
# 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
|
# 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.
|
# just-opened session with no turns (hence no rate_limits) yet.
|
||||||
@@ -40,11 +47,18 @@ class OpenAICodexProvider(Provider):
|
|||||||
label = "OpenAI Codex"
|
label = "OpenAI Codex"
|
||||||
accent = "10a37f" # brand green
|
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._home = Path(codex_home) if codex_home else _codex_home()
|
||||||
|
self._transport = transport # tests inject an httpx.MockTransport
|
||||||
|
|
||||||
async def poll(self) -> ProviderStatus:
|
async def poll(self) -> ProviderStatus:
|
||||||
from daemon.claude_usage_daemon_windows import log
|
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:
|
try:
|
||||||
snap = await asyncio.to_thread(self._latest_rate_limits)
|
snap = await asyncio.to_thread(self._latest_rate_limits)
|
||||||
except Exception as e: # never break the loop (SC#5)
|
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 ProviderStatus(ok=False, st="no data")
|
||||||
return self._to_status(snap)
|
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 -------------------------------------------------
|
# -- local rollout reading -------------------------------------------------
|
||||||
|
|
||||||
def _latest_rate_limits(self) -> dict | None:
|
def _latest_rate_limits(self) -> dict | None:
|
||||||
|
|||||||
@@ -10,6 +10,8 @@ import json
|
|||||||
import os
|
import os
|
||||||
import time
|
import time
|
||||||
|
|
||||||
|
import httpx
|
||||||
|
|
||||||
from daemon.providers.openai_codex import OpenAICodexProvider
|
from daemon.providers.openai_codex import OpenAICodexProvider
|
||||||
|
|
||||||
|
|
||||||
@@ -103,3 +105,76 @@ def test_skips_file_without_rate_limits(tmp_path):
|
|||||||
def test_no_sessions_dir_reports_no_data(tmp_path):
|
def test_no_sessions_dir_reports_no_data(tmp_path):
|
||||||
st = _run(OpenAICodexProvider(codex_home=str(tmp_path)).poll())
|
st = _run(OpenAICodexProvider(codex_home=str(tmp_path)).poll())
|
||||||
assert st.ok is False and st.st == "no data"
|
assert st.ok is False and st.st == "no data"
|
||||||
|
|
||||||
|
|
||||||
|
# -- live ChatGPT usage endpoint (preferred over rollouts) --------------------
|
||||||
|
|
||||||
|
def _write_auth(home, access_token="tok", account_id="acct"):
|
||||||
|
home.mkdir(parents=True, exist_ok=True)
|
||||||
|
(home / "auth.json").write_text(
|
||||||
|
json.dumps({"tokens": {"access_token": access_token, "account_id": account_id}}),
|
||||||
|
encoding="utf-8")
|
||||||
|
|
||||||
|
|
||||||
|
def _live_resp(pri_pct, pri_reset, sec_pct, sec_reset, plan="plus", limit_reached=False):
|
||||||
|
return httpx.Response(200, json={
|
||||||
|
"plan_type": plan,
|
||||||
|
"rate_limit": {"allowed": True, "limit_reached": limit_reached,
|
||||||
|
"primary_window": {"used_percent": pri_pct, "reset_at": pri_reset},
|
||||||
|
"secondary_window": {"used_percent": sec_pct, "reset_at": sec_reset}},
|
||||||
|
})
|
||||||
|
|
||||||
|
|
||||||
|
def test_live_usage_preferred_over_rollout(tmp_path):
|
||||||
|
_write_auth(tmp_path)
|
||||||
|
# a rollout is also present, but the live endpoint must win
|
||||||
|
now = time.time()
|
||||||
|
_write_rollout(tmp_path, "a", [_rl(99.0, now + 600, 99.0, now + 600)])
|
||||||
|
seen = {}
|
||||||
|
|
||||||
|
def handler(req: httpx.Request) -> httpx.Response:
|
||||||
|
seen["auth"] = req.headers.get("authorization")
|
||||||
|
seen["acct"] = req.headers.get("chatgpt-account-id")
|
||||||
|
seen["url"] = str(req.url)
|
||||||
|
return _live_resp(5, now + 3600, 20, now + 7200)
|
||||||
|
|
||||||
|
p = OpenAICodexProvider(codex_home=str(tmp_path), transport=httpx.MockTransport(handler))
|
||||||
|
st = _run(p.poll())
|
||||||
|
assert st.ok is True and st.st == "Plus"
|
||||||
|
assert st.s == 5.0 and 58 <= st.sr <= 60
|
||||||
|
assert st.w == 20.0 and 118 <= st.wr <= 120
|
||||||
|
assert seen["auth"] == "Bearer tok" and seen["acct"] == "acct"
|
||||||
|
assert seen["url"] == "https://chatgpt.com/backend-api/wham/usage"
|
||||||
|
|
||||||
|
|
||||||
|
def test_live_limit_reached_is_limited(tmp_path):
|
||||||
|
_write_auth(tmp_path)
|
||||||
|
now = time.time()
|
||||||
|
|
||||||
|
def handler(req):
|
||||||
|
return _live_resp(100, now + 60, 50, now + 600, limit_reached=True)
|
||||||
|
|
||||||
|
p = OpenAICodexProvider(codex_home=str(tmp_path), transport=httpx.MockTransport(handler))
|
||||||
|
st = _run(p.poll())
|
||||||
|
assert st.st == "limited" and st.ok is True
|
||||||
|
|
||||||
|
|
||||||
|
def test_falls_back_to_rollout_on_token_expiry(tmp_path):
|
||||||
|
_write_auth(tmp_path)
|
||||||
|
now = time.time()
|
||||||
|
_write_rollout(tmp_path, "a", [_rl(64.0, now + 600, 0.0, now + 600)])
|
||||||
|
|
||||||
|
def handler(req):
|
||||||
|
return httpx.Response(401, json={"detail": "unauthorized"})
|
||||||
|
|
||||||
|
p = OpenAICodexProvider(codex_home=str(tmp_path), transport=httpx.MockTransport(handler))
|
||||||
|
st = _run(p.poll())
|
||||||
|
assert st.ok is True and st.s == 64.0 # from the rollout snapshot
|
||||||
|
|
||||||
|
|
||||||
|
def test_no_auth_uses_rollout_without_network(tmp_path):
|
||||||
|
now = time.time()
|
||||||
|
_write_rollout(tmp_path, "a", [_rl(30.0, now + 600, 0.0, now + 600)])
|
||||||
|
# no auth.json => live path skipped entirely; no transport needed
|
||||||
|
st = _run(OpenAICodexProvider(codex_home=str(tmp_path)).poll())
|
||||||
|
assert st.ok is True and st.s == 30.0
|
||||||
|
|||||||
Reference in New Issue
Block a user