clawdmeter/daemon/tests/test_openai_codex.py

"""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