312 lines
10 KiB
Python
Executable File
312 lines
10 KiB
Python
Executable File
#!/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())
|