omarchycn/bin/omarchy-agent-usage-codex

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

Local stats come from native Codex CLI session files, pi/omp sessions that
ran through openai-codex, and opencode sessions that ran on an OpenAI
provider; rate limits and the plan come from the Codex app-server RPC. The agents
panel only ever reads the JSON this prints.
"""

import argparse
import json
import os
import select
import shutil
import sqlite3
import subprocess
import time
from datetime import datetime, timedelta, timezone
from pathlib import Path

AGENT_ID = "codex"
AGENT_NAME = "Codex"
AUTH_HELP = "Run `codex login` to authenticate."


def local_day(value):
  if value is None:
    return datetime.now().strftime("%Y-%m-%d")
  if isinstance(value, (int, float)):
    # pi message timestamps are milliseconds; Codex timestamps are usually seconds.
    if value > 10_000_000_000:
      value = value / 1000
    return datetime.fromtimestamp(value).strftime("%Y-%m-%d")
  text = str(value)
  try:
    if text.endswith("Z"):
      dt = datetime.fromisoformat(text[:-1] + "+00:00")
    else:
      dt = datetime.fromisoformat(text)
    if dt.tzinfo is not None:
      dt = dt.astimezone()
    return dt.strftime("%Y-%m-%d")
  except Exception:
    return datetime.now().strftime("%Y-%m-%d")


def number(value):
  try:
    return int(value or 0)
  except Exception:
    return 0


def model_name(raw):
  value = str(raw or "codex")
  return value if value else "codex"


def runtime_env():
  home = str(Path.home())
  path_parts = [
    os.environ.get("PATH", ""),
    f"{home}/.local/bin",
    f"{home}/.npm-global/bin",
    f"{home}/.local/share/mise/shims",
  ]
  env = os.environ.copy()
  env["PATH"] = os.pathsep.join(part for part in path_parts if part)
  return env


ENV = runtime_env()


def find_command(name):
  return shutil.which(name, path=ENV.get("PATH"))


now = datetime.now()
today = now.strftime("%Y-%m-%d")
recent_dates = [(now - timedelta(days=offset)).strftime("%Y-%m-%d") for offset in range(6, -1, -1)]
recent = {day: {"date": day, "messageCount": 0} for day in recent_dates}
today_tokens_by_model = {}
model_usage = {}
today_sessions = set()
active_days = set()

today_prompts = 0
today_total_tokens = 0
total_prompts = 0
total_sessions = set()
seen_pi_messages = set()


def add_usage(day, session_key, model, input_tokens, output_tokens, cache_read, cache_write):
  global today_prompts, today_total_tokens, total_prompts
  total = input_tokens + output_tokens + cache_read + cache_write
  total_prompts += 1
  total_sessions.add(session_key)
  active_days.add(day)

  bucket = model_usage.setdefault(model, {
    "inputTokens": 0,
    "outputTokens": 0,
    "cacheReadInputTokens": 0,
    "cacheCreationInputTokens": 0,
  })
  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_prompts += 1
    today_sessions.add(session_key)
    today_total_tokens += total
    today_tokens_by_model[model] = today_tokens_by_model.get(model, 0) + total


def scan_pi_sessions():
  roots = [
    Path.home() / ".pi" / "agent" / "sessions",
    Path.home() / ".omp" / "agent" / "sessions",
  ]
  rg = find_command("rg") or "rg"
  for root in roots:
    if not root.exists():
      continue
    try:
      proc = subprocess.Popen(
        [rg, "--json", "-e", r'"provider"\s*:\s*"openai-codex"', "-e", r'"api"\s*:\s*"openai-codex', str(root)],
        stdout=subprocess.PIPE,
        stderr=subprocess.DEVNULL,
        text=True,
        errors="replace",
        env=ENV,
      )
    except FileNotFoundError:
      return

    assert proc.stdout is not None
    for raw in proc.stdout:
      try:
        event = json.loads(raw)
        if event.get("type") != "match":
          continue
        line = event.get("data", {}).get("lines", {}).get("text", "")
        path = event.get("data", {}).get("path", {}).get("text", "pi-session")
        entry = json.loads(line)
      except Exception:
        continue

      if entry.get("type") != "message":
        continue
      message_key = path + ":" + str(entry.get("id") or "")
      if message_key in seen_pi_messages:
        continue
      seen_pi_messages.add(message_key)
      message = entry.get("message") or {}
      if message.get("role") != "assistant":
        continue
      provider = str(message.get("provider") or "")
      api = str(message.get("api") or "")
      if provider != "openai-codex" and not api.startswith("openai-codex"):
        continue

      usage = message.get("usage") or {}
      if not usage:
        continue
      total = number(usage.get("totalTokens"))
      input_tokens = number(usage.get("input"))
      output_tokens = number(usage.get("output"))
      cache_read = number(usage.get("cacheRead"))
      cache_write = number(usage.get("cacheWrite"))
      if total and not (input_tokens or output_tokens or cache_read or cache_write):
        input_tokens = total
      if not (input_tokens or output_tokens or cache_read or cache_write):
        continue

      day = local_day(entry.get("timestamp") or message.get("timestamp"))
      session_key = path
      add_usage(day, session_key, model_name(message.get("model")), input_tokens, output_tokens, cache_read, cache_write)

    try:
      proc.wait(timeout=1)
    except Exception:
      proc.kill()


def scan_opencode_sessions():
  # A subscription burned entirely through opencode leaves no native session
  # files, but opencode records per-message provider, model, and token usage
  # in its own database. Read-only: opencode may be writing right now.
  db = Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) / "opencode" / "opencode.db"
  if not db.is_file():
    return
  try:
    conn = sqlite3.connect(db.resolve().as_uri() + "?mode=ro", uri=True, timeout=2)
  except sqlite3.Error:
    return
  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
        # "openai-local" gateway is not this subscription.
        if not isinstance(entry, dict) or entry.get("role") != "assistant":
          continue
        if str(entry.get("providerID") or "") != "openai":
          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"))
        if not (input_tokens or output_tokens or cache_read or cache_write):
          continue
        day = local_day((entry.get("time") or {}).get("created"))
        model = model_name(str(entry.get("modelID") or "").rstrip("/").split("/")[-1])
      except Exception:
        continue
      add_usage(day, "opencode:" + str(session_id), model, input_tokens, output_tokens, cache_read, cache_write)
  except sqlite3.Error:
    pass
  finally:
    conn.close()


def scan_native_codex_sessions():
  codex_home = Path(os.environ.get("CODEX_HOME") or (Path.home() / ".codex"))
  roots = [codex_home / "sessions", codex_home / "archived_sessions"]
  files = []
  cutoff = time.time() - 30 * 24 * 60 * 60
  for root in roots:
    if not root.exists():
      continue
    for path in root.rglob("*.jsonl"):
      try:
        if path.stat().st_mtime >= cutoff:
          files.append(path)
      except OSError:
        pass

  for path in files:
    current_model = "codex"
    try:
      with path.open(errors="replace") as handle:
        for raw in handle:
          try:
            entry = json.loads(raw)
          except Exception:
            continue
          if entry.get("type") == "turn_context":
            payload = entry.get("payload") or {}
            current_model = model_name(payload.get("model") or payload.get("model_slug") or current_model)
            continue
          payload = entry.get("payload") or entry
          if entry.get("type") == "response_item" and isinstance(payload, dict):
            payload = payload.get("payload") or payload
          if not isinstance(payload, dict):
            continue
          if payload.get("type") != "token_count":
            continue
          info = payload.get("info") or {}
          # total_token_usage is cumulative for the session. Adding every
          # snapshot makes usage grow quadratically, so count the last turn.
          usage = info.get("last_token_usage") or {}
          cache_read = number(usage.get("cached_input_tokens"))
          cache_write = number(usage.get("cache_write_input_tokens"))
          # Cached tokens are included in input_tokens, and reasoning tokens
          # are included in output_tokens. Keep the cache split without
          # counting either category twice.
          input_tokens = max(0, number(usage.get("input_tokens")) - cache_read - cache_write)
          output_tokens = number(usage.get("output_tokens"))
          if not (input_tokens or output_tokens or cache_read or cache_write):
            continue
          day = local_day(entry.get("timestamp") or path.stat().st_mtime)
          add_usage(day, str(path), current_model, input_tokens, output_tokens, cache_read, cache_write)
    except Exception:
      continue


def rpc_request(proc, request_id, method, params=None, timeout=8):
  payload = {"id": request_id, "method": method, "params": params or {}}
  proc.stdin.write(json.dumps(payload) + "\n")
  proc.stdin.flush()
  deadline = time.time() + timeout
  while time.time() < deadline:
    ready, _, _ = select.select([proc.stdout], [], [], 0.25)
    if not ready:
      continue
    line = proc.stdout.readline()
    if not line:
      break
    try:
      message = json.loads(line)
    except Exception:
      continue
    if message.get("id") == request_id:
      return message
  raise TimeoutError(method)


def limit_window(window):
  if not isinstance(window, dict):
    return None
  used = window.get("usedPercent")
  if used is None:
    return None
  mins = number(window.get("windowDurationMins"))
  if mins == 10080:
    label = "Weekly (7-day)"
  elif mins and mins % 60 == 0:
    label = f"{mins // 60}h window"
  elif mins:
    label = f"{mins}m window"
  else:
    label = "Limit"
  reset = window.get("resetsAt")
  return {
    "label": label,
    "percent": float(used) / 100.0,
    "resetsAt": datetime.fromtimestamp(number(reset), timezone.utc).isoformat() if reset else "",
  }


def fetch_codex_rpc():
  result = {"limits": [], "tierLabel": "", "usageStatusText": "", "authHelpText": AUTH_HELP}
  codex = find_command("codex")
  if not codex:
    result["usageStatusText"] = "Codex unavailable"
    result["authHelpText"] = "codex not found in PATH"
    return result

  try:
    proc = subprocess.Popen(
      [codex, "-s", "read-only", "-a", "untrusted", "app-server"],
      stdin=subprocess.PIPE,
      stdout=subprocess.PIPE,
      stderr=subprocess.DEVNULL,
      text=True,
      env=ENV,
    )
  except Exception as exc:
    result["usageStatusText"] = "Codex unavailable"
    result["authHelpText"] = str(exc)
    return result

  try:
    rpc_request(proc, 1, "initialize", {"clientInfo": {"name": "omarchy-agent-usage", "version": "1"}}, timeout=8)
    proc.stdin.write(json.dumps({"method": "initialized", "params": {}}) + "\n")
    proc.stdin.flush()
    account_msg = rpc_request(proc, 2, "account/read", timeout=4)
    limits_msg = rpc_request(proc, 3, "account/rateLimits/read", timeout=4)

    account = (account_msg.get("result") or {}).get("account") or {}
    limits = (limits_msg.get("result") or {}).get("rateLimits") or {}
    plan = limits.get("planType") or account.get("planType") or account.get("type") or ""
    result["tierLabel"] = str(plan) if plan else ""

    for window in (limits.get("primary"), limits.get("secondary")):
      entry = limit_window(window)
      if entry:
        result["limits"].append(entry)
  except Exception as exc:
    result["usageStatusText"] = "Codex limits unavailable"
    result["authHelpText"] = str(exc)
  finally:
    try:
      proc.terminate()
      proc.wait(timeout=1)
    except Exception:
      try:
        proc.kill()
      except Exception:
        pass
  return result


def main():
  parser = argparse.ArgumentParser()
  # Local stats and limits come from the same cheap scan, so there is no
  # cache to force past and no faster limits-only path. The flags exist so
  # every collector accepts the same invocation.
  parser.add_argument("--force", action="store_true")
  parser.add_argument("--limits-only", action="store_true")
  parser.parse_args()

  scan_pi_sessions()
  scan_native_codex_sessions()
  scan_opencode_sessions()
  rpc = fetch_codex_rpc()

  record = {
    "schemaVersion": 1,
    "id": AGENT_ID,
    "name": AGENT_NAME,
    "updatedAt": datetime.now(timezone.utc).isoformat(),
    "ready": True,
    "hasLocalStats": True,
    "todayPrompts": today_prompts,
    "todaySessions": len(today_sessions),
    "todayTotalTokens": today_total_tokens,
    "todayTokensByModel": today_tokens_by_model,
    "recentDays": [recent[day] for day in recent_dates],
    "totalPrompts": total_prompts,
    "totalSessions": len(total_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),
    "modelUsage": model_usage,
  }
  record.update(rpc)
  print(json.dumps(record, separators=(",", ":")))


if __name__ == "__main__":
  main()