usage.py

import json
import os
import re
import shutil
import subprocess
from datetime import datetime

_SESSION_RE = re.compile(r"Current session: (\d+)% used . resets (.+?) \(")
_WEEK_RE = re.compile(r"Current week \(all models\): (\d+)% used . resets (.+?) \(")

# ponytail: one line per call, append-only - `awk -F, '{s+=$2} END {print s}'` for a running total
COST_LOG = os.path.join(os.path.dirname(os.path.abspath(__file__)), "usage_cost.log")


def _parse_reset(text):
    dated = f"{datetime.now().year} {text}"
    try:
        return datetime.strptime(dated, "%Y %b %d, %I:%M%p")
    except ValueError:
        return datetime.strptime(dated, "%Y %b %d, %I%p")  # on-the-hour times omit ":00"


# ponytail: when neither a 5h window nor a week window has started yet, `claude -p /usage`
# drops the "Current session"/"Current week" lines entirely and prints only the
# contributing-usage breakdown - this is normal output, not a parse failure. The week %
# is genuinely unknown in that state (not 0 - a week can be mid-cycle with real usage but
# no line printed), so both come back as None rather than a guessed number.
def parse_usage(text):
    """Returns {'session': (pct, reset_dt), 'week': (pct, reset_dt)}; either pair is
    (None, None) if that window's line is absent from the output."""
    result = {}
    for key, regex in (("session", _SESSION_RE), ("week", _WEEK_RE)):
        m = regex.search(text)
        result[key] = (int(m.group(1)), _parse_reset(m.group(2))) if m else (None, None)
    return result


def get_usage():
    """Returns {'session': (pct, reset_dt), 'week': (pct, reset_dt), 'cost_usd': float}."""
    # ponytail: CreateProcess (what subprocess.run uses without shell=True) doesn't
    # apply PATHEXT, so a plain "claude" argv entry silently fails to find an
    # npm-install's claude.cmd shim - shutil.which does honor PATHEXT and resolves it.
    claude = shutil.which("claude")
    if not claude:
        raise RuntimeError("claude CLI not found on PATH")
    r = subprocess.run(
        [claude, "-p", "/usage", "--output-format", "json"],
        capture_output=True, text=True, timeout=60, encoding="utf-8",
    )
    data = json.loads(r.stdout)
    cost = data.get("total_cost_usd", 0.0)
    text = data.get("result", "")

    with open(COST_LOG, "a", encoding="utf-8") as f:
        f.write(f"{datetime.now().isoformat(timespec='seconds')},{cost}\n")

    if "limits usage" not in text:
        raise RuntimeError(f"couldn't parse /usage output (cost ${cost}):\n{text}")

    usage = parse_usage(text)
    usage["cost_usd"] = cost
    return usage


def format_remaining(reset_dt, now=None):
    if reset_dt is None:
        return "--"
    now = now or datetime.now()
    delta = reset_dt - now
    if delta.total_seconds() < 0:
        return "0h 0m"
    mins = int(delta.total_seconds() // 60)
    return f"{mins // 60}h {mins % 60}m"


if __name__ == "__main__":
    # self-check: a captured no-window response (real output from this machine at
    # 08:11:47 - both lines absent even though the week had 26% used moments earlier)
    _NO_WINDOW_TEXT = """You are currently using your subscription to power your Claude Code usage

What's contributing to your limits usage?
Approximate, based on local sessions on this machine — does not include other devices or claude.ai.

Last 24h · 1111 requests · 23 sessions
  42% of your usage was at >150k context
"""
    assert parse_usage(_NO_WINDOW_TEXT) == {"session": (None, None), "week": (None, None)}
    assert format_remaining(None) == "--"

    # self-check: a normal response with both windows active
    _BOTH_TEXT = (
        "Current session: 12% used · resets Jul 29, 1:10pm (Europe/London)\n"
        "Current week (all models): 26% used · resets Aug 4, 1am (Europe/London)\n"
    )
    parsed = parse_usage(_BOTH_TEXT)
    assert parsed["session"][0] == 12 and parsed["week"][0] == 26
    assert parsed["session"][1].hour == 13 and parsed["week"][1].day == 4
    print("self-check OK")

    usage = get_usage()
    for key in ("session", "week"):
        pct, reset_dt = usage[key]
        pct_str = "--" if pct is None else f"{pct}%"
        print(f"{key}: {pct_str} used, resets {reset_dt}, remaining {format_remaining(reset_dt)}")