omarchy/bin/omarchy-agent-usage-claude

#!/usr/bin/python3
# omarchy:summary=Print the Claude Code usage record as JSON
# omarchy:args=[--force] [--limits-only]
# omarchy:hidden=true
"""Collect Claude Code usage into one display-ready JSON record.

Everything the agents panel shows for Claude comes from this one
command: local transcript stats from ~/.claude/projects, the stats-cache and
history fallbacks for machines without transcripts, pi/omp and opencode
sessions that ran on an Anthropic provider, and the authoritative rate
limits from Anthropic's OAuth usage endpoint. The panel itself only ever
reads the JSON this prints; it never talks to disk formats or endpoints.
"""

from __future__ import annotations

import argparse
import datetime as dt
import fcntl
import hashlib
import json
import os
import re
import sqlite3
import sys
import tempfile
import time
import urllib.error
import urllib.request
from pathlib import Path
from typing import Any

AGENT_ID = "claude"
AGENT_NAME = "Claude Code"
AUTH_HELP = "Run `claude auth login` to restore authoritative usage."
USAGE_ENDPOINT = "https://api.anthropic.com/api/oauth/usage"
PROBE_MIN_INTERVAL_SECONDS = 15
# Anthropic rate-limits its usage endpoint readily. A refused re-check of
# numbers this recent still describes the account; only older ones are
# reported as stale.
CURRENT_ENOUGH_SECONDS = 900


def config_dir() -> Path:
  return expand_path(os.environ.get("CLAUDE_CONFIG_DIR") or "~/.claude")


def expand_path(value: str) -> Path:
  return Path(os.path.expandvars(os.path.expanduser(value))).resolve()


def cache_root() -> Path:
  root = Path(os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache")) / "omarchy" / "agent-usage"
  root.mkdir(parents=True, exist_ok=True)
  return root


def date_string(value: dt.date) -> str:
  return value.strftime("%Y-%m-%d")


def recent_date_strings() -> list[str]:
  today = dt.datetime.now().date()
  return [date_string(today - dt.timedelta(days=offset)) for offset in range(6, -1, -1)]


def local_date_string() -> str:
  return date_string(dt.datetime.now().date())


def local_date_from_timestamp(value: Any) -> str:
  if value is None:
    return local_date_string()

  if isinstance(value, (int, float)):
    try:
      seconds = float(value) / 1000.0 if float(value) > 10_000_000_000 else float(value)
      return date_string(dt.datetime.fromtimestamp(seconds).date())
    except Exception:
      return local_date_string()

  raw = str(value).strip()
  if not raw:
    return local_date_string()

  # Claude JSONL timestamps are usually ISO-8601. Python accepts offsets but
  # not a trailing Z until we normalize it to +00:00.
  try:
    parsed = dt.datetime.fromisoformat(raw.replace("Z", "+00:00"))
    if parsed.tzinfo is not None:
      parsed = parsed.astimezone()
    return date_string(parsed.date())
  except Exception:
    return local_date_string()


def usage_token(usage: dict[str, Any], snake_key: str, camel_key: str) -> int:
  value = usage.get(snake_key, usage.get(camel_key, 0))
  try:
    return round(float(value or 0))
  except Exception:
    return 0


def number(value: Any) -> int:
  try:
    n = float(value or 0)
    return round(n) if n == n else 0
  except Exception:
    return 0


def empty_bucket() -> dict[str, int]:
  return {
    "inputTokens": 0,
    "outputTokens": 0,
    "cacheReadInputTokens": 0,
    "cacheCreationInputTokens": 0,
  }


# ---------------------------------------------------------------- local scan


def scan_projects(projects_path: Path) -> dict[str, Any]:
  today = local_date_string()
  recent_dates = recent_date_strings()
  recent = {day: {"date": day, "messageCount": 0} for day in recent_dates}

  seen: set[str] = set()
  sessions: set[str] = set()
  active_days: set[str] = set()
  today_sessions: set[str] = set()
  today_tokens: dict[str, int] = {}
  usage_by_model: dict[str, dict[str, int]] = {}
  prompts = 0
  today_prompt_count = 0
  today_token_total = 0

  files = projects_path.rglob("*.jsonl") if projects_path.is_dir() else []
  for path in files:
    try:
      with path.open("r", encoding="utf-8", errors="replace") as handle:
        for line_number, line in enumerate(handle, 1):
          # Cheap pre-filter before JSON parsing keeps files with unrelated
          # lines inexpensive.
          if '"usage":' not in line:
            continue

          try:
            entry = json.loads(line)
          except Exception:
            continue

          message = entry.get("message") if isinstance(entry.get("message"), dict) else {}
          if entry.get("type") != "assistant" and message.get("role") != "assistant":
            continue

          usage = message.get("usage") or entry.get("usage")
          if not isinstance(usage, dict):
            continue

          message_id = message.get("id") or entry.get("messageId") or ""
          unique_key = str(message_id) if message_id else f"{path}:{entry.get('uuid') or entry.get('requestId') or line_number}"
          if unique_key in seen:
            continue
          seen.add(unique_key)

          input_tokens = usage_token(usage, "input_tokens", "inputTokens")
          output_tokens = usage_token(usage, "output_tokens", "outputTokens")
          cache_read = usage_token(usage, "cache_read_input_tokens", "cacheReadInputTokens")
          cache_write = usage_token(usage, "cache_creation_input_tokens", "cacheCreationInputTokens")
          total = input_tokens + output_tokens + cache_read + cache_write
          if total <= 0:
            continue

          model = str(message.get("model") or entry.get("model") or "claude")
          day = local_date_from_timestamp(entry.get("timestamp") or message.get("timestamp"))
          session_key = str(entry.get("sessionId") or path)
          sessions.add(session_key)
          active_days.add(day)
          prompts += 1

          bucket = usage_by_model.setdefault(model, empty_bucket())
          bucket["inputTokens"] += input_tokens
          bucket["outputTokens"] += output_tokens
          bucket["cacheReadInputTokens"] += cache_read
          bucket["cacheCreationInputTokens"] += cache_write

          if day in recent:
            # recentDays.messageCount is actually a token total, despite the
            # legacy name shared with synced snapshots.
            recent[day]["messageCount"] += total

          if day == today:
            today_prompt_count += 1
            today_sessions.add(session_key)
            today_token_total += total
            today_tokens[model] = today_tokens.get(model, 0) + total
    except Exception as exc:
      print(f"Ignoring unreadable Claude project file {path}: {exc}", file=sys.stderr)

  return {
    "todayPrompts": today_prompt_count,
    "todaySessions": len(today_sessions),
    "todayTotalTokens": today_token_total,
    "todayTokensByModel": today_tokens,
    "recentDays": [recent[day] for day in recent_dates],
    "modelUsage": usage_by_model,
    "totalPrompts": prompts,
    "totalSessions": len(sessions),
    # Days with any recorded usage, for the all-time "N days" summary. The
    # dates travel too: merging snapshots from several machines needs their
    # union, which a count alone cannot give.
    "activeDays": len(active_days),
    "activeDates": sorted(active_days),
  }


def scan_cache_paths(projects_path: Path) -> tuple[Path, Path]:
  digest = hashlib.sha1(str(projects_path).encode("utf-8")).hexdigest()[:16]
  root = cache_root()
  return root / f"claude-scan-{digest}.json", root / f"claude-scan-{digest}.lock"


def read_fresh_json(path: Path, max_age_seconds: float) -> dict[str, Any] | None:
  if max_age_seconds <= 0 or not path.exists():
    return None
  try:
    if time.time() - path.stat().st_mtime <= max_age_seconds:
      return json.loads(path.read_text(encoding="utf-8"))
  except Exception:
    return None
  return None


def write_json(path: Path, payload: dict[str, Any]) -> None:
  # A temp name unique to this writer, not derived from the target: several
  # collectors can run at once (the update command backgrounds one per agent,
  # the panel refreshes on its own), and a shared temp path means the second
  # replace finds the first one's file already moved away.
  handle_fd, tmp_name = tempfile.mkstemp(dir=path.parent, prefix=path.name + ".", suffix=".tmp")
  tmp = Path(tmp_name)
  try:
    with os.fdopen(handle_fd, "w", encoding="utf-8") as handle:
      handle.write(json.dumps(payload, separators=(",", ":"), sort_keys=True) + "\n")
    # mkstemp opens at 0600; these caches were world-readable before.
    tmp.chmod(0o644)
    tmp.replace(path)
  except BaseException:
    tmp.unlink(missing_ok=True)
    raise


def cached_scan(projects_path: Path, max_age_seconds: float) -> dict[str, Any]:
  cache_file, lock_file = scan_cache_paths(projects_path)

  cached = read_fresh_json(cache_file, max_age_seconds)
  if cached is not None:
    return cached

  with lock_file.open("w") as lock:
    fcntl.flock(lock, fcntl.LOCK_EX)
    cached = read_fresh_json(cache_file, max_age_seconds)
    if cached is not None:
      return cached
    summary = scan_projects(projects_path)
    write_json(cache_file, summary)
    return summary


# ------------------------------------------------------------- local fallback
#
# A machine without transcripts on disk can still know its history: Claude
# Code keeps aggregate counters in stats-cache.json and per-prompt history in
# history.jsonl. Only consulted when the project scan comes back empty.


def stats_cache_fallback(claude_dir: Path) -> dict[str, Any] | None:
  try:
    data = json.loads((claude_dir / "stats-cache.json").read_text(encoding="utf-8"))
  except Exception:
    return None

  today = local_date_string()
  daily_model_tokens = data.get("dailyModelTokens") or []
  today_tokens = {}
  for entry in daily_model_tokens:
    if isinstance(entry, dict) and entry.get("date") == today:
      today_tokens = entry.get("tokensByModel") or {}
      break

  daily_activity = [day for day in (data.get("dailyActivity") or []) if isinstance(day, dict)]
  active_dates = sorted({str(day.get("date")) for day in daily_activity if number(day.get("messageCount")) > 0 and day.get("date")})
  today_prompts, today_sessions = today_prompts_from_history(claude_dir)

  return {
    "todayPrompts": today_prompts,
    "todaySessions": today_sessions,
    "todayTotalTokens": sum(number(v) for v in today_tokens.values()),
    "todayTokensByModel": today_tokens,
    "recentDays": daily_activity[-7:],
    "modelUsage": data.get("modelUsage") or {},
    "totalPrompts": number(data.get("totalMessages")),
    "totalSessions": number(data.get("totalSessions")),
    "activeDays": len(active_dates),
    "activeDates": active_dates,
  }


def today_prompts_from_history(claude_dir: Path) -> tuple[int, int]:
  prompts = 0
  sessions: set[str] = set()
  start_of_day = dt.datetime.combine(dt.datetime.now().date(), dt.time.min).timestamp() * 1000
  try:
    with (claude_dir / "history.jsonl").open("r", encoding="utf-8", errors="replace") as handle:
      lines = handle.readlines()
  except Exception:
    return 0, 0

  for line in reversed(lines):
    line = line.strip()
    if not line:
      continue
    try:
      entry = json.loads(line)
    except Exception:
      continue
    if number(entry.get("timestamp")) < start_of_day:
      break
    prompts += 1
    if entry.get("sessionId"):
      sessions.add(str(entry.get("sessionId")))
  return prompts, len(sessions)


# --------------------------------------------------------------- pi and omp
#
# These agents can consume a Claude subscription without writing native
# Claude Code transcripts. Their compatible JSONL session formats carry the
# provider, model, and token usage on every assistant message.


def scan_pi_usage(max_age_seconds: float) -> dict[str, Any] | None:
  roots = [
    Path.home() / ".pi" / "agent" / "sessions",
    Path.home() / ".omp" / "agent" / "sessions",
  ]
  cache_file = cache_root() / "claude-pi-sessions.json"
  cached = read_fresh_json(cache_file, max_age_seconds)
  if cached is not None:
    return cached.get("stats")

  today = local_date_string()
  recent_dates = recent_date_strings()
  recent = {day: {"date": day, "messageCount": 0} for day in recent_dates}
  sessions: set[str] = set()
  active_days: set[str] = set()
  today_sessions: set[str] = set()
  today_tokens: dict[str, int] = {}
  usage_by_model: dict[str, dict[str, int]] = {}
  seen: set[str] = set()
  prompts = 0
  today_prompt_count = 0
  today_token_total = 0

  for root in roots:
    files = root.rglob("*.jsonl") if root.is_dir() else []
    for path in files:
      try:
        with path.open("r", encoding="utf-8", errors="replace") as handle:
          for line_number, line in enumerate(handle, 1):
            if '"usage"' not in line or '"assistant"' not in line:
              continue
            try:
              entry = json.loads(line)
              message = entry.get("message") if isinstance(entry.get("message"), dict) else {}
              if entry.get("type") != "message" or message.get("role") != "assistant":
                continue
              provider = str(message.get("provider") or "")
              if provider != "anthropic":
                continue
              unique_key = f"{path}:{entry.get('id') or line_number}"
              if unique_key in seen:
                continue
              seen.add(unique_key)
              usage = message.get("usage") or {}
              input_tokens = usage_token(usage, "input", "inputTokens")
              output_tokens = usage_token(usage, "output", "outputTokens")
              cache_read = usage_token(usage, "cacheRead", "cache_read_input_tokens")
              cache_write = usage_token(usage, "cacheWrite", "cache_creation_input_tokens")
              total = input_tokens + output_tokens + cache_read + cache_write
              if total <= 0:
                total = number(usage.get("totalTokens"))
                input_tokens = total
              if total <= 0:
                continue
              model = str(message.get("model") or "claude")
              day = local_date_from_timestamp(entry.get("timestamp") or message.get("timestamp"))
            except Exception:
              continue

            session_key = str(path)
            sessions.add(session_key)
            active_days.add(day)
            prompts += 1
            bucket = usage_by_model.setdefault(model, empty_bucket())
            bucket["inputTokens"] += input_tokens
            bucket["outputTokens"] += output_tokens
            bucket["cacheReadInputTokens"] += cache_read
            bucket["cacheCreationInputTokens"] += cache_write
            if day in recent:
              recent[day]["messageCount"] += total
            if day == today:
              today_prompt_count += 1
              today_sessions.add(session_key)
              today_token_total += total
              today_tokens[model] = today_tokens.get(model, 0) + total
      except OSError:
        continue

  stats = None
  if prompts > 0:
    stats = {
      "todayPrompts": today_prompt_count,
      "todaySessions": len(today_sessions),
      "todayTotalTokens": today_token_total,
      "todayTokensByModel": today_tokens,
      "recentDays": [recent[day] for day in recent_dates],
      "modelUsage": usage_by_model,
      "totalPrompts": prompts,
      "totalSessions": len(sessions),
      "activeDays": len(active_days),
      "activeDates": sorted(active_days),
    }
  write_json(cache_file, {"stats": stats})
  return stats


# ---------------------------------------------------------------- opencode
#
# A Claude subscription burned entirely through opencode never writes a
# transcript under ~/.claude, but opencode records per-message provider,
# model, and token usage in its own database. Scan it for Anthropic-provider
# messages and merge the result into whatever the transcript scan found.


def scan_opencode_usage(max_age_seconds: float) -> dict[str, Any] | None:
  db = Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) / "opencode" / "opencode.db"
  if not db.is_file():
    return None

  # Same freshness contract as the transcript scan: --limits-only promises to
  # reuse recent local stats, and a big opencode history walked on every panel
  # open would break that promise.
  cache_file = cache_root() / f"claude-opencode-{hashlib.sha1(str(db).encode('utf-8')).hexdigest()[:16]}.json"
  cached = read_fresh_json(cache_file, max_age_seconds)
  if cached is not None:
    return cached.get("stats")

  today = local_date_string()
  recent_dates = recent_date_strings()
  recent = {day: {"date": day, "messageCount": 0} for day in recent_dates}
  sessions: set[str] = set()
  active_days: set[str] = set()
  today_sessions: set[str] = set()
  today_tokens: dict[str, int] = {}
  usage_by_model: dict[str, dict[str, int]] = {}
  prompts = 0
  today_prompt_count = 0
  today_token_total = 0

  try:
    # Read-only: opencode may be writing right now.
    conn = sqlite3.connect(db.resolve().as_uri() + "?mode=ro", uri=True, timeout=2)
  except sqlite3.Error:
    return None
  try:
    conn.execute("PRAGMA query_only = ON")
    for session_id, raw in conn.execute("SELECT session_id, data FROM message"):
      # One malformed row must not abort the scan, so every shape assumption
      # lives inside the try.
      try:
        entry = json.loads(raw)
        # Exact match: opencode provider ids are free-form, and a custom
        # "anthropic-proxy" gateway is not this subscription.
        if not isinstance(entry, dict) or entry.get("role") != "assistant":
          continue
        if str(entry.get("providerID") or "") != "anthropic":
          continue
        tokens = entry.get("tokens") or {}
        cache = tokens.get("cache") or {}
        input_tokens = number(tokens.get("input"))
        # opencode keeps thinking tokens out of output; both are generated.
        output_tokens = number(tokens.get("output")) + number(tokens.get("reasoning"))
        cache_read = number(cache.get("read"))
        cache_write = number(cache.get("write"))
        total = input_tokens + output_tokens + cache_read + cache_write
        if total <= 0:
          continue

        created = number((entry.get("time") or {}).get("created"))
        day = dt.datetime.fromtimestamp(created / 1000).strftime("%Y-%m-%d") if created > 0 else today
        model = str(entry.get("modelID") or "claude").rstrip("/").split("/")[-1]
      except Exception:
        continue
      session_key = "opencode:" + str(session_id)
      sessions.add(session_key)
      active_days.add(day)
      prompts += 1

      bucket = usage_by_model.setdefault(model, empty_bucket())
      bucket["inputTokens"] += input_tokens
      bucket["outputTokens"] += output_tokens
      bucket["cacheReadInputTokens"] += cache_read
      bucket["cacheCreationInputTokens"] += cache_write

      if day in recent:
        recent[day]["messageCount"] += total
      if day == today:
        today_prompt_count += 1
        today_sessions.add(session_key)
        today_token_total += total
        today_tokens[model] = today_tokens.get(model, 0) + total
  except sqlite3.Error:
    return None
  finally:
    conn.close()

  stats = None
  if prompts > 0:
    stats = {
      "todayPrompts": today_prompt_count,
      "todaySessions": len(today_sessions),
      "todayTotalTokens": today_token_total,
      "todayTokensByModel": today_tokens,
      "recentDays": [recent[day] for day in recent_dates],
      "modelUsage": usage_by_model,
      "totalPrompts": prompts,
      "totalSessions": len(sessions),
      "activeDays": len(active_days),
      "activeDates": sorted(active_days),
    }
  write_json(cache_file, {"stats": stats})
  return stats


def merge_stats(base: dict[str, Any], extra: dict[str, Any]) -> dict[str, Any]:
  merged = dict(base)
  for key in ("todayPrompts", "todaySessions", "todayTotalTokens", "totalPrompts", "totalSessions"):
    merged[key] = number(base.get(key)) + number(extra.get(key))

  combined = dict(base.get("todayTokensByModel") or {})
  for model, count in (extra.get("todayTokensByModel") or {}).items():
    combined[model] = number(combined.get(model)) + number(count)
  merged["todayTokensByModel"] = combined

  usage = {model: dict(bucket) for model, bucket in (base.get("modelUsage") or {}).items()}
  for model, bucket in (extra.get("modelUsage") or {}).items():
    target = usage.setdefault(model, empty_bucket())
    for field, count in (bucket or {}).items():
      target[field] = number(target.get(field)) + number(count)
  merged["modelUsage"] = usage

  by_date: dict[str, int] = {}
  for source in (base.get("recentDays") or [], extra.get("recentDays") or []):
    for day in source:
      date = str((day or {}).get("date") or "")
      if date:
        by_date[date] = by_date.get(date, 0) + number((day or {}).get("messageCount"))
  merged["recentDays"] = [{"date": date, "messageCount": by_date[date]} for date in sorted(by_date)]

  # Sources overlap in time, so union dates rather than summing counts. A
  # fallback that only knows a count still bounds the answer from below.
  dates = set(base.get("activeDates") or []) | set(extra.get("activeDates") or [])
  merged["activeDates"] = sorted(dates)
  merged["activeDays"] = max(len(dates), number(base.get("activeDays")), number(extra.get("activeDays")))
  return merged


# ------------------------------------------------------------------- limits


# The access token, its expiry, and the display-safe plan label from the
# CLI's login. Nothing else leaves the credential store: the token goes
# nowhere but the Authorization header of the limits probe, and only the
# plan label may travel into the printed record.
def oauth_login(claude_dir: Path) -> tuple[str, int, str]:
  try:
    data = json.loads((claude_dir / ".credentials.json").read_text(encoding="utf-8"))
  except Exception:
    return "", 0, ""
  login = data.get("claudeAiOauth")
  if not isinstance(login, dict):
    return "", 0, ""
  plan = plan_label(str(login.get("rateLimitTier") or ""), str(login.get("subscriptionType") or ""))
  return str(login.get("accessToken") or ""), number(login.get("expiresAt")), plan


def plan_label(tier: str, subscription: str) -> str:
  if tier:
    match = re.search(r"max_(\d+x)", tier, re.IGNORECASE)
    if match:
      return "Max " + match.group(1)
  if subscription:
    return subscription[0].upper() + subscription[1:]
  return ""


def parse_utilization(value: Any) -> float:
  try:
    return float(str(value).strip().replace("%", ""))
  except Exception:
    return float("nan")


def normalize_utilization(value: Any, percent_scale: bool) -> float:
  n = parse_utilization(value)
  if not (n >= 0):
    return -1.0
  # Anthropic's OAuth usage endpoint currently reports percentages (for
  # example 37.0 or 1.0). Older payloads sometimes used fractions (0.37).
  # A payload containing any value >= 1 is percent-scaled, so 1.0 renders
  # as 1%, not 100%.
  if percent_scale or n > 1:
    return min(1.0, n / 100.0)
  return min(1.0, n)


def normalize_reset_at(value: Any) -> str:
  if value is None:
    return ""
  raw = str(value).strip()
  if raw == "":
    return ""
  if raw.isdigit():
    ts = int(raw)
    if ts < 1e12:
      ts *= 1000
    try:
      return dt.datetime.fromtimestamp(ts / 1000, dt.timezone.utc).isoformat()
    except Exception:
      return raw
  try:
    parsed = dt.datetime.fromisoformat(raw.replace("Z", "+00:00"))
    return parsed.isoformat()
  except Exception:
    return raw


def usage_bucket(payload: dict[str, Any], key: str) -> dict[str, Any] | None:
  bucket = payload.get(key)
  return bucket if isinstance(bucket, dict) else None


# An entry's `kind` names its window the way the flat buckets' keys do
# ("weekly_scoped", "five_hour_scoped"). The panel reads a window out of free
# text, which cannot survive a model name like "Opus 5 (1M context)" — the
# "1M" reads as a one-minute window — so the window is settled here instead
# and travels as an explicit title. It is capitalized the way the flat windows
# title themselves, so "Fable Weekly" sits beside "Weekly" rather than under it.
def scoped_window(kind: str) -> str:
  text = kind.lower()
  if "month" in text:
    return "Monthly"
  if "week" in text or "day" in text:
    return "Weekly"
  if "hour" in text or "session" in text:
    return "Session"
  return ""


# Alongside the flat buckets, the payload carries a `limits` array, and that
# array is the only place a model-scoped allowance shows up — a weekly window
# that only Fable draws from, say. The matching legacy keys
# (`seven_day_opus`, `seven_day_sonnet`, …) stayed behind at null, so a
# collector that reads buckets alone silently drops a limit the account is
# actually spending against. A model can hold more than one scoped window, and
# only the pair of model and window tells them apart, so both make the title
# and both make the key that keeps a repeat out.
def scoped_limits(payload: dict[str, Any], percent_scale: bool) -> list[dict[str, Any]]:
  entries = payload.get("limits")
  if not isinstance(entries, list):
    return []
  out: list[dict[str, Any]] = []
  seen: set[tuple[str, str]] = set()
  for entry in entries:
    if not isinstance(entry, dict):
      continue
    scope = entry.get("scope")
    model = scope.get("model") if isinstance(scope, dict) else None
    if not isinstance(model, dict):
      continue
    # A display name is what the panel wants, but an entry carrying only an id
    # still names a window worth showing.
    name = str(model.get("display_name") or model.get("id") or "").strip()
    kind = str(entry.get("kind") or "").strip()
    if name == "" or (name, kind) in seen:
      continue
    percent = normalize_utilization(entry.get("percent"), percent_scale)
    if percent < 0:
      continue
    seen.add((name, kind))
    window = scoped_window(kind)
    title = name + " " + window if window else name
    out.append({
      "label": title,
      "title": title,
      "percent": percent,
      "resetsAt": normalize_reset_at(entry.get("resets_at")),
    })
  return out


def probe_limits(access_token: str) -> dict[str, Any]:
  request = urllib.request.Request(
    USAGE_ENDPOINT,
    headers={
      "Authorization": "Bearer " + access_token,
      "anthropic-beta": "oauth-2025-04-20",
      "Accept": "application/json",
    },
  )
  try:
    with urllib.request.urlopen(request, timeout=10) as response:
      payload = json.loads(response.read().decode("utf-8", errors="replace"))
  except urllib.error.HTTPError as error:
    retry_after = error.headers.get("retry-after", "") if error.headers else ""
    if error.code == 429:
      help_text = "Anthropic's usage endpoint is rate limiting checks right now" + (
        f" (retry after {retry_after}s)" if retry_after else ""
      ) + ". Local Claude Code stats are still shown."
    else:
      help_text = f"Anthropic's usage endpoint returned status {error.code}. Local Claude Code stats are still shown."
    return {"ok": False, "helpText": help_text}
  except Exception:
    # A transport failure reached no server at all — no route, no DNS. Any
    # real answer, including an error status, is a server we should stop
    # pestering; this is not.
    return {
      "ok": False,
      "transport": True,
      "helpText": "Couldn't reach Anthropic's usage endpoint. Retrying shortly. Local Claude Code stats are still shown.",
    }

  weekly = usage_bucket(payload, "seven_day_oauth_apps") or usage_bucket(payload, "seven_day")
  session = usage_bucket(payload, "five_hour")
  raw = [session.get("utilization") if session else None, weekly.get("utilization") if weekly else None]
  # One payload speaks one convention, so the scoped entries settle the scale
  # alongside the buckets rather than assuming their own.
  entries = payload.get("limits")
  if isinstance(entries, list):
    raw += [entry.get("percent") for entry in entries if isinstance(entry, dict)]
  percent_scale = any(parse_utilization(v) >= 1 for v in raw)

  limits = []
  if session is not None:
    percent = normalize_utilization(session.get("utilization"), percent_scale)
    if percent >= 0:
      limits.append({"label": "Session (5-hour)", "percent": percent, "resetsAt": normalize_reset_at(session.get("resets_at"))})
  if weekly is not None:
    percent = normalize_utilization(weekly.get("utilization"), percent_scale)
    if percent >= 0:
      limits.append({"label": "Weekly (7-day)", "percent": percent, "resetsAt": normalize_reset_at(weekly.get("resets_at"))})
  limits.extend(scoped_limits(payload, percent_scale))

  if not limits:
    return {"ok": False, "helpText": "Anthropic's usage endpoint returned no limits. Local Claude Code stats are still shown."}
  return {"ok": True, "limits": limits}


# A cached percentage outlives the probe that measured it, but only until its
# window rolls over: once a window has reset, the figure describes a period
# that is over, and a stale 78% would misreport an allowance that is now
# untouched. A window with no reset time, or one that will not parse, is kept
# — an unreadable timestamp is no reason to throw away a real number.
def limit_window_open(entry: dict[str, Any], now: dt.datetime) -> bool:
  raw = str(entry.get("resetsAt") or "")
  if raw == "":
    return True
  try:
    resets_at = dt.datetime.fromisoformat(raw.replace("Z", "+00:00"))
  except Exception:
    return True
  if resets_at.tzinfo is None:
    resets_at = resets_at.replace(tzinfo=dt.timezone.utc)
  return resets_at > now


def usable_cached_limits(cached: dict[str, Any]) -> list[dict[str, Any]]:
  entries = cached.get("limits")
  if not isinstance(entries, list):
    return []
  now = dt.datetime.now(dt.timezone.utc)
  return [entry for entry in entries if isinstance(entry, dict) and limit_window_open(entry, now)]


def collect_limits(access_token: str, expires_at_ms: int, force: bool, cache_name: str = "claude-limits.json") -> dict[str, Any]:
  # `live` says the numbers came from Anthropic just now (or within the reuse
  # window) rather than from a cache kept past a failed or impossible probe.
  result = {"limits": [], "usageStatusText": "", "authHelpText": AUTH_HELP, "live": False}

  # When the numbers were measured, so a stale account can say how old
  # they are.
  result["fetchedAtMs"] = 0

  # A panel that is opened and shut repeatedly must not turn into a request
  # per flick, so recent probe results are reused for a short window — and
  # kept as the answer of record when a later probe fails. Every account
  # keeps its own, so one account's allowance never stands in for another's.
  probe_cache = cache_root() / cache_name
  cached = read_fresh_json(probe_cache, float("inf")) or {}
  fallback = usable_cached_limits(cached)
  result["fetchedAtMs"] = number(cached.get("fetchedAtMs"))

  # Probing needs a live token and only the Claude Code CLI can mint one: it
  # refreshes the credential file when it runs, so a machine left alone long
  # enough finds the saved token lapsed. Say so — an empty limits list with
  # nothing else set hides the whole section and explains nothing — and keep
  # showing the last numbers whose window has not since reset.
  if access_token == "":
    result["limits"] = fallback
    result["usageStatusText"] = "Waiting for auth"
    return result
  if expires_at_ms > 0 and expires_at_ms <= time.time() * 1000:
    result["limits"] = fallback
    result["usageStatusText"] = "Sign-in expired"
    result["authHelpText"] = (
      "Claude Code's saved sign-in expired"
      + (" — showing the last known limits." if fallback else ".")
      + " Start Claude Code, or run `claude auth login`, to refresh it."
    )
    return result

  # --force is a person asking for fresh numbers, so it skips the reuse window
  # entirely; the interval is there to absorb repeated panel opens, not to
  # overrule someone who pressed refresh.
  fetched_at = number(cached.get("fetchedAtMs")) / 1000
  if fallback and not force and time.time() - fetched_at < PROBE_MIN_INTERVAL_SECONDS:
    result["limits"] = fallback
    result["live"] = True
    return result

  probe = probe_limits(access_token)
  if probe["ok"]:
    result["limits"] = probe["limits"]
    result["live"] = True
    result["fetchedAtMs"] = round(time.time() * 1000)
    write_json(probe_cache, {"fetchedAtMs": result["fetchedAtMs"], "limits": probe["limits"]})
    return result

  # The first probe after login often fires before DHCP has handed out a
  # route. Ask the shell to try again sooner than its regular interval.
  if probe.get("transport"):
    result["retryAdvised"] = True
  if fallback:
    result["limits"] = fallback
    result["live"] = time.time() - result["fetchedAtMs"] / 1000 < CURRENT_ENOUGH_SECONDS
  else:
    result["usageStatusText"] = "Claude limits unavailable"
    result["authHelpText"] = probe["helpText"]
  return result


# ---------------------------------------------------------------- accounts


# The accounts `omarchy agent account` registered, read straight from its
# registry. Only a registry holding a second account matters: with one, the
# record is exactly what it always was.
def registered_accounts() -> list[dict[str, Any]]:
  state = Path(os.environ.get("XDG_STATE_HOME") or Path.home() / ".local" / "state")
  try:
    registry = json.loads((state / "omarchy" / "agents" / "accounts" / "claude.json").read_text(encoding="utf-8"))
  except Exception:
    return []
  accounts = [a for a in registry.get("accounts") or [] if isinstance(a, dict) and a.get("id")]
  if len(accounts) < 2:
    return []
  active_id = registry.get("active") or accounts[0]["id"]
  if not any(a["id"] == active_id for a in accounts):
    active_id = accounts[0]["id"]
  for account in accounts:
    account["active"] = account["id"] == active_id
    account["switch"] = {
      "mode": "auto" if registry.get("switch") == "auto" else "manual",
      "threshold": registry.get("threshold") or 95,
    }
  return accounts


# A parked account whose sign-in lapsed can't be probed, and once every window
# it last saw has reset, the cache has nothing current to offer. That is not
# "unknown" — a reset window is an untouched allowance, and switching needs to
# know this account is the one with the most room. Each window it last saw
# comes back at 0%.
def replenished_limits(cache_name: str) -> list[dict[str, Any]]:
  cached = read_fresh_json(cache_root() / cache_name, float("inf")) or {}
  entries = cached.get("limits")
  if not isinstance(entries, list):
    return []
  now = dt.datetime.now(dt.timezone.utc)
  return [
    dict(entry, percent=0.0, resetsAt="")
    for entry in entries
    if isinstance(entry, dict) and not limit_window_open(entry, now)
  ]


def current_account_id(home: Path) -> str:
  path = Path.home() / ".claude.json" if home == expand_path("~/.claude") else home / ".claude.json"
  try:
    account = json.loads(path.read_text(encoding="utf-8")).get("oauthAccount") or {}
  except Exception:
    return ""
  return str(account.get("accountUuid") or "")


def account_limits(account: dict[str, Any], force: bool) -> dict[str, Any]:
  home = expand_path(account["home"]) if account.get("home") else expand_path("~/.claude")
  access_token, expires_at_ms, plan = oauth_login(home)
  # Keyed by the subscription, not the label: forget "Work" and add a different
  # "Work", and the new one must never fall back on the old one's allowance.
  # Who the home is signed in as now, not when it was added: a `claude auth
  # login` run in it by hand makes it a different subscription.
  identity = current_account_id(home)
  key = re.sub(r"[^A-Za-z0-9_-]", "", identity or str(account.get("accountId") or "")) or account["id"]
  # The primary home is keyed the same way once it says who it's signed in
  # as, so signing it in to another subscription never inherits the last
  # one's numbers; until then it keeps the single-account cache.
  if account.get("primary") and not identity:
    cache_name = "claude-limits.json"
  else:
    cache_name = f"claude-limits-{key}.json"
  limits = collect_limits(access_token, expires_at_ms, force, cache_name)
  if not limits["live"] and not limits["limits"]:
    limits["limits"] = replenished_limits(cache_name)
  return {
    "id": account["id"],
    "label": str(account.get("label") or account["id"]),
    "email": str(account.get("email") or ""),
    "plan": plan or str(account.get("plan") or ""),
    "active": account["active"],
    "primary": bool(account.get("primary")),
    "limits": limits["limits"],
    "stale": not limits["live"],
    "fetchedAt": limits["fetchedAtMs"],
    "usageStatusText": limits["usageStatusText"],
    "authHelpText": limits["authHelpText"],
    "retryAdvised": bool(limits.get("retryAdvised")),
  }


# -------------------------------------------------------------------- record


def main() -> int:
  parser = argparse.ArgumentParser()
  parser.add_argument("--force", action="store_true", help="rescan transcripts and re-probe limits, ignoring caches")
  parser.add_argument("--limits-only", action="store_true", help="reuse any recent transcript scan; only the limits probe must be fresh")
  parser.add_argument("--cache-seconds", type=float, default=20)
  args = parser.parse_args()

  claude_dir = config_dir()
  scan_age = 0 if args.force else (900 if args.limits_only else args.cache_seconds)
  stats = cached_scan(claude_dir / "projects", scan_age)

  if number(stats.get("totalPrompts")) <= 0:
    fallback = stats_cache_fallback(claude_dir)
    if fallback is not None:
      stats = fallback
    else:
      # No transcripts and no aggregate cache, but history.jsonl alone can
      # still put numbers on today.
      today_prompts, today_sessions = today_prompts_from_history(claude_dir)
      if today_prompts or today_sessions:
        stats = dict(stats, todayPrompts=today_prompts, todaySessions=today_sessions)

  pi_usage = scan_pi_usage(scan_age)
  if pi_usage is not None:
    stats = merge_stats(stats, pi_usage)

  opencode = scan_opencode_usage(scan_age)
  if opencode is not None:
    stats = merge_stats(stats, opencode)

  registered = registered_accounts()
  accounts = [account_limits(account, args.force) for account in registered]
  current = next((a for a in accounts if a["active"]), None)
  if current:
    # The top-level fields keep describing the account new sessions use, so
    # the bar icon and anything else reading one set of limits stays right.
    plan = current["plan"]
    limits = {key: current[key] for key in ("limits", "usageStatusText", "authHelpText", "retryAdvised")}
    stale, fetched_at = current["stale"], current["fetchedAt"]
  else:
    access_token, expires_at_ms, plan = oauth_login(claude_dir)
    limits = collect_limits(access_token, expires_at_ms, args.force)
    stale, fetched_at = not limits["live"], limits["fetchedAtMs"]

  record = {
    "schemaVersion": 1,
    "id": AGENT_ID,
    "name": AGENT_NAME,
    "updatedAt": dt.datetime.now(dt.timezone.utc).isoformat(),
    "ready": number(stats.get("totalPrompts")) > 0 or len(limits["limits"]) > 0,
    "hasLocalStats": True,
    "tierLabel": plan,
    "usageStatusText": limits["usageStatusText"],
    "authHelpText": limits["authHelpText"],
    "limits": limits["limits"],
    # Numbers kept past a failed probe, and when they were measured, so the
    # panel can tell them apart from ones checked just now.
    "limitsStale": stale,
    "limitsFetchedAt": fetched_at,
  }
  if limits.get("retryAdvised"):
    record["retryAdvised"] = True
  if accounts:
    record["accountSwitch"] = registered[0]["switch"]
    record["accounts"] = accounts
  record.update(stats)
  print(json.dumps(record, separators=(",", ":"), sort_keys=True))
  return 0


if __name__ == "__main__":
  raise SystemExit(main())