codex-usage.sh

#!/usr/bin/env bash

set -euo pipefail

# Codex quota report.
# Requirements: codex, jq, python3
# Optional: GNU coreutils (timeout/gtimeout) as a watchdog

# Defaults
WORK_START="11:00"
WORK_END="18:00"
INCLUDE_WEEKENDS="false"

# CLI argument parser
while [[ $# -gt 0 ]]; do
  case "$1" in
    -s|--start)
      WORK_START="$2"
      shift 2
      ;;
    -e|--end)
      WORK_END="$2"
      shift 2
      ;;
    -w|--weekends)
      INCLUDE_WEEKENDS="true"
      shift
      ;;
    -h|--help)
      echo "Usage: $0 [options]"
      echo "Options:"
      echo "  -s, --start TIME     Workday start time (default: 11:00)"
      echo "  -e, --end TIME       Workday end time (default: 18:00)"
      echo "  -w, --weekends       Include weekends in the work model"
      exit 0
      ;;
    *)
      echo "Unknown option: $1" >&2
      exit 1
      ;;
  esac
done

export WORK_START
export WORK_END
export INCLUDE_WEEKENDS

REQUEST_TIMEOUT=20
SERVER_WAIT=3

for command_name in codex jq python3; do
  if ! command -v "$command_name" >/dev/null 2>&1; then
    echo "Error: missing required command: $command_name" >&2
    exit 1
  fi
done

# Optional watchdog: GNU coreutils only (gtimeout on macOS via brew).
codex_cmd=(codex app-server)
if command -v gtimeout >/dev/null 2>&1; then
  codex_cmd=(gtimeout "${REQUEST_TIMEOUT}s" "${codex_cmd[@]}")
elif command -v timeout >/dev/null 2>&1; then
  codex_cmd=(timeout "${REQUEST_TIMEOUT}s" "${codex_cmd[@]}")
fi

stderr_log="$(mktemp -t codex_quota.XXXXXX)"
trap 'rm -f "$stderr_log"' EXIT

response="$(
  {
    printf '%s\n' \
      '{"jsonrpc":"2.0","method":"initialize","id":1,"params":{"clientInfo":{"name":"quota_check","title":"Quota Check","version":"1.0"}}}' \
      '{"jsonrpc":"2.0","method":"initialized","params":{}}' \
      '{"jsonrpc":"2.0","method":"account/rateLimits/read","id":2,"params":{}}'
    sleep "$SERVER_WAIT"
  } | "${codex_cmd[@]}" 2>"$stderr_log" |
    jq -c --unbuffered 'select(.id == 2)' |
    head -n 1
)" || true

if [[ -z "$response" ]]; then
  echo "Error: Codex returned no quota data." >&2
  echo "Check authentication and the command: codex app-server" >&2
  if [[ -s "$stderr_log" ]]; then
    echo "--- codex stderr ---" >&2
    tail -n 20 "$stderr_log" >&2
  fi
  exit 1
fi

python3 - "$response" <<'PY'
import datetime as dt
import json
import os
import sys

# ── Config ───────────────────────────────────────────────

def parse_time(t_str, default):
    try:
        h, m = map(int, t_str.split(':'))
        return dt.time(h, m)
    except Exception:
        return default

WORK_START = parse_time(os.environ.get("WORK_START", ""), dt.time(11, 0))
WORK_END = parse_time(os.environ.get("WORK_END", ""), dt.time(18, 0))
INCLUDE_WEEKENDS = os.environ.get("INCLUDE_WEEKENDS", "false").lower() == "true"

WORKDAY_HOURS = (
    dt.datetime.combine(dt.date.min, WORK_END)
    - dt.datetime.combine(dt.date.min, WORK_START)
).total_seconds() / 3600.0

MIN_HOURS_FOR_PACE = 1.0
MIN_WEEKLY_MINS = 24 * 60       # below this the weekday model is meaningless
WEEKLY_MIN_MINS = 6 * 24 * 60   # accept 6..8 days as "weekly" without a warning
WEEKLY_MAX_MINS = 8 * 24 * 60
USED_TOLERANCE = 105.0          # allow small overage before treating it as garbage

BAR_WIDTH = 36
BAR_PREFIX = "   Usage   "  # keeps the plan marker aligned with the bar
MAX_CALENDAR_CELLS = 10

WEEKDAYS = ["Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun"]


def die(message):
    print(f"Error: {message}", file=sys.stderr)
    sys.exit(1)


# ── Parse ────────────────────────────────────────────────

try:
    payload = json.loads(sys.argv[1])
except (IndexError, ValueError) as exc:
    die(f"could not parse the Codex response ({exc}).")

if isinstance(payload.get("error"), dict):
    message = payload["error"].get("message") or "unknown error"
    print(f"Codex error: {message}", file=sys.stderr)
    sys.exit(1)

result = payload.get("result") or {}
limits = (result.get("rateLimitsByLimitId") or {}).get("codex")

if not isinstance(limits, dict):
    limits = result.get("rateLimits")

if not isinstance(limits, dict):
    die("no Codex limit found in the response.")

primary = limits.get("primary") or {}

plan_raw = (
    limits.get("planType")
    or result.get("planType")
    or (result.get("account") or {}).get("planType")
    or "unknown"
)
plan = str(plan_raw)[:32] if isinstance(plan_raw, str) else "unknown"


def parse_reset(value):
    """resetsAt may be epoch seconds, epoch millis or an ISO-8601 string."""
    if isinstance(value, bool) or value is None:
        return None

    if isinstance(value, (int, float)):
        seconds = float(value)
    else:
        text = str(value).strip()
        if not text:
            return None
        try:
            seconds = float(text)
        except ValueError:
            try:
                parsed = dt.datetime.fromisoformat(text.replace("Z", "+00:00"))
            except ValueError:
                return None
            if parsed.tzinfo is None:
                parsed = parsed.astimezone()
            return parsed.astimezone()

    if seconds > 1e11:      # milliseconds
        seconds /= 1000.0

    try:
        return dt.datetime.fromtimestamp(seconds, dt.timezone.utc).astimezone()
    except (OverflowError, OSError, ValueError):
        return None


used = primary.get("usedPercent")
window_minutes = primary.get("windowDurationMins")
reset = parse_reset(primary.get("resetsAt"))

if used is None or window_minutes is None or reset is None:
    die("Codex returned incomplete quota data.")

try:
    used = float(used)
    window_minutes = int(window_minutes)
except (TypeError, ValueError):
    die("Codex returned malformed quota data.")

# Clamp rather than abort: a small overage (100.4%) is a rounding artifact,
# not a reason to refuse the report.
if used < 0.0 or used > USED_TOLERANCE:
    die(f"usedPercent out of range: {used}")

used = min(100.0, max(0.0, used))

if window_minutes < MIN_WEEKLY_MINS:
    die(
        f"quota window is {window_minutes / 60.0:.1f}h, not weekly — "
        "the weekday/weekend model no longer applies; "
        "update this script"
    )

now = dt.datetime.now(dt.timezone.utc).astimezone()
start = reset - dt.timedelta(minutes=window_minutes)
remaining = max(0.0, 100.0 - used)

if not (WEEKLY_MIN_MINS <= window_minutes <= WEEKLY_MAX_MINS):
    print(
        f"Warning: unexpected quota window: {window_minutes / 1440.0:.2f} days "
        f"({window_minutes} min); this report assumes a weekly window.",
        file=sys.stderr,
    )

# ── Soft Weekend Constraint Helpers ─────────────────────

def is_workday(day):
    # Weekdays are always workdays
    if day.weekday() < 5:
        return True
    # If explicitly enabled, weekends are always workdays
    if INCLUDE_WEEKENDS:
        return True
    # Soft constraint: past or current weekend days in this window are counted
    # so that historical weekend usage doesn't break calculations or skew pace.
    if day <= now.date():
        return True
    return False


# ── Colors ───────────────────────────────────────────────

is_tty = sys.stdout.isatty() and not os.environ.get("NO_COLOR")

RESET = "\033[0m" if is_tty else ""
BOLD = "\033[1m" if is_tty else ""
DIM = "\033[2m" if is_tty else ""
GREEN = "\033[32m" if is_tty else ""
YELLOW = "\033[33m" if is_tty else ""
RED = "\033[31m" if is_tty else ""
CYAN = "\033[36m" if is_tty else ""
GRAY = "\033[90m" if is_tty else ""


def bar(value, color):
    value = min(100.0, max(0.0, value))
    filled = round(BAR_WIDTH * value / 100.0)

    return color + "█" * filled + GRAY + "░" * (BAR_WIDTH - filled) + RESET


def plan_marker_line(plan_value):
    """Line with the plan marker, aligned under the bar."""
    pos = round(BAR_WIDTH * min(100.0, max(0.0, plan_value)) / 100.0)
    pos = min(BAR_WIDTH - 1, pos)

    return (
        " " * (len(BAR_PREFIX) + pos)
        + f"{CYAN}▲ plan {plan_value:.0f}%{RESET}"
    )


def format_dt(value, include_year=False):
    date_format = "%b %d, %Y" if include_year else "%b %d"
    return (
        f"{WEEKDAYS[value.weekday()]}, "
        f"{value.strftime(date_format)} "
        f"at {value:%H:%M}"
    )


def format_days(days):
    hours = days * WORKDAY_HOURS

    if hours < 1.0:
        return "less than an hour"
    if hours < WORKDAY_HOURS - 0.5:
        return f"~{hours:.0f}h of work"
    return f"~{days:.1f} workdays"


def status_of(ratio):
    """Single source of truth for the verdict text AND the bar color."""
    if ratio is None:
        return (
            CYAN,
            YELLOW,
            "⏳ Not enough data — pace visible after "
            f"{MIN_HOURS_FOR_PACE:.0f}h of work",
        )
    if ratio > 1.20:
        return (RED, RED, f"🔴 Burning {(ratio - 1) * 100:.0f}% faster than plan")
    if ratio > 1.05:
        return (
            YELLOW,
            YELLOW,
            f"⚠️  Burning {(ratio - 1) * 100:.0f}% faster than plan",
        )
    if ratio < 0.95:
        return (
            GREEN,
            GREEN,
            f"✅ Spending {(1 - ratio) * 100:.0f}% slower than plan",
        )
    return (GREEN, GREEN, "✅ Right on plan")


# ── Working-time model ───────────────────────────────────


def combine(day, clock):
    """Naive combine + astimezone -> correct local offset for that date."""
    return dt.datetime.combine(day, clock).astimezone()


def work_seconds_between(begin, end):
    if end <= begin:
        return 0.0

    total = 0.0
    day = begin.date()

    while day <= end.date():
        if is_workday(day):
            overlap_start = max(begin, combine(day, WORK_START))
            overlap_end = min(end, combine(day, WORK_END))

            if overlap_end > overlap_start:
                total += (overlap_end - overlap_start).total_seconds()

        day += dt.timedelta(days=1)

    return total


def work_segments(begin, end):
    """List of (date, working hours) per day between begin and end."""
    segments = []

    if end <= begin:
        return segments

    day = begin.date()

    while day <= end.date():
        if is_workday(day):
            seg_start = max(begin, combine(day, WORK_START))
            seg_end = min(end, combine(day, WORK_END))

            if seg_end > seg_start:
                hours = (seg_end - seg_start).total_seconds() / 3600.0
                segments.append((day, hours))

        day += dt.timedelta(days=1)

    return segments


def next_work_start(moment, limit):
    cursor = moment

    for _ in range(14):
        if is_workday(cursor.date()):
            day_start = combine(cursor.date(), WORK_START)
            day_end = combine(cursor.date(), WORK_END)

            if cursor < day_start:
                candidate = day_start
            elif cursor < day_end:
                candidate = cursor
            else:
                cursor = combine(
                    cursor.date() + dt.timedelta(days=1), WORK_START
                )
                continue

            if candidate < limit:
                return candidate
            return None

        cursor = combine(cursor.date() + dt.timedelta(days=1), WORK_START)

    return None


def add_work_seconds(begin, seconds, limit):
    if seconds <= 0:
        return begin

    cursor = begin
    seconds_left = seconds

    for _ in range(30):
        slot_start = next_work_start(cursor, limit)

        if slot_start is None:
            return None

        slot_end = min(combine(slot_start.date(), WORK_END), limit)
        available = max(0.0, (slot_end - slot_start).total_seconds())

        if seconds_left <= available:
            return slot_start + dt.timedelta(seconds=seconds_left)

        seconds_left -= available
        cursor = combine(slot_start.date() + dt.timedelta(days=1), WORK_START)

    return None


total_work_hours = work_seconds_between(start, reset) / 3600.0
elapsed_work_hours = work_seconds_between(start, min(now, reset)) / 3600.0
remaining_work_hours = work_seconds_between(max(now, start), reset) / 3600.0

if total_work_hours <= 0:
    die("no planned working hours in the current window.")

planned_used = min(
    100.0, max(0.0, 100.0 * elapsed_work_hours / total_work_hours)
)
delta_points = used - planned_used

baseline_per_hour = 100.0 / total_work_hours
baseline_per_day = baseline_per_hour * WORKDAY_HOURS

actual_per_hour = None
pace_ratio = None

if elapsed_work_hours >= MIN_HOURS_FOR_PACE:
    actual_per_hour = used / elapsed_work_hours
    pace_ratio = actual_per_hour / baseline_per_hour


def window_calendar(future_segments, covered_hours):
    """Calendar spanning the real window: past days, weekends, forecast."""
    cover = {}
    left = covered_hours

    for day, hours in future_segments:
        if left >= hours - 1e-9:
            cover[day] = "full"
            left -= hours
        elif left > 0.1:
            cover[day] = "part"
            left = 0.0
        else:
            cover[day] = "none"

    today = now.date()

    def symbol(day):
        if not is_workday(day):
            return f"{GRAY}·{RESET}"
        if day in cover:
            state = cover[day]
            if state == "full":
                return f"{GREEN}■{RESET}"
            if state == "part":
                return f"{YELLOW}◧{RESET}"
            return f"{RED}□{RESET}"
        # workday with no future hours in the window: already worked
        return f"{DIM}{CYAN}■{RESET}"

    span = (reset.date() - start.date()).days + 1
    span = max(1, min(span, MAX_CALENDAR_CELLS))
    days = [start.date() + dt.timedelta(days=i) for i in range(span)]

    cells = []
    for day in days:
        label = WEEKDAYS[day.weekday()]
        label = (
            f"{BOLD}{label}{RESET}" if day == today
            else f"{DIM}{label}{RESET}"
        )
        cells.append(f"{label} {symbol(day)}")

    return "  ".join(cells)


# ── Output ───────────────────────────────────────────────

bar_color, verdict_color, verdict = status_of(pace_ratio)

schedule_label = "every day" if INCLUDE_WEEKENDS else "weekdays"

print()
print(f"{BOLD}{CYAN}CODEX{RESET}  {DIM}{plan} · {schedule_label} {WORK_START:%H:%M}–{WORK_END:%H:%M}{RESET}")
print(f"{GRAY}{'─' * (len(BAR_PREFIX) + BAR_WIDTH + 8)}{RESET}")
print()

print(f"  {verdict_color}{BOLD}{verdict}{RESET}")
print()

print(f"{BAR_PREFIX}{bar(used, bar_color)}   {BOLD}{used:.0f}%{RESET}")
print(plan_marker_line(planned_used))

if abs(delta_points) > 0.5:
    schedule_days = abs(delta_points) / baseline_per_day
    label = "Over plan by" if delta_points > 0 else "Under plan by"
    print(f"  {DIM}{label}{RESET} {BOLD}{format_days(schedule_days)}{RESET}")

print()

# ── Will it last until reset ─────────────────────────────

print(f"  {BOLD}Will it last until reset?{RESET}")

if remaining <= 0:
    print(f"  {RED}{BOLD}🔴 Quota already exhausted{RESET}")

elif remaining_work_hours <= 0:
    print(
        f"  {GREEN}{BOLD}✅ Yes{RESET} "
        f"{DIM}— no working hours left before reset{RESET}"
    )

elif actual_per_hour is None or actual_per_hour <= 0:
    print(f"  {DIM}Forecast will appear once there is enough data{RESET}")

else:
    work_hours_to_empty = remaining / actual_per_hour
    segments = work_segments(max(now, start), reset)
    enough = work_hours_to_empty >= remaining_work_hours - 1e-9

    if enough:
        projected_left = max(
            0.0, remaining - actual_per_hour * remaining_work_hours
        )
        projected_days = projected_left / baseline_per_day
        print(
            f"  {GREEN}{BOLD}✅ Yes{RESET}  "
            f"{DIM}· with {format_days(projected_days)} of slack "
            f"left at reset{RESET}"
        )
    else:
        depletion = add_work_seconds(now, work_hours_to_empty * 3600.0, reset)
        print(
            f"  {RED}{BOLD}🔴 No"
            + (f" — runs out {format_dt(depletion)}" if depletion else "")
            + f"{RESET}"
        )

    print()
    print(f"  {window_calendar(segments, work_hours_to_empty)}")

print()
print(
    f"  {DIM}Quota left:{RESET} {BOLD}{remaining:.0f}%{RESET}"
    f"{DIM}  ·  resets {format_dt(reset, include_year=True)}{RESET}"
)

print()
print(
    f"{DIM}"
    f"Plan accrues only on {schedule_label} {WORK_START:%H:%M}–{WORK_END:%H:%M}. "
    "Forecast is based on the average pace since the window started."
    f"{RESET}"
)
print()
PY