omarchycn/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, 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 sys
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


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:
  tmp = path.with_suffix(path.suffix + ".tmp")
  tmp.write_text(json.dumps(payload, separators=(",", ":"), sort_keys=True) + "\n", encoding="utf-8")
  tmp.replace(path)


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)


# ------------------------------------------------------------------- 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


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]
  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"))})

  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}


def collect_limits(access_token: str, expires_at_ms: int, force: bool) -> dict[str, Any]:
  result = {"limits": [], "usageStatusText": "", "authHelpText": AUTH_HELP}

  if access_token == "":
    result["usageStatusText"] = "Waiting for auth"
    return result
  if expires_at_ms > 0 and expires_at_ms <= time.time() * 1000:
    return result

  # 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.
  probe_cache = cache_root() / "claude-limits.json"
  cached = read_fresh_json(probe_cache, float("inf")) or {}
  fetched_at = number(cached.get("fetchedAtMs")) / 1000
  min_interval = 0 if force else PROBE_MIN_INTERVAL_SECONDS
  if cached.get("limits") and time.time() - fetched_at < max(min_interval, PROBE_MIN_INTERVAL_SECONDS):
    result["limits"] = cached["limits"]
    return result

  probe = probe_limits(access_token)
  if probe["ok"]:
    result["limits"] = probe["limits"]
    write_json(probe_cache, {"fetchedAtMs": round(time.time() * 1000), "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 cached.get("limits"):
    result["limits"] = cached["limits"]
  else:
    result["usageStatusText"] = "Claude limits unavailable"
    result["authHelpText"] = probe["helpText"]
  return result


# -------------------------------------------------------------------- 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)

  access_token, expires_at_ms, plan = oauth_login(claude_dir)
  limits = collect_limits(access_token, expires_at_ms, args.force)

  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"],
  }
  if limits.get("retryAdvised"):
    record["retryAdvised"] = True
  record.update(stats)
  print(json.dumps(record, separators=(",", ":"), sort_keys=True))
  return 0


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