Add Hermes, Antigravity, Z.ai, Ollama Cloud, MiniMax, OpenRouter, and DeepSeek to the agents panel

bin/omarchy-agent-usage-agy

@@ -0,0 +1,399 @@
+#!/usr/bin/python3
+# omarchy:summary=Print the Antigravity usage record as JSON
+# omarchy:args=[--force] [--limits-only]
+# omarchy:hidden=true
+"""Collect Google Antigravity usage into one display-ready JSON record.
+
+The id is agy, the name omarchy-default-agent and omarchy-agent know the CLI
+by, so the record and the default agent agree on which agent this is.
+
+Local stats come from Antigravity's history log, conversation summaries, and
+conversation transcripts under ~/.gemini/antigravity-cli/. The plan and
+limits come from the Google Cloud Code API (v1internal:loadCodeAssist and
+v1internal:retrieveUserQuotaSummary), asked with the sign-in agy keeps in the
+Secret Service keyring. The agents panel only ever reads the JSON this prints.
+"""
+
+import argparse
+import base64
+import json
+import os
+import sqlite3
+import subprocess
+import tempfile
+import time
+import urllib.error
+import urllib.request
+from datetime import datetime, timedelta, timezone
+from pathlib import Path
+
+AGENT_ID = "agy"
+AGENT_NAME = "Antigravity"
+AUTH_HELP = "Run `agy` to sign in."
+EXPIRED_HELP = "Antigravity's saved sign-in expired. Start `agy` to refresh it."
+# A panel opened and shut repeatedly reuses an answer for about a minute.
+PROBE_REUSE_SECONDS = 60
+# A limits-only refresh needs fresh limits, not a fresh history scan.
+LIMITS_ONLY_SCAN_REUSE_SECONDS = 900
+
+
+def app_root():
+  return Path(os.environ.get("AGY_DIR") or (Path.home() / ".gemini" / "antigravity-cli"))
+
+
+def cache_root():
+  return Path(os.environ.get("XDG_CACHE_HOME") or (Path.home() / ".cache")) / "omarchy" / "agent-usage"
+
+
+def read_json(path):
+  try:
+    return json.loads(Path(path).read_text(encoding="utf-8"))
+  except Exception:
+    return None
+
+
+# A temporary file of its own per write, so overlapping runs can't rename
+# each other's away.
+def write_json(path, payload):
+  path = Path(path)
+  path.parent.mkdir(parents=True, exist_ok=True)
+  fd, tmp = tempfile.mkstemp(prefix=path.name + ".", dir=str(path.parent))
+  try:
+    with os.fdopen(fd, "w", encoding="utf-8") as handle:
+      json.dump(payload, handle)
+    os.replace(tmp, path)
+  except Exception:
+    Path(tmp).unlink(missing_ok=True)
+    raise
+
+
+def parse_time(value):
+  if value is None or value == "":
+    return None
+  if isinstance(value, (int, float)) or str(value).isdigit():
+    stamp = float(value)
+    return datetime.fromtimestamp(stamp / 1000 if stamp > 1e12 else stamp, timezone.utc)
+  try:
+    parsed = datetime.fromisoformat(str(value).replace("Z", "+00:00"))
+  except ValueError:
+    return None
+  return parsed if parsed.tzinfo else parsed.replace(tzinfo=timezone.utc)
+
+
+def local_day(value):
+  parsed = parse_time(value)
+  return parsed.astimezone().date().isoformat() if parsed else ""
+
+
+# The sign-in agy keeps in the keyring: its access token, when that lapses,
+# and which kind of account it is.
+def login():
+  try:
+    raw = subprocess.check_output(
+      ["secret-tool", "lookup", "service", "gemini", "username", "antigravity"],
+      text=True, stderr=subprocess.DEVNULL, timeout=3,
+    )
+    data = json.loads(raw) if raw.strip() else None
+  except Exception:
+    data = None
+  if not isinstance(data, dict):
+    return None
+
+  token = data.get("token") if isinstance(data.get("token"), dict) else {}
+  expires = parse_time(token.get("expiry") or data.get("expiry"))
+  return {
+    "accessToken": str(token.get("access_token") or data.get("access_token") or ""),
+    "expiresAt": expires.timestamp() if expires else 0,
+    "authMethod": str(data.get("auth_method") or "consumer"),
+    "email": email_of(data.get("id_token")),
+  }
+
+
+def email_of(id_token):
+  try:
+    claims = str(id_token).split(".")[1]
+    return str(json.loads(base64.urlsafe_b64decode(claims + "=" * (-len(claims) % 4))).get("email") or "")
+  except Exception:
+    return ""
+
+
+def normalize_tier_label(name, tier_id=""):
+  for word in ("ultra", "pro", "free"):
+    if word in tier_id.lower() or word in name.lower():
+      return word.capitalize()
+  if name.lower().startswith("google ai "):
+    return name[10:].strip()
+  return name or tier_id.replace("-tier", "").capitalize()
+
+
+# A paid tier names the plan; without one, the current tier does.
+def extract_tier(data):
+  for key in ("paidTier", "currentTier"):
+    tier = data.get(key)
+    if isinstance(tier, dict):
+      name = str(tier.get("name") or "").strip()
+      tier_id = str(tier.get("id") or "").strip()
+      if name or tier_id:
+        return normalize_tier_label(name, tier_id)
+  return ""
+
+
+def used(bucket):
+  return round(max(0.0, min(1.0, 1.0 - float(bucket.get("remainingFraction", 1.0)))), 4)
+
+
+def reset_of(bucket):
+  parsed = parse_time(bucket.get("resetTime"))
+  return parsed.isoformat() if parsed else ""
+
+
+# The Gemini group's 5-hour and weekly buckets are the plan's own limits; the
+# other groups (Claude and GPT models) are scoped allowances, shown once
+# they have been touched.
+def extract_limits(data):
+  groups = [g for g in data.get("groups") or [] if isinstance(g, dict)]
+  if not groups:
+    return []
+  gemini = next((g for g in groups if "gemini" in str(g.get("displayName") or "").lower()), groups[0])
+
+  limits = []
+  session = weekly = None
+  for bucket in gemini.get("buckets") or []:
+    if not isinstance(bucket, dict):
+      continue
+    window = str(bucket.get("window") or "").lower()
+    bucket_id = str(bucket.get("bucketId") or "").lower()
+    name = str(bucket.get("displayName") or "").lower()
+    if window == "5h" or "5h" in bucket_id or "five" in name:
+      session = bucket
+    elif window == "weekly" or "weekly" in bucket_id or "week" in name:
+      weekly = bucket
+  if session:
+    limits.append({"label": "Session (5-hour)", "title": "Session", "percent": used(session), "resetsAt": reset_of(session)})
+  if weekly:
+    limits.append({"label": "Weekly (7-day)", "title": "Weekly", "percent": used(weekly), "resetsAt": reset_of(weekly)})
+
+  for group in groups:
+    if group is gemini:
+      continue
+    group_name = str(group.get("displayName") or "").strip()
+    short = "Claude/GPT" if "claude" in group_name.lower() else (group_name or "3P")
+    for bucket in group.get("buckets") or []:
+      if isinstance(bucket, dict) and used(bucket) > 0:
+        title = f"{short} {'Weekly' if 'week' in str(bucket.get('window') or '').lower() else 'Session'}"
+        limits.append({"label": title, "title": title, "percent": used(bucket), "resetsAt": reset_of(bucket)})
+  return limits
+
+
+def post(endpoint, token):
+  base = (os.environ.get("CLOUD_CODE_URL") or "https://daily-cloudcode-pa.googleapis.com").rstrip("/")
+  request = urllib.request.Request(
+    f"{base}/v1internal:{endpoint}", data=b"{}",
+    headers={"Authorization": f"Bearer {token}", "Content-Type": "application/json", "User-Agent": "antigravity"},
+  )
+  with urllib.request.urlopen(request, timeout=8) as response:
+    return json.loads(response.read().decode("utf-8", errors="replace"))
+
+
+# A window kept from an earlier check is dropped once it has reset, since
+# what it said no longer holds.
+def open_windows(limits):
+  now = datetime.now(timezone.utc)
+  return [limit for limit in limits or [] if not parse_time(limit.get("resetsAt")) or parse_time(limit.get("resetsAt")) > now]
+
+
+def collect_limits(entry, force):
+  result = {"limits": [], "tierLabel": "", "usageStatusText": "", "authHelpText": AUTH_HELP, "live": False, "fetchedAtMs": 0}
+  if not entry:
+    result["usageStatusText"] = "Waiting for auth"
+    return result
+
+  cache = cache_root() / "agy-limits.json"
+  cached = read_json(cache) or {}
+  fetched_at = int(cached.get("fetchedAtMs") or 0)
+  kept = open_windows(cached.get("limits"))
+  result.update(fetchedAtMs=fetched_at, tierLabel=str(cached.get("tierLabel") or ""))
+
+  if not entry["accessToken"]:
+    result.update(limits=kept, usageStatusText="Waiting for auth")
+    return result
+
+  # agy refreshes its token while it runs; one left to lapse can't be used
+  # until agy starts again.
+  if entry["expiresAt"] and entry["expiresAt"] <= time.time():
+    result.update(limits=kept, usageStatusText="Sign-in expired", authHelpText=EXPIRED_HELP)
+    return result
+
+  if kept and not force and time.time() - fetched_at / 1000 < PROBE_REUSE_SECONDS:
+    result.update(limits=kept, live=True)
+    return result
+
+  tier, limits = "", None
+  try:
+    tier = extract_tier(post("loadCodeAssist", entry["accessToken"]))
+    limits = extract_limits(post("retrieveUserQuotaSummary", entry["accessToken"])) or None
+  except urllib.error.HTTPError as error:
+    if error.code in (401, 403):
+      result.update(usageStatusText="Sign-in expired", authHelpText=EXPIRED_HELP)
+  except Exception:
+    pass
+
+  if tier:
+    result["tierLabel"] = tier
+  if limits:
+    result.update(limits=limits, live=True, fetchedAtMs=round(time.time() * 1000))
+    write_json(cache, {"fetchedAtMs": result["fetchedAtMs"], "limits": limits, "tierLabel": result["tierLabel"]})
+  else:
+    result["limits"] = kept
+    if not result["usageStatusText"] and not kept:
+      result["usageStatusText"] = "Antigravity limits unavailable"
+  return result
+
+
+def model_bucket(models, model):
+  return models.setdefault(model, {"inputTokens": 0, "outputTokens": 0, "cacheReadInputTokens": 0, "cacheCreationInputTokens": 0})
+
+
+def history_entries(root):
+  try:
+    lines = (root / "history.jsonl").read_text(encoding="utf-8", errors="ignore").splitlines()
+  except OSError:
+    return
+  for line in lines:
+    try:
+      entry = json.loads(line)
+    except ValueError:
+      continue
+    if isinstance(entry, dict):
+      yield entry
+
+
+# Every prompt from history.jsonl, every conversation from it and from the
+# conversation summaries, and the tokens each model step in a transcript
+# reports. Steps that report no usage add no tokens.
+def local_stats(root):
+  today = datetime.now().astimezone().date()
+  week = [(today - timedelta(days=offset)).isoformat() for offset in range(6, -1, -1)]
+  today = today.isoformat()
+  by_day = {day: 0 for day in week}
+  prompts = today_prompts = today_tokens = 0
+  days, sessions, today_sessions = set(), set(), set()
+  today_by_model, models = {}, {}
+
+  for entry in history_entries(root):
+    day = local_day(entry.get("timestamp"))
+    if not day:
+      continue
+    prompts += 1
+    days.add(day)
+    conversation = entry.get("conversationId") or entry.get("conversation_id")
+    if conversation:
+      sessions.add(conversation)
+    if day == today:
+      today_prompts += 1
+      if conversation:
+        today_sessions.add(conversation)
+
+  summaries = root / "conversation_summaries.db"
+  if summaries.is_file():
+    try:
+      conn = sqlite3.connect(summaries.resolve().as_uri() + "?mode=ro", uri=True, timeout=2)
+      try:
+        for conversation, modified in conn.execute("SELECT conversation_id, last_modified_time FROM conversation_summaries"):
+          if conversation:
+            sessions.add(conversation)
+            if local_day(modified) == today:
+              today_sessions.add(conversation)
+      finally:
+        conn.close()
+    except sqlite3.Error:
+      pass
+
+  for transcript in (root / "brain").glob("*/.system_generated/logs/transcript.jsonl"):
+    try:
+      lines = transcript.read_text(encoding="utf-8", errors="ignore").splitlines()
+    except OSError:
+      continue
+    for line in lines:
+      try:
+        step = json.loads(line)
+      except ValueError:
+        continue
+      usage = step.get("usage") if isinstance(step, dict) else None
+      if not isinstance(usage, dict):
+        continue
+      uncached = int(usage.get("prompt_token_count") or usage.get("input_tokens") or 0)
+      output = int(usage.get("candidates_token_count") or usage.get("output_tokens") or 0)
+      model = str(step.get("model") or "unknown")
+      bucket = model_bucket(models, model)
+      bucket["inputTokens"] += uncached
+      bucket["outputTokens"] += output
+      day = local_day(step.get("created_at"))
+      if day in by_day:
+        by_day[day] += uncached + output
+      if day == today:
+        today_tokens += uncached + output
+        today_by_model[model] = today_by_model.get(model, 0) + uncached + output
+
+  return {
+    "hasLocalStats": prompts > 0 or bool(sessions) or bool(models),
+    "todayPrompts": today_prompts,
+    "todaySessions": len(today_sessions),
+    "todayTotalTokens": today_tokens,
+    "todayTokensByModel": today_by_model,
+    "recentDays": [{"date": day, "messageCount": by_day[day]} for day in week],
+    "totalPrompts": prompts,
+    "totalSessions": len(sessions),
+    "activeDays": len(days),
+    "activeDates": sorted(days),
+    "modelUsage": models,
+  }
+
+
+# A scan from before midnight would count yesterday's prompts as today's.
+def cached_local_stats(root, max_age):
+  cache = cache_root() / "agy-stats.json"
+  cached = read_json(cache) or {}
+  today = datetime.now().astimezone().date().isoformat()
+  if cached.get("root") == str(root) and cached.get("day") == today and time.time() - float(cached.get("at") or 0) < max_age:
+    return cached["stats"]
+  stats = local_stats(root)
+  write_json(cache, {"root": str(root), "day": today, "at": time.time(), "stats": stats})
+  return stats
+
+
+def main():
+  parser = argparse.ArgumentParser()
+  parser.add_argument("--force", action="store_true")
+  parser.add_argument("--limits-only", action="store_true")
+  args = parser.parse_args()
+
+  # Nobody signed in to Antigravity here: no limits and no plan, so without
+  # local history the panel skips the record.
+  entry = login()
+  limits = collect_limits(entry, args.force)
+  tier = limits["tierLabel"]
+  if not tier and entry and entry["authMethod"] == "enterprise":
+    tier = "Enterprise"
+  elif not tier and entry and entry["authMethod"] in ("cloud", "vertex"):
+    tier = "Cloud"
+
+  record = {
+    "schemaVersion": 1,
+    "id": AGENT_ID,
+    "name": AGENT_NAME,
+    "updatedAt": datetime.now(timezone.utc).isoformat(),
+    "ready": entry is not None,
+    "tierLabel": tier,
+    "usageStatusText": limits["usageStatusText"],
+    "authHelpText": limits["authHelpText"],
+    "limits": limits["limits"],
+    "limitsStale": not limits["live"],
+    "limitsFetchedAt": limits["fetchedAtMs"],
+  }
+  record.update(cached_local_stats(app_root(), LIMITS_ONLY_SCAN_REUSE_SECONDS if args.limits_only and not args.force else 0))
+  print(json.dumps(record, separators=(",", ":")))
+
+
+if __name__ == "__main__":
+  main()