"""v3 M2 — OpenAI Codex provider: read rate-limit % from the Codex session rollouts. Codex writes a token_count event per turn whose payload.rate_limits mirrors Claude's model (primary=5h, secondary=weekly). The provider surfaces the freshest such snapshot; a window whose reset time has passed is reported as a fresh 0%. """ import asyncio import json import os import time import httpx from daemon.providers.openai_codex import OpenAICodexProvider def _run(coro): return asyncio.get_event_loop().run_until_complete(coro) def _write_rollout(home, name, rate_limits_list, mtime=None): """Write a rollout JSONL with one token_count event per rate_limits dict.""" d = home / "sessions" / "2026" / "07" / "10" d.mkdir(parents=True, exist_ok=True) p = d / f"rollout-{name}.jsonl" lines = [] for i, rl in enumerate(rate_limits_list): lines.append(json.dumps({ "timestamp": f"2026-07-10T00:00:0{i}Z", "type": "event_msg", "payload": {"type": "token_count", "info": {"total_token_usage": {"total_tokens": 1}}, "rate_limits": rl}, })) p.write_text("\n".join(lines), encoding="utf-8") if mtime is not None: os.utime(p, (mtime, mtime)) return p def _rl(pri_used, pri_reset, sec_used, sec_reset, plan="plus", reached=None): return {"limit_id": "codex", "primary": {"used_percent": pri_used, "window_minutes": 300, "resets_at": pri_reset}, "secondary": {"used_percent": sec_used, "window_minutes": 10080, "resets_at": sec_reset}, "plan_type": plan, "rate_limit_reached_type": reached} def test_maps_windows_with_future_reset(tmp_path): now = time.time() _write_rollout(tmp_path, "a", [_rl(42.0, now + 3600, 10.0, now + 7200)]) st = _run(OpenAICodexProvider(codex_home=str(tmp_path)).poll()) assert st.ok is True assert st.s == 42.0 and 58 <= st.sr <= 60 assert st.w == 10.0 and 118 <= st.wr <= 120 assert st.st == "Plus" def test_past_reset_is_fresh_zero(tmp_path): now = time.time() # primary window already elapsed => 0% fresh, unknown reset; secondary still open _write_rollout(tmp_path, "a", [_rl(99.0, now - 100, 55.0, now + 600)]) st = _run(OpenAICodexProvider(codex_home=str(tmp_path)).poll()) assert st.s == 0.0 and st.sr == -1 assert st.w == 55.0 and 8 <= st.wr <= 10 def test_rate_limit_reached_marks_limited(tmp_path): now = time.time() _write_rollout(tmp_path, "a", [_rl(100.0, now + 60, 80.0, now + 600, reached="primary")]) st = _run(OpenAICodexProvider(codex_home=str(tmp_path)).poll()) assert st.st == "limited" and st.ok is True def test_last_snapshot_in_file_wins(tmp_path): now = time.time() _write_rollout(tmp_path, "a", [_rl(10.0, now + 600, 1.0, now + 600), _rl(73.0, now + 600, 2.0, now + 600)]) st = _run(OpenAICodexProvider(codex_home=str(tmp_path)).poll()) assert st.s == 73.0 # the later event, not the first def test_newest_file_wins(tmp_path): now = time.time() _write_rollout(tmp_path, "old", [_rl(11.0, now + 600, 0.0, now + 600)], mtime=now - 500) _write_rollout(tmp_path, "new", [_rl(88.0, now + 600, 0.0, now + 600)], mtime=now) st = _run(OpenAICodexProvider(codex_home=str(tmp_path)).poll()) assert st.s == 88.0 def test_skips_file_without_rate_limits(tmp_path): now = time.time() # newest file has a token_count with null rate_limits; older file has the data _write_rollout(tmp_path, "older", [_rl(64.0, now + 600, 0.0, now + 600)], mtime=now - 500) p = _write_rollout(tmp_path, "newer", [], mtime=now) p.write_text(json.dumps({ "timestamp": "2026-07-10T01:00:00Z", "type": "event_msg", "payload": {"type": "token_count", "info": {}, "rate_limits": None}, }), encoding="utf-8") os.utime(p, (now, now)) st = _run(OpenAICodexProvider(codex_home=str(tmp_path)).poll()) assert st.ok is True and st.s == 64.0 def test_no_sessions_dir_reports_no_data(tmp_path): st = _run(OpenAICodexProvider(codex_home=str(tmp_path)).poll()) 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