init

bin/codex-waybar

@@ -0,0 +1,311 @@
+#!/usr/bin/env python3
+"""Waybar custom module for Codex account rate limits."""
+
+from __future__ import annotations
+
+import argparse
+import html
+import json
+import os
+import selectors
+import subprocess
+import sys
+import time
+from datetime import datetime
+from pathlib import Path
+from string import Formatter
+from typing import Any
+
+FIVE_HOUR_MINS = 300
+WEEKLY_MINS = 10_080
+DEFAULT_FORMAT = "5H {five_hour_used}% · W {weekly_used}% · {five_hour_reset_in}"
+DEFAULT_TOOLTIP = (
+    "Codex ({plan})\n"
+    "5-hour: {five_hour_used}% used, resets {five_hour_reset_at} ({five_hour_reset_in})\n"
+    "Weekly: {weekly_used}% used, resets {weekly_reset_at} ({weekly_reset_in})"
+)
+ALLOWED_FIELDS = {
+    "five_hour_used",
+    "weekly_used",
+    "five_hour_reset_in",
+    "weekly_reset_in",
+    "five_hour_reset_at",
+    "weekly_reset_at",
+    "plan",
+}
+
+
+class UsageError(RuntimeError):
+    pass
+
+
+def cache_path() -> Path:
+    root = Path(os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache"))
+    return root / "codex-waybar" / "rate-limits.json"
+
+
+def parse_json_lines(output: str, response_id: int = 2) -> dict[str, Any]:
+    for line in output.splitlines():
+        try:
+            message = json.loads(line)
+        except json.JSONDecodeError:
+            continue
+        if message.get("id") == response_id:
+            if "error" in message:
+                raise UsageError(str(message["error"]))
+            result = message.get("result")
+            if isinstance(result, dict):
+                return result
+            raise UsageError("Codex returned an invalid rate-limit result")
+    raise UsageError("Codex did not return rate-limit data")
+
+
+def fetch_rate_limits(timeout: float, codex_command: str = "codex") -> dict[str, Any]:
+    requests = [
+        {
+            "id": 1,
+            "method": "initialize",
+            "params": {
+                "clientInfo": {"name": "codex-waybar", "version": "1.0.0"},
+            },
+        },
+        {"method": "initialized"},
+        {"id": 2, "method": "account/rateLimits/read"},
+    ]
+    payload = "\n".join(json.dumps(item, separators=(",", ":")) for item in requests) + "\n"
+    process = None
+    try:
+        process = subprocess.Popen(
+            [codex_command, "app-server"],
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+            bufsize=1,
+        )
+        assert process.stdin is not None
+        assert process.stdout is not None
+        process.stdin.write(payload)
+        process.stdin.flush()
+
+        deadline = time.monotonic() + timeout
+        selector = selectors.DefaultSelector()
+        selector.register(process.stdout, selectors.EVENT_READ)
+        while time.monotonic() < deadline:
+            events = selector.select(deadline - time.monotonic())
+            if not events:
+                break
+            line = process.stdout.readline()
+            if not line:
+                break
+            try:
+                message = json.loads(line)
+            except json.JSONDecodeError:
+                continue
+            if message.get("id") == 2:
+                return parse_json_lines(line)
+        raise UsageError("Codex app-server timed out or returned no rate-limit data")
+    except FileNotFoundError as exc:
+        raise UsageError(f"{codex_command!r} was not found") from exc
+    finally:
+        if process is not None:
+            if process.stdin:
+                try:
+                    process.stdin.close()
+                except OSError:
+                    pass
+            if process.poll() is None:
+                process.terminate()
+            try:
+                process.wait(timeout=2)
+            except subprocess.TimeoutExpired:
+                process.kill()
+                process.wait()
+
+
+def select_snapshot(result: dict[str, Any]) -> dict[str, Any]:
+    by_id = result.get("rateLimitsByLimitId")
+    if isinstance(by_id, dict):
+        codex = by_id.get("codex")
+        if isinstance(codex, dict):
+            return codex
+        for value in by_id.values():
+            if isinstance(value, dict):
+                return value
+    snapshot = result.get("rateLimits")
+    if isinstance(snapshot, dict):
+        return snapshot
+    raise UsageError("No Codex rate-limit snapshot is available")
+
+
+def select_windows(snapshot: dict[str, Any]) -> tuple[dict[str, Any], dict[str, Any]]:
+    windows = [
+        window
+        for key in ("primary", "secondary")
+        if isinstance((window := snapshot.get(key)), dict)
+    ]
+    by_duration = {window.get("windowDurationMins"): window for window in windows}
+    five_hour = by_duration.get(FIVE_HOUR_MINS)
+    weekly = by_duration.get(WEEKLY_MINS)
+    if not isinstance(five_hour, dict) or not isinstance(weekly, dict):
+        raise UsageError("The 5-hour or weekly Codex limit window is unavailable")
+    return five_hour, weekly
+
+
+def load_cache(path: Path) -> dict[str, Any] | None:
+    try:
+        value = json.loads(path.read_text())
+    except (OSError, json.JSONDecodeError):
+        return None
+    return value if isinstance(value, dict) else None
+
+
+def save_cache(path: Path, result: dict[str, Any]) -> None:
+    try:
+        path.parent.mkdir(parents=True, exist_ok=True)
+        temporary = path.with_suffix(".tmp")
+        temporary.write_text(json.dumps({"saved_at": int(time.time()), "result": result}))
+        temporary.replace(path)
+    except OSError:
+        pass
+
+
+def human_countdown(resets_at: Any, now: int) -> str:
+    if not isinstance(resets_at, int):
+        return "unknown"
+    remaining = max(0, resets_at - now)
+    days, remainder = divmod(remaining, 86_400)
+    hours, remainder = divmod(remainder, 3_600)
+    minutes = (remainder + 59) // 60
+    if minutes == 60:
+        hours += 1
+        minutes = 0
+    if hours == 24:
+        days += 1
+        hours = 0
+    if days:
+        return f"{days}d {hours}h"
+    if hours:
+        return f"{hours}h {minutes:02d}m"
+    return f"{minutes}m"
+
+
+def reset_time(resets_at: Any) -> str:
+    if not isinstance(resets_at, int):
+        return "unknown"
+    return datetime.fromtimestamp(resets_at).astimezone().strftime("%Y-%m-%d %H:%M")
+
+
+def validate_template(template: str) -> None:
+    try:
+        fields = {field for _, field, _, _ in Formatter().parse(template) if field}
+    except ValueError as exc:
+        raise UsageError(f"Invalid format template: {exc}") from exc
+    unknown = fields - ALLOWED_FIELDS
+    if unknown:
+        raise UsageError(f"Unknown format field(s): {', '.join(sorted(unknown))}")
+
+
+def format_values(snapshot: dict[str, Any], now: int) -> tuple[dict[str, Any], int]:
+    five_hour, weekly = select_windows(snapshot)
+    five_used = int(five_hour.get("usedPercent", 0))
+    weekly_used = int(weekly.get("usedPercent", 0))
+    values = {
+        "five_hour_used": five_used,
+        "weekly_used": weekly_used,
+        "five_hour_reset_in": human_countdown(five_hour.get("resetsAt"), now),
+        "weekly_reset_in": human_countdown(weekly.get("resetsAt"), now),
+        "five_hour_reset_at": reset_time(five_hour.get("resetsAt")),
+        "weekly_reset_at": reset_time(weekly.get("resetsAt")),
+        "plan": snapshot.get("planType") or "unknown",
+    }
+    return values, max(five_used, weekly_used)
+
+
+def output_json(
+    result: dict[str, Any],
+    text_format: str,
+    tooltip_format: str,
+    warning: int,
+    critical: int,
+    stale: bool = False,
+    stale_reason: str | None = None,
+    now: int | None = None,
+) -> dict[str, Any]:
+    snapshot = select_snapshot(result)
+    values, percentage = format_values(snapshot, now or int(time.time()))
+    classes = ["critical" if percentage >= critical else "warning" if percentage >= warning else "normal"]
+    if stale:
+        classes.append("stale")
+    tooltip = tooltip_format.format_map(values)
+    if stale:
+        tooltip += f"\n\nCached data: {stale_reason or 'refresh failed'}"
+    return {
+        "text": html.escape(text_format.format_map(values)),
+        "tooltip": html.escape(tooltip),
+        "class": classes,
+        "percentage": percentage,
+    }
+
+
+def unavailable_json(reason: str) -> dict[str, Any]:
+    return {
+        "text": "Codex ?",
+        "tooltip": html.escape(f"Codex usage unavailable\n{reason}"),
+        "class": ["unavailable"],
+        "percentage": 0,
+    }
+
+
+def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--format", default=DEFAULT_FORMAT, help="Waybar text template")
+    parser.add_argument("--tooltip-format", default=DEFAULT_TOOLTIP, help="Waybar tooltip template")
+    parser.add_argument("--warning", type=int, default=70, help="Warning threshold percentage")
+    parser.add_argument("--critical", type=int, default=90, help="Critical threshold percentage")
+    parser.add_argument("--timeout", type=float, default=10, help="Codex app-server timeout in seconds")
+    parser.add_argument("--codex-command", default="codex", help=argparse.SUPPRESS)
+    parser.add_argument("--no-cache", action="store_true", help="Disable cached fallback data")
+    return parser.parse_args(argv)
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = parse_args(argv)
+    try:
+        validate_template(args.format)
+        validate_template(args.tooltip_format)
+        if not 0 <= args.warning <= args.critical <= 100:
+            raise UsageError("thresholds must satisfy 0 <= warning <= critical <= 100")
+    except UsageError as exc:
+        print(json.dumps(unavailable_json(str(exc)), separators=(",", ":")))
+        return 2
+
+    path = cache_path()
+    try:
+        result = fetch_rate_limits(args.timeout, args.codex_command)
+        save_cache(path, result)
+        output = output_json(result, args.format, args.tooltip_format, args.warning, args.critical)
+    except UsageError as exc:
+        cached = None if args.no_cache else load_cache(path)
+        cached_result = cached.get("result") if isinstance(cached, dict) else None
+        try:
+            if not isinstance(cached_result, dict):
+                raise UsageError(str(exc))
+            output = output_json(
+                cached_result,
+                args.format,
+                args.tooltip_format,
+                args.warning,
+                args.critical,
+                stale=True,
+                stale_reason=str(exc),
+            )
+        except UsageError:
+            output = unavailable_json(str(exc))
+
+    print(json.dumps(output, separators=(",", ":")))
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

bin/codex-waybar-configure

@@ -0,0 +1,124 @@
+#!/usr/bin/env python3
+"""Safely add the Codex usage module to an existing Waybar JSONC config."""
+
+from __future__ import annotations
+
+import argparse
+import re
+import shutil
+import sys
+from datetime import datetime
+from pathlib import Path
+
+MODULE_NAME = "custom/codex-usage"
+MODULE_BLOCK = '''  "custom/codex-usage": {
+    "exec": "codex-waybar",
+    "return-type": "json",
+    "interval": 60,
+    "tooltip": true
+  }
+'''
+CSS_BLOCK = """
+
+/* Codex usage module */
+#custom-codex-usage {
+  min-width: 12px;
+  margin: 0 7.5px;
+}
+
+#custom-codex-usage.warning {
+  color: #d8a657;
+}
+
+#custom-codex-usage.critical,
+#custom-codex-usage.unavailable {
+  color: #a55555;
+}
+
+#custom-codex-usage.stale {
+  opacity: 0.65;
+}
+"""
+
+
+class ConfigureError(RuntimeError):
+    pass
+
+
+def backup(path: Path) -> Path:
+    stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
+    destination = path.with_name(f"{path.name}.bak.{stamp}")
+    shutil.copy2(path, destination)
+    return destination
+
+
+def insert_module_reference(text: str, section: str) -> str:
+    pattern = re.compile(rf'("modules-{section}"\s*:\s*\[)(.*?)(\])', re.DOTALL)
+    matches = list(pattern.finditer(text))
+    if len(matches) != 1:
+        raise ConfigureError(f'Could not uniquely locate "modules-{section}"')
+    match = matches[0]
+    if f'"{MODULE_NAME}"' in match.group(2):
+        return text
+    return text[: match.end(1)] + f'"{MODULE_NAME}", ' + text[match.end(1) :]
+
+
+def insert_module_block(text: str) -> str:
+    if re.search(rf'"{re.escape(MODULE_NAME)}"\s*:', text):
+        return text
+    close = text.rfind("}")
+    if close < 0:
+        raise ConfigureError("Could not locate the top-level closing brace")
+    before = text[:close].rstrip()
+    comma = "" if before.endswith("{") else ","
+    return f"{before}{comma}\n{MODULE_BLOCK}{text[close:]}"
+
+
+def configure(config_path: Path, style_path: Path, section: str) -> tuple[Path | None, Path | None]:
+    if not config_path.is_file():
+        raise ConfigureError(f"Waybar config not found: {config_path}")
+    original = config_path.read_text()
+    updated = insert_module_block(insert_module_reference(original, section))
+    config_backup = None
+    if updated != original:
+        config_backup = backup(config_path)
+        config_path.write_text(updated)
+
+    style_backup = None
+    if style_path.exists():
+        style = style_path.read_text()
+        if "#custom-codex-usage" not in style:
+            style_backup = backup(style_path)
+            style_path.write_text(style.rstrip() + CSS_BLOCK)
+    else:
+        style_path.parent.mkdir(parents=True, exist_ok=True)
+        style_path.write_text(CSS_BLOCK.lstrip())
+    return config_backup, style_backup
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    default_dir = Path.home() / ".config" / "waybar"
+    parser.add_argument("--section", choices=("left", "center", "right"), required=True)
+    parser.add_argument("--config", type=Path, default=default_dir / "config.jsonc")
+    parser.add_argument("--style", type=Path, default=default_dir / "style.css")
+    return parser.parse_args()
+
+
+def main() -> int:
+    args = parse_args()
+    try:
+        config_backup, style_backup = configure(args.config, args.style, args.section)
+    except (OSError, ConfigureError) as exc:
+        print(f"codex-waybar-configure: {exc}", file=sys.stderr)
+        return 1
+    print(f"Configured {args.config}")
+    if config_backup:
+        print(f"Backup: {config_backup}")
+    if style_backup:
+        print(f"Backup: {style_backup}")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())