dotfiles/hyprland/.config/waybar/scripts/openai-rate.py

#!/usr/bin/env python3
import json
import math
import os
import sys
import urllib.error
import urllib.request
from argparse import ArgumentParser
from pathlib import Path


DEFAULT_BASE_URL = "https://chatgpt.com/backend-api"


def waybar(text, klass="good", tooltip=None):
    print(json.dumps({"text": text, "class": klass, "tooltip": tooltip or text}))
    sys.exit(0)


def load_auth():
    auth_path = Path(os.environ.get("OPENAI_AUTH_FILE", "~/.codex/auth.json")).expanduser()
    auth = {}
    if auth_path.is_file():
        try:
            auth = json.loads(auth_path.read_text())
        except (OSError, json.JSONDecodeError):
            waybar("AI auth unreadable", "critical", f"Could not parse {auth_path}")

    tokens = auth.get("tokens") or {}
    access_token = tokens.get("access_token")
    account_id = tokens.get("account_id")
    if access_token:
        return access_token, account_id, "~/.codex/auth.json:tokens.access_token"

    api_key = auth.get("OPENAI_API_KEY") or os.environ.get("OPENAI_API_KEY")
    if api_key:
        source = "~/.codex/auth.json:OPENAI_API_KEY" if auth.get("OPENAI_API_KEY") else "env:OPENAI_API_KEY"
        return api_key, account_id, source

    waybar(
        "AI auth missing",
        "warning",
        "No Codex access token found in ~/.codex/auth.json",
    )


def usage_url():
    base = os.environ.get("CODEX_RATE_BASE_URL") or os.environ.get("OPENAI_BASE_URL") or DEFAULT_BASE_URL
    base = base.rstrip("/")
    if base.endswith("/backend-api/codex"):
        base = base.removesuffix("/codex")
    if base.startswith(("https://chatgpt.com", "https://chat.openai.com")) and "/backend-api" not in base:
        base += "/backend-api"
    return f"{base}/wham/usage" if "/backend-api" in base else f"{base}/api/codex/usage"


def fetch_usage(token, account_id):
    headers = {
        "Authorization": f"Bearer {token}",
        "Accept": "application/json",
        "User-Agent": "codex-waybar-rate",
    }
    if account_id:
        headers["ChatGPT-Account-ID"] = account_id

    req = urllib.request.Request(usage_url(), headers=headers)
    try:
        with urllib.request.urlopen(req, timeout=15) as res:
            return json.load(res)
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode("utf-8", "replace")[:240]
        waybar("AI rate failed", "critical", f"Usage endpoint returned HTTP {exc.code}\n{detail}")
    except (OSError, TimeoutError) as exc:
        waybar("AI rate failed", "critical", f"Could not reach Codex usage endpoint\n{exc}")
    except json.JSONDecodeError as exc:
        waybar("AI rate unreadable", "critical", f"Usage endpoint returned invalid JSON\n{exc}")


def window_label(seconds):
    if not isinstance(seconds, int) or seconds <= 0:
        return "usage"

    minutes = math.ceil(seconds / 60)
    if minutes < 90:
        return f"{minutes}m"

    hours = round(minutes / 60)
    if hours < 36:
        return f"{hours}h"

    days = round(hours / 24)
    if days < 365:
        return f"{days}d"

    return f"{round(days / 365)}y"


def reset_text(seconds):
    if not isinstance(seconds, int) or seconds < 0:
        return None
    if seconds < 90:
        return f"{seconds}s"
    minutes = round(seconds / 60)
    if minutes < 90:
        return f"{minutes}m"
    hours = round(minutes / 60)
    if hours < 48:
        return f"{hours}h"
    return f"{round(hours / 24)}d"


def usage_class(used, limit_reached=False):
    if limit_reached or used >= 95:
        return "critical"
    if used >= 90:
        return "warning"
    if used >= 75:
        return "regular"
    return "good"


def format_window(name, window):
    if not isinstance(window, dict):
        return None

    used = window.get("used_percent")
    if not isinstance(used, int):
        return None

    label = window_label(window.get("limit_window_seconds"))
    left = max(0, 100 - used)
    reset = reset_text(window.get("reset_after_seconds"))
    tooltip = f"{name} {label}: {used}% used"
    if reset:
        tooltip += f", resets in {reset}"
    return {
        "label": label,
        "left": left,
        "text": f"{left}% {label}",
        "used": used,
        "tooltip": tooltip,
    }


def credit_text(credits):
    if not isinstance(credits, dict) or not credits.get("has_credits"):
        return None
    if credits.get("unlimited"):
        return "credits unlimited"
    balance = credits.get("balance")
    if balance is None:
        return "credits enabled"
    try:
        return f"credits {round(float(balance))}"
    except (TypeError, ValueError):
        return f"credits {balance}"


def spend_control_text(spend_control):
    limit = (spend_control or {}).get("individual_limit")
    if not isinstance(limit, dict):
        return None

    remaining = limit.get("remaining_percent")
    used = limit.get("used")
    cap = limit.get("limit")
    reset = reset_text(limit.get("reset_after_seconds"))
    if not isinstance(remaining, int):
        return None

    text = f"monthly credits: {100 - remaining}% used"
    if used and cap:
        text += f" ({used}/{cap})"
    if reset:
        text += f", resets in {reset}"
    return text


def selected_window_output(data, source, selector):
    rate_limit = data.get("rate_limit") or {}
    window = format_window(selector, rate_limit.get(f"{selector}_window"))
    if not window:
        waybar("n/a", "warning", f"No {selector} usage window returned\nsource: {source}")

    klass = usage_class(window["used"], rate_limit.get("limit_reached"))
    tooltip = [window["tooltip"], f"{window['left']}% remaining", f"source: {source}"]
    waybar(f"{window['left']}% {window['label']}", klass, "\n".join(tooltip))


def combined_output(data, source):
    rate_limit = data.get("rate_limit") or {}
    windows = [
        format_window("primary", rate_limit.get("primary_window")),
        format_window("secondary", rate_limit.get("secondary_window")),
    ]
    windows = [window for window in windows if window]

    parts = [window["text"] for window in windows]
    credits = credit_text(data.get("credits"))
    if credits:
        parts.append(credits)

    if not parts:
        waybar("AI n/a", "warning", f"No displayable usage data\nsource: {source}")

    worst_used = max(window["used"] for window in windows) if windows else 0
    klass = usage_class(worst_used, rate_limit.get("limit_reached"))

    tooltip = [window["tooltip"] for window in windows]
    spend = spend_control_text(data.get("spend_control"))
    if credits:
        tooltip.append(credits)
    if spend:
        tooltip.append(spend)
    tooltip.append(f"source: {source}")

    waybar(f"AI {'  '.join(parts)}", klass, "\n".join(tooltip))


def parse_args():
    parser = ArgumentParser(description="Waybar Codex rate-limit widget")
    parser.add_argument(
        "--window",
        choices=("primary", "secondary", "all"),
        default="all",
        help="usage window to display",
    )
    return parser.parse_args()


def main():
    args = parse_args()
    token, account_id, source = load_auth()
    data = fetch_usage(token, account_id)

    if args.window == "all":
        combined_output(data, source)
    else:
        selected_window_output(data, source, args.window)


if __name__ == "__main__":
    main()