codex_reset_tracker.py

#!/usr/bin/env python3
import argparse
import json
import os
import urllib.error
import urllib.request
from datetime import datetime, timezone
from pathlib import Path


BASE_URL = "https://chatgpt.com/backend-api"
AUTH_JSON = Path.home() / ".codex" / "auth.json"


def load_json(path):
    try:
        with open(path, "r", encoding="utf-8") as f:
            return json.load(f)
    except OSError as e:
        raise SystemExit(f"Could not read {path}: {e.strerror}") from e
    except json.JSONDecodeError as e:
        raise SystemExit(f"Invalid JSON in {path}: {e}") from e


def parse_time(value):
    if value in (None, ""):
        return None
    if isinstance(value, (int, float)):
        try:
            return datetime.fromtimestamp(value / 1000 if value > 10_000_000_000 else value, timezone.utc)
        except (OSError, OverflowError, ValueError):
            return None
    if isinstance(value, str):
        text = value.replace("Z", "+00:00")
        try:
            dt = datetime.fromisoformat(text)
            return dt if dt.tzinfo else dt.replace(tzinfo=timezone.utc)
        except ValueError:
            return None
    return None


def display_zone(name):
    if name.lower() == "local":
        return datetime.now().astimezone().tzinfo
    if name.upper() == "UTC":
        return timezone.utc
    try:
        from zoneinfo import ZoneInfo

        return ZoneInfo(name)
    except Exception as e:
        raise SystemExit(f"Unknown timezone {name!r}. Use UTC, local, or an IANA name like America/Vancouver.") from e


def fmt(dt, tz):
    return "unknown" if dt is None else dt.astimezone(tz).strftime("%Y-%m-%d %H:%M:%S %Z")


def fmt_reset(dt, tz, now):
    if dt is None:
        return "unknown"
    local = dt.astimezone(tz)
    if local.date() == now.astimezone(tz).date():
        return local.strftime("%I:%M %p").lstrip("0")
    return local.strftime("%b %d %I:%M %p").replace(" 0", " ").lstrip("0")


def normalize_credits(data, now):
    nested = data.get("rate_limit_reset_credits") if isinstance(data, dict) else None
    payload = nested if isinstance(nested, dict) else data if isinstance(data, dict) else {}
    credits = payload.get("credits") if isinstance(payload.get("credits"), list) else []
    rows = []
    for credit in credits:
        if not isinstance(credit, dict):
            continue
        expires = parse_time(credit.get("expires_at"))
        granted = parse_time(credit.get("granted_at"))
        rows.append({
            "title": credit.get("title") or "Usage reset",
            "status": credit.get("status") or "unknown",
            "granted_at": granted,
            "expires_at": expires,
            "days_left": None if expires is None else (expires - now).total_seconds() / 86400,
        })
    rows.sort(key=lambda r: (r["status"] != "available", r["expires_at"] is None, r["expires_at"] or datetime.max.replace(tzinfo=timezone.utc)))
    try:
        available_count = int(payload.get("available_count") or 0)
    except (TypeError, ValueError):
        available_count = 0
    return rows, available_count


def visible_credits(rows):
    return [r for r in rows if r["status"] == "available"]


def hidden_credit_counts(rows):
    counts = {}
    for row in rows:
        if row["status"] == "available":
            continue
        key = "redeemed" if row["status"] in {"redeemed", "used"} else row["status"]
        counts[key] = counts.get(key, 0) + 1
    return counts


def duration_label(seconds):
    if seconds in (None, ""):
        return "Usage"
    try:
        minutes = float(seconds) / 60
    except (TypeError, ValueError):
        return "Usage"
    if minutes >= 10079:
        weeks = max(1, round(minutes / (7 * 1440)))
        return "Weekly" if weeks == 1 else f"{weeks} weeks"
    if minutes >= 1439:
        return f"{round(minutes / 1440)}d"
    if minutes >= 60:
        return f"{round(minutes / 60)}h"
    return f"{max(1, round(minutes))}m"


def usage_windows(data):
    rows = []

    def add_windows(group, rate_limit):
        if not isinstance(rate_limit, dict):
            return
        for key in ("primary_window", "secondary_window"):
            window = rate_limit.get(key)
            if not isinstance(window, dict):
                continue
            used = window.get("used_percent")
            if not isinstance(used, (int, float)):
                continue
            rows.append({
                "group": group,
                "label": duration_label(window.get("limit_window_seconds")),
                "remaining_percent": max(0, min(100, 100 - used)),
                "used_percent": used,
                "resets_at": parse_time(window.get("reset_at")),
                "blocked": used >= 100 or rate_limit.get("limit_reached") is True or rate_limit.get("allowed") is False,
            })

    if isinstance(data, dict):
        add_windows("Usage remaining", data.get("rate_limit"))
        for item in data.get("additional_rate_limits") or []:
            if isinstance(item, dict):
                add_windows(item.get("limit_name") or "Additional limit", item.get("rate_limit"))
    return rows


def usage_credits(data):
    credits = data.get("credits") if isinstance(data, dict) else None
    if not isinstance(credits, dict):
        return None
    if credits.get("unlimited") is True:
        return "unlimited"
    balance = credits.get("balance")
    if isinstance(balance, (int, float)):
        return f"{balance:g}"
    if credits.get("has_credits") is False:
        return "0"
    return None


def plan_name(data):
    plan = data.get("plan_type") if isinstance(data, dict) else None
    names = {
        "free": "Free",
        "go": "Go",
        "plus": "Plus",
        "pro": "Pro",
        "team": "Team",
        "self_serve_business_usage_based": "Team usage-based",
        "business": "Business",
        "enterprise": "Enterprise",
        "enterprise_cbp_usage_based": "Enterprise usage-based",
        "enterprise_cbp_automation": "Enterprise automation",
        "edu": "Education",
        "education": "Education",
        "hc": "Enterprise",
    }
    return names.get(plan, plan)


def load_auth(auth_json):
    token = os.environ.get("CODEX_BEARER_TOKEN")
    account_id = os.environ.get("CODEX_ACCOUNT_ID")
    if token:
        return token, account_id
    path = Path(auth_json).expanduser()
    if not path.exists():
        return None, None
    data = load_json(path)
    tokens = data.get("tokens") if isinstance(data, dict) else {}
    if not isinstance(tokens, dict):
        tokens = {}
    return tokens.get("access_token"), tokens.get("account_id")


def request_json(path, auth_json, timeout):
    token, account_id = load_auth(auth_json)
    if not token:
        raise SystemExit(f"Not logged in. Run Codex first, or set CODEX_BEARER_TOKEN. Looked for {auth_json}.")
    headers = {
        "Accept": "application/json",
        "OpenAI-Beta": "codex-1",
        "OAI-Language": "en",
        "User-Agent": "codex-reset-tracker/0.1",
        "originator": "codex-reset-tracker",
        "Authorization": f"Bearer {token}",
    }
    if account_id:
        headers["ChatGPT-Account-ID"] = account_id
    req = urllib.request.Request(BASE_URL + path, headers=headers)
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            try:
                return json.loads(resp.read().decode("utf-8"))
            except json.JSONDecodeError as e:
                raise SystemExit(f"Invalid JSON from {path}: {e}") from e
    except urllib.error.HTTPError as e:
        raise SystemExit(f"HTTP {e.code} from {path}") from e
    except urllib.error.URLError as e:
        raise SystemExit(f"Network error from {path}: {e}") from e


def json_default(value):
    if isinstance(value, datetime):
        return value.isoformat()
    raise TypeError(f"Object of type {type(value).__name__} is not JSON serializable")


def print_report(rows, available_count, usage_rows, credits_remaining, plan, tz, warn_days, now):
    print("Codex Reset Tracker")
    if plan:
        print(f"Plan: {plan}")
    print(f"Available resets: {available_count}")
    if credits_remaining is not None:
        print(f"Usage credits: {credits_remaining}")
    print()
    print("Usage remaining:")
    if usage_rows:
        for row in usage_rows:
            suffix = " blocked" if row["blocked"] else ""
            group = "" if row["group"] == "Usage remaining" else f"{row['group']} "
            print(f"- {group}{row['label']}: {row['remaining_percent']:.0f}% remaining, resets {fmt_reset(row['resets_at'], tz, now)}{suffix}")
    else:
        print("- none listed by the API")
    shown_rows = visible_credits(rows)
    hidden_counts = hidden_credit_counts(rows)
    if not shown_rows:
        print()
        print("Available reset credits: none listed by the API")
        if hidden_counts:
            print("Hidden credits: " + ", ".join(f"{k} {v}" for k, v in sorted(hidden_counts.items())))
        return
    print()
    print("Available reset credits:")
    if hidden_counts:
        print("Hidden credits: " + ", ".join(f"{k} {v}" for k, v in sorted(hidden_counts.items())))
    for i, row in enumerate(shown_rows, 1):
        days = row["days_left"]
        left = "unknown" if days is None else f"{days:.1f} days"
        marker = ""
        if days is not None and row["status"] == "available":
            if days < 0:
                marker = " expired"
            elif days <= warn_days:
                marker = " expiring-soon"
        print(f"{i}. {row['title']} [{row['status']}]{marker}")
        print(f"   granted: {fmt(row['granted_at'], tz)}")
        print(f"   expires: {fmt(row['expires_at'], tz)}")
        print(f"   time left: {left}")


def main(argv=None):
    parser = argparse.ArgumentParser(description="Show Codex usage reset credits and rate limit reset expiry dates.")
    parser.add_argument("--auth-json", default=str(AUTH_JSON), help="Path to Codex auth.json.")
    parser.add_argument("--credits-json", help="Use a saved /wham/rate-limit-reset-credits response.")
    parser.add_argument("--usage-json", help="Use a saved /wham/usage response.")
    parser.add_argument("--timezone", default="local", help="local, UTC, or an IANA timezone.")
    parser.add_argument("--warn-days", type=float, default=3)
    parser.add_argument("--timeout", type=float, default=30)
    parser.add_argument("--json", action="store_true", help="Print normalized JSON.")
    args = parser.parse_args(argv)

    credits = load_json(args.credits_json) if args.credits_json else request_json("/wham/rate-limit-reset-credits", args.auth_json, args.timeout)
    usage = load_json(args.usage_json) if args.usage_json else request_json("/wham/usage", args.auth_json, args.timeout)

    now = datetime.now(timezone.utc)
    rows, available_count = normalize_credits(credits, now)
    if args.json:
        print(json.dumps({
            "available_count": available_count,
            "hidden_credits": hidden_credit_counts(rows),
            "plan": plan_name(usage),
            "usage_credits": usage_credits(usage),
            "usage": usage_windows(usage),
            "credits": visible_credits(rows),
        }, default=json_default, indent=2, sort_keys=True))
        return 0

    print_report(rows, available_count, usage_windows(usage), usage_credits(usage), plan_name(usage), display_zone(args.timezone), args.warn_days, now)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())