From 56814ef8212bd980044cc73d4363cf026970c6d5 Mon Sep 17 00:00:00 2001 From: Mena <99997681+mbot11@users.noreply.github.com> Date: Thu, 1 Oct 2026 21:52:06 -0400 Subject: [PATCH 01/19] Add Hermes usage collector for the agents panel Co-authored-by: Scott Cashman --- bin/omarchy-agent-usage-hermes | 173 ++++++++++++++++++++++++ test/shell.d/agent-usage-hermes-test.sh | 82 +++++++++++ 2 files changed, 255 insertions(+) create mode 100755 bin/omarchy-agent-usage-hermes create mode 100755 test/shell.d/agent-usage-hermes-test.sh diff --git a/bin/omarchy-agent-usage-hermes b/bin/omarchy-agent-usage-hermes new file mode 100755 index 00000000000..0b692086352 --- /dev/null +++ b/bin/omarchy-agent-usage-hermes @@ -0,0 +1,173 @@ +#!/usr/bin/python3 +# omarchy:summary=Print the Hermes Agent usage record as JSON +# omarchy:args=[--write] +# omarchy:hidden=true +"""Collect Hermes Agent usage into one display-ready JSON record. + +Everything the agents panel shows for Hermes comes from this one command: +local session stats from ~/.hermes/state.db (sessions, session_model_usage, +messages) and nothing else. It prints the same record contract as the +packaged claude/codex collectors; with --write it updates the record file +atomically instead, so a user systemd timer can keep the panel fresh -- +omarchy-agent-usage-update only discovers packaged collectors. + +Chat sessions that arrived over messaging gateways count too: Hermes treats +them as first-class sessions and so does this record. +""" + +from __future__ import annotations + +import json +import os +import shutil +import sqlite3 +import sys +import tempfile +from datetime import datetime, timedelta +from pathlib import Path + +DB_PATH = Path.home() / ".hermes" / "state.db" +USAGE_DIR = Path(os.environ.get("XDG_STATE_HOME") or Path.home() / ".local/state") / "omarchy/agents/usage" +RECENT_DAYS = 7 + + +def day_bounds(day: datetime.date) -> tuple[float, float]: + start = datetime(day.year, day.month, day.day).astimezone() + return start.timestamp(), start.timestamp() + 86400 + + +def collect() -> dict: + now = datetime.now().astimezone() + record = { + "schemaVersion": 1, + "id": "hermes", + "name": "Hermes", + "updatedAt": now.isoformat(), + "ready": shutil.which("hermes") is not None, + "hasLocalStats": False, + "todayPrompts": 0, + "todaySessions": 0, + "todayTotalTokens": 0, + "todayTokensByModel": {}, + "recentDays": [], + "totalPrompts": 0, + "totalSessions": 0, + "activeDays": 0, + "activeDates": [], + "modelUsage": {}, + "limits": [], + "tierLabel": "", + "usageStatusText": "" if DB_PATH.exists() else "Hermes state not found", + "authHelpText": "", + } + + if not DB_PATH.exists(): + return record + + try: + conn = sqlite3.connect(f"file:{DB_PATH}?mode=ro", uri=True) + except sqlite3.Error: + record["usageStatusText"] = "Hermes state unreadable" + return record + + try: + today_start, _ = day_bounds(now.date()) + week_start, _ = day_bounds(now.date() - timedelta(days=RECENT_DAYS - 1)) + + record["hasLocalStats"] = True + record["totalSessions"] = conn.execute("SELECT COUNT(*) FROM sessions").fetchone()[0] + record["todaySessions"] = conn.execute( + "SELECT COUNT(*) FROM sessions WHERE started_at >= ?", (today_start,) + ).fetchone()[0] + record["totalPrompts"] = conn.execute( + "SELECT COUNT(*) FROM messages WHERE role = 'user'" + ).fetchone()[0] + record["todayPrompts"] = conn.execute( + "SELECT COUNT(*) FROM messages WHERE role = 'user' AND timestamp >= ?", (today_start,) + ).fetchone()[0] + + for model, tokens in conn.execute( + """SELECT smu.model, + SUM(smu.input_tokens + smu.output_tokens + + smu.cache_read_tokens + smu.cache_write_tokens) + FROM session_model_usage smu + JOIN sessions s ON s.id = smu.session_id + WHERE s.started_at >= ? + GROUP BY smu.model""", + (today_start,), + ): + if tokens: + record["todayTokensByModel"][model] = int(tokens) + record["todayTotalTokens"] = sum(record["todayTokensByModel"].values()) + + per_day = dict( + conn.execute( + """SELECT date(timestamp, 'unixepoch', 'localtime'), COUNT(*) + FROM messages + WHERE role = 'user' AND timestamp >= ? + GROUP BY 1""", + (week_start,), + ) + ) + record["recentDays"] = [ + { + "date": (now.date() - timedelta(days=offset)).isoformat(), + "messageCount": int(per_day.get((now.date() - timedelta(days=offset)).isoformat(), 0)), + } + for offset in range(RECENT_DAYS - 1, -1, -1) + ] + + active_dates = [ + row[0] + for row in conn.execute( + "SELECT DISTINCT date(started_at, 'unixepoch', 'localtime') FROM sessions ORDER BY 1" + ) + if row[0] + ] + record["activeDays"] = len(active_dates) + record["activeDates"] = active_dates[-400:] + + for model, inp, outp, cache_read, cache_write in conn.execute( + """SELECT model, SUM(input_tokens), SUM(output_tokens), + SUM(cache_read_tokens), SUM(cache_write_tokens) + FROM session_model_usage GROUP BY model""" + ): + record["modelUsage"][model] = { + "inputTokens": int(inp or 0), + "outputTokens": int(outp or 0), + "cacheReadInputTokens": int(cache_read or 0), + "cacheCreationInputTokens": int(cache_write or 0), + } + except sqlite3.Error: + record["usageStatusText"] = "Hermes state unreadable" + finally: + conn.close() + + return record + + +def main() -> int: + record = collect() + payload = json.dumps(record, separators=(",", ":")) + + if "--write" not in sys.argv[1:]: + print(payload) + return 0 + + USAGE_DIR.mkdir(parents=True, exist_ok=True) + fd, tmp = tempfile.mkstemp(prefix=".hermes.", dir=USAGE_DIR) + try: + with os.fdopen(fd, "w") as handle: + handle.write(payload + "\n") + os.replace(tmp, USAGE_DIR / "hermes.json") + except OSError: + try: + os.unlink(tmp) + except OSError: + pass + return 1 + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/test/shell.d/agent-usage-hermes-test.sh b/test/shell.d/agent-usage-hermes-test.sh new file mode 100755 index 00000000000..230779d1672 --- /dev/null +++ b/test/shell.d/agent-usage-hermes-test.sh @@ -0,0 +1,82 @@ +#!/bin/bash + +source "$(dirname "$0")/base-test.sh" + +require_command jq +require_command python3 + +TEST_HOME=$(mktemp -d) +trap 'rm -rf "$TEST_HOME"' EXIT + +mkdir -p "$TEST_HOME/.hermes" +HOME="$TEST_HOME" python3 - <<'PY' +import os +import sqlite3 +import time +from pathlib import Path + +path = Path(os.environ["HOME"]) / ".hermes" / "state.db" +conn = sqlite3.connect(path) +conn.executescript( + """ + CREATE TABLE sessions ( + id TEXT PRIMARY KEY, + started_at REAL NOT NULL + ); + CREATE TABLE messages ( + id INTEGER PRIMARY KEY, + session_id TEXT NOT NULL, + role TEXT NOT NULL, + timestamp REAL NOT NULL + ); + CREATE TABLE session_model_usage ( + session_id TEXT NOT NULL, + model TEXT NOT NULL, + input_tokens INTEGER NOT NULL DEFAULT 0, + output_tokens INTEGER NOT NULL DEFAULT 0, + cache_read_tokens INTEGER NOT NULL DEFAULT 0, + cache_write_tokens INTEGER NOT NULL DEFAULT 0, + PRIMARY KEY (session_id, model) + ); + """ +) +now = time.time() +conn.execute("INSERT INTO sessions VALUES (?, ?)", ("session-1", now)) +conn.executemany( + "INSERT INTO messages VALUES (?, ?, ?, ?)", + [(1, "session-1", "user", now), (2, "session-1", "assistant", now)], +) +conn.executemany( + "INSERT INTO session_model_usage VALUES (?, ?, ?, ?, ?, ?)", + [ + ("session-1", "gpt-6-astra", 10, 20, 30, 40), + ("session-1", "gpt-6-sol", 11, 21, 31, 41), + ], +) +conn.commit() +conn.close() +PY + +record=$(HOME="$TEST_HOME" "$ROOT/bin/omarchy-agent-usage-hermes") || + fail "Hermes collector exits successfully" +pass "Hermes collector exits successfully" + +printf '%s\n' "$record" | jq -e '.modelUsage["gpt-6-astra"].inputTokens == 10 and .modelUsage["gpt-6-sol"].inputTokens == 11' >/dev/null || + fail "Hermes collector preserves usage for each model" +pass "Hermes collector preserves usage for each model" + +printf '%s\n' "$record" | jq -e '.todayTokensByModel["gpt-6-sol"] == 104 and .todayTotalTokens == 204' >/dev/null || + fail "Hermes collector reports today's per-model totals" +pass "Hermes collector reports today's per-model totals" + +printf '%s\n' "$record" | jq -e '.totalPrompts == 1 and .todayPrompts == 1 and .totalSessions == 1 and .todaySessions == 1' >/dev/null || + fail "Hermes collector reports prompt and session totals" +pass "Hermes collector reports prompt and session totals" + +HOME="$TEST_HOME" XDG_STATE_HOME="" "$ROOT/bin/omarchy-agent-usage-hermes" --write || + fail "Hermes collector writes an atomic usage record" +pass "Hermes collector writes an atomic usage record" + +[[ $(jq -r '.modelUsage | keys | join(",")' "$TEST_HOME/.local/state/omarchy/agents/usage/hermes.json") == "gpt-6-astra,gpt-6-sol" ]] || + fail "written Hermes record retains all model names" +pass "written Hermes record retains all model names" From af8c27e7355680c78e9e7301584452fdcb503dd6 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Thu, 1 Oct 2026 21:53:20 -0400 Subject: [PATCH 02/19] Read Hermes usage from HERMES_HOME and every named profile Profile discovery follows #6647 by Mustafa Okur. Drops the --write mode, reports the week's tokens by day like the other collectors, and lists Hermes in the panel's defaults, README, and the manual. Co-authored-by: Mustafa Okur Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-hermes | 270 +++++++++++------------- manual/17-ai.md | 2 +- shell/plugins/agents/README.md | 4 +- shell/plugins/agents/manifest.json | 5 +- test/shell.d/agent-usage-hermes-test.sh | 99 +++++---- 5 files changed, 185 insertions(+), 195 deletions(-) diff --git a/bin/omarchy-agent-usage-hermes b/bin/omarchy-agent-usage-hermes index 0b692086352..33051b27e92 100755 --- a/bin/omarchy-agent-usage-hermes +++ b/bin/omarchy-agent-usage-hermes @@ -1,173 +1,145 @@ #!/usr/bin/python3 # omarchy:summary=Print the Hermes Agent usage record as JSON -# omarchy:args=[--write] +# omarchy:args=[--force] [--limits-only] # omarchy:hidden=true """Collect Hermes Agent usage into one display-ready JSON record. -Everything the agents panel shows for Hermes comes from this one command: -local session stats from ~/.hermes/state.db (sessions, session_model_usage, -messages) and nothing else. It prints the same record contract as the -packaged claude/codex collectors; with --write it updates the record file -atomically instead, so a user systemd timer can keep the panel fresh -- -omarchy-agent-usage-update only discovers packaged collectors. +Everything the agents panel shows for Hermes is local session stats, read +only from the state.db in Hermes' home (HERMES_HOME, else ~/.hermes) and in +each named profile under ~/.hermes/profiles. Hermes has no subscription of +its own, so the record has no limits. Chat sessions that arrived over messaging gateways count too: Hermes treats them as first-class sessions and so does this record. """ -from __future__ import annotations - +import argparse import json import os import shutil import sqlite3 -import sys -import tempfile -from datetime import datetime, timedelta +from datetime import datetime, timedelta, timezone from pathlib import Path -DB_PATH = Path.home() / ".hermes" / "state.db" -USAGE_DIR = Path(os.environ.get("XDG_STATE_HOME") or Path.home() / ".local/state") / "omarchy/agents/usage" RECENT_DAYS = 7 -def day_bounds(day: datetime.date) -> tuple[float, float]: - start = datetime(day.year, day.month, day.day).astimezone() - return start.timestamp(), start.timestamp() + 86400 - - -def collect() -> dict: - now = datetime.now().astimezone() - record = { - "schemaVersion": 1, - "id": "hermes", - "name": "Hermes", - "updatedAt": now.isoformat(), - "ready": shutil.which("hermes") is not None, - "hasLocalStats": False, - "todayPrompts": 0, - "todaySessions": 0, - "todayTotalTokens": 0, - "todayTokensByModel": {}, - "recentDays": [], - "totalPrompts": 0, - "totalSessions": 0, - "activeDays": 0, - "activeDates": [], - "modelUsage": {}, - "limits": [], - "tierLabel": "", - "usageStatusText": "" if DB_PATH.exists() else "Hermes state not found", - "authHelpText": "", - } - - if not DB_PATH.exists(): - return record - - try: - conn = sqlite3.connect(f"file:{DB_PATH}?mode=ro", uri=True) - except sqlite3.Error: - record["usageStatusText"] = "Hermes state unreadable" - return record - +# Hermes anchors named profiles under ~/.hermes whatever HERMES_HOME says, so +# they are looked for there. A profile linked to the home counts once. +def state_paths(): + home = Path(os.environ.get("HERMES_HOME") or (Path.home() / ".hermes")) + candidates = [home / "state.db"] + profiles = Path.home() / ".hermes" / "profiles" + if profiles.is_dir(): + candidates.extend(sorted(path / "state.db" for path in profiles.iterdir() if path.is_dir())) + + paths, seen = [], set() + for candidate in candidates: + if not candidate.is_file(): + continue + resolved = candidate.resolve() + if resolved not in seen: + seen.add(resolved) + paths.append(candidate) + return paths + + +def day_start(day): + return datetime(day.year, day.month, day.day).astimezone().timestamp() + + +# One database's share of the stats, added into the running totals. +def scan(path, totals, today, week): + conn = sqlite3.connect(path.resolve().as_uri() + "?mode=ro", uri=True, timeout=2) + try: + today_start = day_start(today) + week_start = day_start(week[0]) + totals["sessions"] += conn.execute("SELECT COUNT(*) FROM sessions").fetchone()[0] + totals["todaySessions"] += conn.execute("SELECT COUNT(*) FROM sessions WHERE started_at >= ?", (today_start,)).fetchone()[0] + totals["prompts"] += conn.execute("SELECT COUNT(*) FROM messages WHERE role = 'user'").fetchone()[0] + totals["todayPrompts"] += conn.execute("SELECT COUNT(*) FROM messages WHERE role = 'user' AND timestamp >= ?", (today_start,)).fetchone()[0] + + for date in conn.execute("SELECT DISTINCT date(started_at, 'unixepoch', 'localtime') FROM sessions"): + if date[0]: + totals["dates"].add(date[0]) + + # Usage is kept per session and model, so a session's tokens land on the + # day it started. + for day, model, uncached, output, cache_read, cache_write in conn.execute( + """SELECT date(s.started_at, 'unixepoch', 'localtime'), smu.model, + SUM(smu.input_tokens), SUM(smu.output_tokens), + SUM(smu.cache_read_tokens), SUM(smu.cache_write_tokens) + FROM session_model_usage smu + JOIN sessions s ON s.id = smu.session_id + GROUP BY 1, 2""" + ): + bucket = totals["models"].setdefault(model, {"inputTokens": 0, "outputTokens": 0, "cacheReadInputTokens": 0, "cacheCreationInputTokens": 0}) + bucket["inputTokens"] += int(uncached or 0) + bucket["outputTokens"] += int(output or 0) + bucket["cacheReadInputTokens"] += int(cache_read or 0) + bucket["cacheCreationInputTokens"] += int(cache_write or 0) + tokens = int(uncached or 0) + int(output or 0) + int(cache_read or 0) + int(cache_write or 0) + if day in totals["byDay"]: + totals["byDay"][day] += tokens + if day == today.isoformat() and tokens: + totals["todayByModel"][model] = totals["todayByModel"].get(model, 0) + tokens + finally: + conn.close() + + +def collect(): + now = datetime.now().astimezone() + today = now.date() + week = [today - timedelta(days=offset) for offset in range(RECENT_DAYS - 1, -1, -1)] + totals = { + "sessions": 0, "todaySessions": 0, "prompts": 0, "todayPrompts": 0, + "dates": set(), "models": {}, "todayByModel": {}, + "byDay": {day.isoformat(): 0 for day in week}, + } + + paths = state_paths() + status = "" if paths else "Hermes state not found" + read = 0 + for path in paths: try: - today_start, _ = day_bounds(now.date()) - week_start, _ = day_bounds(now.date() - timedelta(days=RECENT_DAYS - 1)) - - record["hasLocalStats"] = True - record["totalSessions"] = conn.execute("SELECT COUNT(*) FROM sessions").fetchone()[0] - record["todaySessions"] = conn.execute( - "SELECT COUNT(*) FROM sessions WHERE started_at >= ?", (today_start,) - ).fetchone()[0] - record["totalPrompts"] = conn.execute( - "SELECT COUNT(*) FROM messages WHERE role = 'user'" - ).fetchone()[0] - record["todayPrompts"] = conn.execute( - "SELECT COUNT(*) FROM messages WHERE role = 'user' AND timestamp >= ?", (today_start,) - ).fetchone()[0] - - for model, tokens in conn.execute( - """SELECT smu.model, - SUM(smu.input_tokens + smu.output_tokens - + smu.cache_read_tokens + smu.cache_write_tokens) - FROM session_model_usage smu - JOIN sessions s ON s.id = smu.session_id - WHERE s.started_at >= ? - GROUP BY smu.model""", - (today_start,), - ): - if tokens: - record["todayTokensByModel"][model] = int(tokens) - record["todayTotalTokens"] = sum(record["todayTokensByModel"].values()) - - per_day = dict( - conn.execute( - """SELECT date(timestamp, 'unixepoch', 'localtime'), COUNT(*) - FROM messages - WHERE role = 'user' AND timestamp >= ? - GROUP BY 1""", - (week_start,), - ) - ) - record["recentDays"] = [ - { - "date": (now.date() - timedelta(days=offset)).isoformat(), - "messageCount": int(per_day.get((now.date() - timedelta(days=offset)).isoformat(), 0)), - } - for offset in range(RECENT_DAYS - 1, -1, -1) - ] - - active_dates = [ - row[0] - for row in conn.execute( - "SELECT DISTINCT date(started_at, 'unixepoch', 'localtime') FROM sessions ORDER BY 1" - ) - if row[0] - ] - record["activeDays"] = len(active_dates) - record["activeDates"] = active_dates[-400:] - - for model, inp, outp, cache_read, cache_write in conn.execute( - """SELECT model, SUM(input_tokens), SUM(output_tokens), - SUM(cache_read_tokens), SUM(cache_write_tokens) - FROM session_model_usage GROUP BY model""" - ): - record["modelUsage"][model] = { - "inputTokens": int(inp or 0), - "outputTokens": int(outp or 0), - "cacheReadInputTokens": int(cache_read or 0), - "cacheCreationInputTokens": int(cache_write or 0), - } + scan(path, totals, today, week) + read += 1 except sqlite3.Error: - record["usageStatusText"] = "Hermes state unreadable" - finally: - conn.close() - - return record - - -def main() -> int: - record = collect() - payload = json.dumps(record, separators=(",", ":")) - - if "--write" not in sys.argv[1:]: - print(payload) - return 0 - - USAGE_DIR.mkdir(parents=True, exist_ok=True) - fd, tmp = tempfile.mkstemp(prefix=".hermes.", dir=USAGE_DIR) - try: - with os.fdopen(fd, "w") as handle: - handle.write(payload + "\n") - os.replace(tmp, USAGE_DIR / "hermes.json") - except OSError: - try: - os.unlink(tmp) - except OSError: - pass - return 1 - return 0 + status = "Hermes state unreadable" + + dates = sorted(totals["dates"]) + return { + "schemaVersion": 1, + "id": "hermes", + "name": "Hermes", + "updatedAt": datetime.now(timezone.utc).isoformat(), + "ready": shutil.which("hermes") is not None, + "hasLocalStats": read > 0, + "todayPrompts": totals["todayPrompts"], + "todaySessions": totals["todaySessions"], + "todayTotalTokens": sum(totals["todayByModel"].values()), + "todayTokensByModel": totals["todayByModel"], + "recentDays": [{"date": day, "messageCount": tokens} for day, tokens in totals["byDay"].items()], + "totalPrompts": totals["prompts"], + "totalSessions": totals["sessions"], + "activeDays": len(dates), + "activeDates": dates[-400:], + "modelUsage": totals["models"], + "limits": [], + "tierLabel": "", + "usageStatusText": status, + "authHelpText": "", + } + + +def main(): + # The shared collector flags are accepted; a local read is always fresh. + parser = argparse.ArgumentParser() + parser.add_argument("--force", action="store_true") + parser.add_argument("--limits-only", action="store_true") + parser.parse_args() + print(json.dumps(collect(), separators=(",", ":"))) if __name__ == "__main__": - sys.exit(main()) + main() diff --git a/manual/17-ai.md b/manual/17-ai.md index ea8a8d65eaf..f9df5ba956a 100644 --- a/manual/17-ai.md +++ b/manual/17-ai.md @@ -38,7 +38,7 @@ There are terminal shortcuts too: `a` runs the default agent inline in the curre ### The agents panel -The top bar has an agents icon. Before you've set up an agent, its panel walks you through signing in to Claude, Codex, or Grok, and the first one becomes your default agent. After that, the panel puts every subscription on one page: each agent and account with its plan and the percentage used of its 5-hour session and weekly limits (or the remaining prepaid balance), under a line that rotates through what you've used across all of them, like tokens this week and your most used model. Under _Make something_ it offers starter prompts for a new theme, plugin, or app that hand the task to your default agent. Apps are built the way Omarchy's own Hype, Monologue, and Omacut are, in C++ and Qt Quick, following the omarchy-app agent skill, and the compiler and Qt pieces they need come installed. Claude Code, Codex, Grok, and Fireworks are covered out of the box. +The top bar has an agents icon. Before you've set up an agent, its panel walks you through signing in to Claude, Codex, or Grok, and the first one becomes your default agent. After that, the panel puts every subscription on one page: each agent and account with its plan and the percentage used of its 5-hour session and weekly limits (or the remaining prepaid balance), under a line that rotates through what you've used across all of them, like tokens this week and your most used model. Under _Make something_ it offers starter prompts for a new theme, plugin, or app that hand the task to your default agent. Apps are built the way Omarchy's own Hype, Monologue, and Omacut are, in C++ and Qt Quick, following the omarchy-app agent skill, and the compiler and Qt pieces they need come installed. Claude Code, Codex, Grok, and Fireworks are covered out of the box, and Hermes shows its sessions and tokens, from every profile, once you've used it. Left-click the bar icon for the panel, right-click to launch your default agent. The usage records behind it are regenerated every 15 minutes by `omarchy agent usage-update`, and the panel can even merge usage from your other machines via a synced folder. See the README under `$OMARCHY_PATH/shell/plugins/agents/` for the full settings. diff --git a/shell/plugins/agents/README.md b/shell/plugins/agents/README.md index c9daebba45e..69beb6675cd 100644 --- a/shell/plugins/agents/README.md +++ b/shell/plugins/agents/README.md @@ -74,6 +74,7 @@ light surfaces — and the bar glyph stands in when there is none. | `codex` | The Codex app-server RPC | native Codex CLI session files (plus pi and opencode sessions) | | `grok` | The credits endpoint behind Grok's `/usage` view (the billing period's included usage) | Each session's `usage.json` (the ledger `grok usage` prints: tokens by model per finished turn), plus `summary.json` for sessions | | `fireworks` | Estimated prepaid balance: configured funding minus rated account costs | Fireworks billing API, grouped by day and model for the last 30 days | +| `hermes` | None (Hermes has no subscription of its own) | `state.db` in Hermes' home and in each named profile under `~/.hermes/profiles`, read-only | When `~/.local/state/omarchy/agents/accounts/.json` registers more than one account, the `claude`, `codex`, and `grok` records @@ -91,7 +92,8 @@ account changed. Claude limits need a signed-in CLI; without credentials the panel says so and falls back to local stats only. A non-default Claude directory is honored via -`CLAUDE_CONFIG_DIR`, Codex via `CODEX_HOME`, Grok via `GROK_HOME`. Grok's +`CLAUDE_CONFIG_DIR`, Codex via `CODEX_HOME`, Grok via `GROK_HOME`, Hermes via +`HERMES_HOME`. Grok's plan comes from the settings it caches in its home, and its limit from the credits endpoint its own `/usage` view reads, asked with each account's sign-in; a sign-in left to lapse shows the last credits until Grok runs diff --git a/shell/plugins/agents/manifest.json b/shell/plugins/agents/manifest.json index 7ca52484b12..6b3e2262650 100644 --- a/shell/plugins/agents/manifest.json +++ b/shell/plugins/agents/manifest.json @@ -5,7 +5,7 @@ "version": "1.0.0", "author": "Omarchy", "license": "MIT", - "description": "Claude Code, Codex, Grok, and Fireworks usage, limits, and pace in a native Omarchy bar panel.", + "description": "Claude Code, Codex, Grok, Fireworks, and Hermes usage, limits, and pace in a native Omarchy bar panel.", "kinds": ["bar-widget"], "activation": "on-demand", "entryPoints": { @@ -22,7 +22,8 @@ "claude": { "enabled": true }, "codex": { "enabled": true }, "grok": { "enabled": true }, - "fireworks": { "enabled": true } + "fireworks": { "enabled": true }, + "hermes": { "enabled": true } }, "refreshIntervalSec": 900, "syncMode": "Off", diff --git a/test/shell.d/agent-usage-hermes-test.sh b/test/shell.d/agent-usage-hermes-test.sh index 230779d1672..9fcd0fcfc47 100755 --- a/test/shell.d/agent-usage-hermes-test.sh +++ b/test/shell.d/agent-usage-hermes-test.sh @@ -8,56 +8,59 @@ require_command python3 TEST_HOME=$(mktemp -d) trap 'rm -rf "$TEST_HOME"' EXIT -mkdir -p "$TEST_HOME/.hermes" -HOME="$TEST_HOME" python3 - <<'PY' -import os +# A Hermes state.db at $1 with one session today, its prompt, and usage for +# the models after it, given as name:base where every count is base plus a +# per-field offset. +make_state() { + mkdir -p "$(dirname "$1")" + python3 - "$@" <<'PY' import sqlite3 +import sys import time -from pathlib import Path -path = Path(os.environ["HOME"]) / ".hermes" / "state.db" +path, models = sys.argv[1], sys.argv[2:] conn = sqlite3.connect(path) conn.executescript( - """ - CREATE TABLE sessions ( - id TEXT PRIMARY KEY, - started_at REAL NOT NULL - ); - CREATE TABLE messages ( - id INTEGER PRIMARY KEY, - session_id TEXT NOT NULL, - role TEXT NOT NULL, - timestamp REAL NOT NULL - ); - CREATE TABLE session_model_usage ( - session_id TEXT NOT NULL, - model TEXT NOT NULL, - input_tokens INTEGER NOT NULL DEFAULT 0, - output_tokens INTEGER NOT NULL DEFAULT 0, - cache_read_tokens INTEGER NOT NULL DEFAULT 0, - cache_write_tokens INTEGER NOT NULL DEFAULT 0, - PRIMARY KEY (session_id, model) - ); - """ + """ + CREATE TABLE sessions ( + id TEXT PRIMARY KEY, + started_at REAL NOT NULL + ); + CREATE TABLE messages ( + id INTEGER PRIMARY KEY, + session_id TEXT NOT NULL, + role TEXT NOT NULL, + timestamp REAL NOT NULL + ); + CREATE TABLE session_model_usage ( + session_id TEXT NOT NULL, + model TEXT NOT NULL, + input_tokens INTEGER NOT NULL DEFAULT 0, + output_tokens INTEGER NOT NULL DEFAULT 0, + cache_read_tokens INTEGER NOT NULL DEFAULT 0, + cache_write_tokens INTEGER NOT NULL DEFAULT 0, + PRIMARY KEY (session_id, model) + ); + """ ) now = time.time() conn.execute("INSERT INTO sessions VALUES (?, ?)", ("session-1", now)) conn.executemany( - "INSERT INTO messages VALUES (?, ?, ?, ?)", - [(1, "session-1", "user", now), (2, "session-1", "assistant", now)], -) -conn.executemany( - "INSERT INTO session_model_usage VALUES (?, ?, ?, ?, ?, ?)", - [ - ("session-1", "gpt-6-astra", 10, 20, 30, 40), - ("session-1", "gpt-6-sol", 11, 21, 31, 41), - ], + "INSERT INTO messages VALUES (?, ?, ?, ?)", + [(1, "session-1", "user", now), (2, "session-1", "assistant", now)], ) +for spec in models: + name, base = spec.split(":") + base = int(base) + conn.execute("INSERT INTO session_model_usage VALUES (?, ?, ?, ?, ?, ?)", ("session-1", name, base, base + 10, base + 20, base + 30)) conn.commit() conn.close() PY +} -record=$(HOME="$TEST_HOME" "$ROOT/bin/omarchy-agent-usage-hermes") || +make_state "$TEST_HOME/.hermes/state.db" gpt-6-astra:10 gpt-6-sol:11 + +record=$(HOME="$TEST_HOME" HERMES_HOME="" "$ROOT/bin/omarchy-agent-usage-hermes" --force) || fail "Hermes collector exits successfully" pass "Hermes collector exits successfully" @@ -69,14 +72,26 @@ printf '%s\n' "$record" | jq -e '.todayTokensByModel["gpt-6-sol"] == 104 and .to fail "Hermes collector reports today's per-model totals" pass "Hermes collector reports today's per-model totals" +printf '%s\n' "$record" | jq -e '.recentDays | length == 7 and .[-1].messageCount == 204' >/dev/null || + fail "Hermes collector reports the last week's tokens by day" +pass "Hermes collector reports the last week's tokens by day" + printf '%s\n' "$record" | jq -e '.totalPrompts == 1 and .todayPrompts == 1 and .totalSessions == 1 and .todaySessions == 1' >/dev/null || fail "Hermes collector reports prompt and session totals" pass "Hermes collector reports prompt and session totals" -HOME="$TEST_HOME" XDG_STATE_HOME="" "$ROOT/bin/omarchy-agent-usage-hermes" --write || - fail "Hermes collector writes an atomic usage record" -pass "Hermes collector writes an atomic usage record" +# Named profiles live under ~/.hermes/profiles even when HERMES_HOME moves the +# default home elsewhere. +make_state "$TEST_HOME/elsewhere/state.db" gpt-6-astra:100 +make_state "$TEST_HOME/.hermes/profiles/work/state.db" gpt-6-sol:1000 +record=$(HOME="$TEST_HOME" HERMES_HOME="$TEST_HOME/elsewhere" "$ROOT/bin/omarchy-agent-usage-hermes") +printf '%s\n' "$record" | jq -e '.modelUsage["gpt-6-astra"].inputTokens == 100 and .modelUsage["gpt-6-sol"].inputTokens == 1000 and .totalSessions == 2 and .totalPrompts == 2' >/dev/null || + fail "Hermes collector reads HERMES_HOME and every named profile" +pass "Hermes collector reads HERMES_HOME and every named profile" -[[ $(jq -r '.modelUsage | keys | join(",")' "$TEST_HOME/.local/state/omarchy/agents/usage/hermes.json") == "gpt-6-astra,gpt-6-sol" ]] || - fail "written Hermes record retains all model names" -pass "written Hermes record retains all model names" +empty_home=$(mktemp -d) +record=$(HOME="$empty_home" HERMES_HOME="" PATH=/usr/bin:/bin "$ROOT/bin/omarchy-agent-usage-hermes") +rm -rf "$empty_home" +printf '%s\n' "$record" | jq -e '.limits == [] and .tierLabel == "" and .totalSessions == 0 and .activeDays == 0 and .modelUsage == {}' >/dev/null || + fail "Hermes collector prints a record the panel skips without Hermes" +pass "Hermes collector prints a record the panel skips without Hermes" From d29582b3989108e2eedb25add50a97a9e6df32c3 Mon Sep 17 00:00:00 2001 From: steph4n <284007938+steph4n-gh@users.noreply.github.com> Date: Thu, 1 Oct 2026 21:54:09 -0400 Subject: [PATCH 03/19] Add Google Antigravity usage collector Add the omarchy-agent-usage-antigravity collector to monitor Google Antigravity usage alongside the other agents in the agents bar panel. Authoritative subscription tier and quota limits (rolling 5-hour session, weekly limits, and scoped 3P model allowances) are probed from the Google Cloud Code API via Secret Service keyring credentials with rolling local history estimation as fallback. --- bin/omarchy-agent-usage-antigravity | 737 ++++++++++++++++++ shell/plugins/agents/README.md | 5 +- .../agents/assets/antigravity-light.svg | 1 + shell/plugins/agents/assets/antigravity.svg | 1 + shell/plugins/agents/manifest.json | 5 +- .../agent-usage-antigravity-limits-test.sh | 136 ++++ .../agent-usage-antigravity-scanner-test.sh | 115 +++ 7 files changed, 997 insertions(+), 3 deletions(-) create mode 100755 bin/omarchy-agent-usage-antigravity create mode 100644 shell/plugins/agents/assets/antigravity-light.svg create mode 100644 shell/plugins/agents/assets/antigravity.svg create mode 100755 test/shell.d/agent-usage-antigravity-limits-test.sh create mode 100755 test/shell.d/agent-usage-antigravity-scanner-test.sh diff --git a/bin/omarchy-agent-usage-antigravity b/bin/omarchy-agent-usage-antigravity new file mode 100755 index 00000000000..4f205e1dd76 --- /dev/null +++ b/bin/omarchy-agent-usage-antigravity @@ -0,0 +1,737 @@ +#!/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. + +Local stats come from Antigravity conversation transcripts and history logs +under ~/.gemini/antigravity-cli/. Limits, tier, and session window calculations +come from the Google Cloud Code API (v1internal:loadCodeAssist and +v1internal:retrieveUserQuotaSummary) with local rolling calculations as fallback. +The agents panel only ever reads the JSON this prints. +""" + +from __future__ import annotations + +import argparse +import base64 +import fcntl +import json +import os +import shutil +import sqlite3 +import subprocess +import sys +import tempfile +import time +import urllib.error +import urllib.request +from datetime import datetime, timedelta, timezone +from pathlib import Path +from typing import Any + +AGENT_ID = "antigravity" +AGENT_NAME = "Antigravity" +AUTH_HELP = "Run `agy` to authenticate." + +SCAN_REUSE_SECONDS = 20 +LIMITS_ONLY_REUSE_SECONDS = 900 +PROBE_MIN_INTERVAL_SECONDS = 15 + + +def app_root() -> Path: + return Path(os.environ.get("AGY_DIR", Path.home() / ".gemini/antigravity-cli")).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 scan_cache_paths() -> tuple[Path, Path]: + root = cache_root() + return root / f"{AGENT_ID}-stats-cache.json", root / f"{AGENT_ID}-stats-cache.lock" + + +def read_fresh_json(path: Path, max_age_seconds: float) -> Any: + try: + stat = path.stat() + except OSError: + return None + age = time.time() - stat.st_mtime + if age > max_age_seconds: + return None + try: + with path.open("r", encoding="utf-8") as f: + return json.load(f) + except Exception: + return None + + +def write_json(path: Path, payload: Any) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + fd, tmp_path = tempfile.mkstemp(prefix=f".{path.name}.", dir=str(path.parent)) + tmp = Path(tmp_path) + try: + with os.fdopen(fd, "w", encoding="utf-8") as f: + json.dump(payload, f, separators=(",", ":")) + f.write("\n") + tmp.chmod(0o644) + tmp.replace(path) + except BaseException: + tmp.unlink(missing_ok=True) + raise + + +def check_auth() -> tuple[bool, str, str, str, int]: + """Check authentication state via Secret Service keyring.""" + try: + raw = subprocess.check_output( + ["secret-tool", "lookup", "service", "gemini", "username", "antigravity"], + text=True, + stderr=subprocess.DEVNULL, + timeout=3, + ) + if raw.strip(): + data = json.loads(raw) + auth_method = data.get("auth_method", "consumer") + token_dict = data.get("token") or {} + access_token = str(token_dict.get("access_token") or data.get("access_token") or "") + expiry_raw = str(token_dict.get("expiry") or data.get("expiry") or "") + expires_at_ms = 0 + if expiry_raw: + try: + if expiry_raw.isdigit(): + expires_at_ms = int(expiry_raw) + if expires_at_ms < 1e12: + expires_at_ms *= 1000 + else: + dt = datetime.fromisoformat(expiry_raw.replace("Z", "+00:00")) + expires_at_ms = int(dt.timestamp() * 1000) + except Exception: + expires_at_ms = 0 + + id_token = data.get("id_token") + email = "" + if id_token: + parts = id_token.split(".") + if len(parts) >= 2: + padded = parts[1] + "=" * ((4 - len(parts[1]) % 4) % 4) + claims = json.loads(base64.urlsafe_b64decode(padded)) + email = claims.get("email", "") + return True, email, auth_method, access_token, expires_at_ms + except Exception: + pass + return False, "", "consumer", "", 0 + + +def normalize_tier_label(tier_name: str, tier_id: str = "") -> str: + tier_id_lower = tier_id.lower() + tier_name_lower = tier_name.lower() + if "ultra" in tier_id_lower or "ultra" in tier_name_lower: + return "Ultra" + if "pro" in tier_id_lower or "pro" in tier_name_lower: + return "Pro" + if "free" in tier_id_lower or "free" in tier_name_lower: + return "Free" + if tier_name.lower().startswith("google ai "): + return tier_name[10:].strip() + if tier_name: + return tier_name + if tier_id: + return tier_id.replace("-tier", "").capitalize() + return "Free" + + +def extract_tier(data: dict[str, Any], auth_method: str = "consumer") -> str: + paid_tier = data.get("paidTier") + if isinstance(paid_tier, dict): + t_name = str(paid_tier.get("name") or "").strip() + t_id = str(paid_tier.get("id") or "").strip() + if t_name or t_id: + return normalize_tier_label(t_name, t_id) + + current_tier = data.get("currentTier") + if isinstance(current_tier, dict): + c_id = str(current_tier.get("id") or "").strip() + c_name = str(current_tier.get("name") or "").strip() + if c_id or c_name: + return normalize_tier_label(c_name, c_id) + + if auth_method == "enterprise": + return "Enterprise" + if auth_method in ("cloud", "vertex"): + return "Cloud" + return "Free" + + +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 datetime.fromtimestamp(ts / 1000, timezone.utc).isoformat() + except Exception: + return raw + try: + parsed = datetime.fromisoformat(raw.replace("Z", "+00:00")) + return parsed.isoformat() + except Exception: + return raw + + +def extract_limits(data: dict[str, Any]) -> list[dict[str, Any]]: + groups = data.get("groups") + if not isinstance(groups, list): + return [] + + gemini_group = None + other_groups = [] + for g in groups: + if not isinstance(g, dict): + continue + name = str(g.get("displayName") or "").lower() + if "gemini" in name or gemini_group is None: + gemini_group = g + else: + other_groups.append(g) + + limits = [] + # Primary group: gemini models + if gemini_group: + buckets = gemini_group.get("buckets") or [] + bucket_5h = None + bucket_weekly = None + for b in buckets: + if not isinstance(b, dict): + continue + w = str(b.get("window") or "").lower() + bid = str(b.get("bucketId") or "").lower() + dname = str(b.get("displayName") or "").lower() + if w == "5h" or "5h" in bid or "five" in dname: + bucket_5h = b + elif w == "weekly" or "weekly" in bid or "week" in dname: + bucket_weekly = b + + # Enforce Session (5-hour) first, then Weekly (7-day) + if bucket_5h: + rem = float(bucket_5h.get("remainingFraction", 1.0)) + used = round(max(0.0, min(1.0, 1.0 - rem)), 4) + limits.append({ + "label": "Session (5-hour)", + "title": "Session", + "percent": used, + "resetsAt": normalize_reset_at(bucket_5h.get("resetTime")), + }) + if bucket_weekly: + rem = float(bucket_weekly.get("remainingFraction", 1.0)) + used = round(max(0.0, min(1.0, 1.0 - rem)), 4) + limits.append({ + "label": "Weekly (7-day)", + "title": "Weekly", + "percent": used, + "resetsAt": normalize_reset_at(bucket_weekly.get("resetTime")), + }) + + # Scoped limits for other groups (e.g. Claude and GPT models) + for g in other_groups: + g_name = str(g.get("displayName") or "").strip() + short_name = "Claude/GPT" if "claude" in g_name.lower() else (g_name or "3P") + for b in (g.get("buckets") or []): + if not isinstance(b, dict): + continue + rem = float(b.get("remainingFraction", 1.0)) + used = round(max(0.0, min(1.0, 1.0 - rem)), 4) + # Only show secondary model limits if there has been usage + if used > 0: + w = str(b.get("window") or "").lower() + w_title = "Weekly" if "week" in w else "Session" + title = f"{short_name} {w_title}" + limits.append({ + "label": title, + "title": title, + "percent": used, + "resetsAt": normalize_reset_at(b.get("resetTime")), + }) + + return limits + + +def probe_antigravity(access_token: str, endpoint_base: str = "") -> dict[str, Any]: + if not endpoint_base: + endpoint_base = os.environ.get("CLOUD_CODE_URL", "https://daily-cloudcode-pa.googleapis.com").rstrip("/") + + headers = { + "Authorization": f"Bearer {access_token}", + "Content-Type": "application/json", + "User-Agent": "antigravity", + } + + result: dict[str, Any] = { + "ok": False, + "tierLabel": "", + "limits": [], + "helpText": "", + "rateLimited": False, + "transport": False, + } + + tier = "" + try: + req = urllib.request.Request(f"{endpoint_base}/v1internal:loadCodeAssist", data=b"{}", headers=headers) + with urllib.request.urlopen(req, timeout=8) as resp: + payload = json.loads(resp.read().decode("utf-8", errors="replace")) + tier = extract_tier(payload) + except urllib.error.HTTPError as err: + if err.code == 429: + result["rateLimited"] = True + elif err.code == 401: + result["helpText"] = "Antigravity authentication expired. Run `agy` to authenticate." + return result + except Exception: + result["transport"] = True + result["helpText"] = "Could not reach Antigravity service." + + limits = [] + try: + req = urllib.request.Request(f"{endpoint_base}/v1internal:retrieveUserQuotaSummary", data=b"{}", headers=headers) + with urllib.request.urlopen(req, timeout=8) as resp: + payload = json.loads(resp.read().decode("utf-8", errors="replace")) + limits = extract_limits(payload) + except urllib.error.HTTPError as err: + if err.code == 429: + result["rateLimited"] = True + elif err.code == 401: + result["helpText"] = "Antigravity authentication expired. Run `agy` to authenticate." + return result + except Exception: + result["transport"] = True + + result["ok"] = bool(tier or limits) + result["tierLabel"] = tier + result["limits"] = limits + return result + + +def limit_window_open(entry: dict[str, Any], now: datetime) -> bool: + raw = str(entry.get("resetsAt") or "") + if raw == "": + return True + try: + resets_at = datetime.fromisoformat(raw.replace("Z", "+00:00")) + except Exception: + return True + if resets_at.tzinfo is None: + resets_at = resets_at.replace(tzinfo=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 = datetime.now(timezone.utc) + return [entry for entry in entries if isinstance(entry, dict) and limit_window_open(entry, now)] + + +def detect_tier_label( + root_dir: Path, + auth_ok: bool, + email: str, + auth_method: str = "consumer", + probed_tier: str = "", + cached_tier: str = "", +) -> str: + if not auth_ok: + return "" + + if os.environ.get("AGY_TIER"): + return os.environ["AGY_TIER"] + + cfg_file = Path.home() / ".config" / "omarchy" / "agents" / "antigravity.json" + if cfg_file.exists(): + try: + with cfg_file.open("r", encoding="utf-8") as f: + data = json.load(f) + if data.get("tier"): + return str(data["tier"]) + except Exception: + pass + + if probed_tier: + return probed_tier + + if cached_tier: + return cached_tier + + if auth_method == "enterprise": + return "Enterprise" + if auth_method in ("cloud", "vertex"): + return "Cloud" + + return "Pro" + + +def collect_local_stats(root_dir: Path, now: datetime) -> dict[str, Any]: + today_str = now.strftime("%Y-%m-%d") + recent_dates = [(now - timedelta(days=i)).strftime("%Y-%m-%d") for i in range(6, -1, -1)] + recent = {d: {"date": d, "messageCount": 0} for d in recent_dates} + + history_file = root_dir / "history.jsonl" + brain_dir = root_dir / "brain" + + today_prompts = 0 + total_prompts = 0 + active_dates: set[str] = set() + today_sessions: set[str] = set() + total_sessions: set[str] = set() + + if history_file.exists(): + try: + with history_file.open("r", encoding="utf-8", errors="ignore") as f: + for line in f: + line = line.strip() + if not line: + continue + try: + entry = json.loads(line) + ts = entry.get("timestamp") + if not ts: + continue + try: + if isinstance(ts, (int, float)): + dt = datetime.fromtimestamp(ts / 1000.0, timezone.utc) + else: + dt = datetime.fromisoformat(str(ts).replace("Z", "+00:00")) + except Exception: + continue + d_str = dt.strftime("%Y-%m-%d") + active_dates.add(d_str) + total_prompts += 1 + conv_id = entry.get("conversationId") or entry.get("conversation_id") + if conv_id: + total_sessions.add(conv_id) + if d_str == today_str: + today_prompts += 1 + if conv_id: + today_sessions.add(conv_id) + except Exception: + continue + except Exception: + pass + + # Fallback to conversation summaries if history.jsonl is sparse + summaries_db = root_dir / "conversation_summaries.db" + if summaries_db.exists(): + try: + conn = sqlite3.connect(f"file:{summaries_db}?mode=ro", uri=True) + cur = conn.cursor() + for row in cur.execute("SELECT conversation_id, last_modified_time FROM conversation_summaries"): + cid, lmt = row[0], str(row[1] or "") + if cid: + total_sessions.add(cid) + if lmt.startswith(today_str): + today_sessions.add(cid) + conn.close() + except Exception: + pass + + # Compute token counts across conversation transcripts + model_usage: dict[str, dict[str, int]] = {} + today_tokens = 0 + today_tokens_by_model: dict[str, int] = {} + + if brain_dir.exists(): + try: + for conv_folder in brain_dir.iterdir(): + if not conv_folder.is_dir(): + continue + tr_file = conv_folder / ".system_generated/logs/transcript.jsonl" + if not tr_file.exists(): + continue + try: + with tr_file.open("r", encoding="utf-8", errors="ignore") as f: + for line in f: + line = line.strip() + if not line: + continue + try: + step = json.loads(line) + m_name = step.get("model") or "gemini-3.8-flash-high" + if m_name not in model_usage: + model_usage[m_name] = { + "inputTokens": 0, + "outputTokens": 0, + "cacheReadInputTokens": 0, + "cacheCreationInputTokens": 0, + } + + step_usage = step.get("usage") + is_model = step.get("source") == "MODEL" or step.get("type") == "PLANNER_RESPONSE" + if isinstance(step_usage, dict): + inp = int(step_usage.get("prompt_token_count") or step_usage.get("input_tokens") or 0) + out = int(step_usage.get("candidates_token_count") or step_usage.get("output_tokens") or 0) + tok = inp + out + model_usage[m_name]["inputTokens"] += inp + model_usage[m_name]["outputTokens"] += out + else: + c_text = step.get("content") or "" + t_text = step.get("thinking") or "" + tok = max(1, (len(c_text) + len(t_text)) // 4) + if is_model: + model_usage[m_name]["outputTokens"] += tok + else: + model_usage[m_name]["inputTokens"] += tok + + created = step.get("created_at") + is_today = False + if created: + step_d = created[:10] + if step_d in recent: + recent[step_d]["messageCount"] += tok + if step_d == today_str: + is_today = True + else: + is_today = True + + if is_today: + today_tokens += tok + today_tokens_by_model[m_name] = today_tokens_by_model.get(m_name, 0) + tok + except Exception: + continue + except Exception: + continue + except Exception: + pass + + return { + "todayPrompts": today_prompts, + "todaySessions": len(today_sessions), + "todayTotalTokens": today_tokens, + "todayTokensByModel": today_tokens_by_model, + "recentDays": [recent[d] for d in recent_dates], + "totalPrompts": total_prompts, + "totalSessions": len(total_sessions), + "activeDays": len(active_dates), + "activeDates": sorted(list(active_dates)), + "modelUsage": model_usage, + } + + +def calculate_limits(root_dir: Path, now: datetime) -> list[dict[str, Any]]: + """Fallback calculation of rolling limits from local history files.""" + history_file = root_dir / "history.jsonl" + five_hours_ago_ts = int((now - timedelta(hours=5)).timestamp() * 1000) + + prompts_5h = 0 + oldest_5h_ts = None + prompts_7d = 0 + + seven_days_ago_ts = int((now - timedelta(days=7)).timestamp() * 1000) + + if history_file.exists(): + try: + with history_file.open("r", encoding="utf-8", errors="ignore") as f: + for line in f: + line = line.strip() + if not line: + continue + try: + entry = json.loads(line) + ts = entry.get("timestamp") + if not ts: + continue + if ts >= seven_days_ago_ts: + prompts_7d += 1 + if ts >= five_hours_ago_ts: + prompts_5h += 1 + if oldest_5h_ts is None or ts < oldest_5h_ts: + oldest_5h_ts = ts + except Exception: + continue + except Exception: + pass + + # Rolling 5h reset + if oldest_5h_ts: + reset_5h = datetime.fromtimestamp(oldest_5h_ts / 1000.0, timezone.utc) + timedelta(hours=5) + else: + reset_5h = now + timedelta(hours=5) + + # Check if a 429 / resource exhausted error happened recently + cli_log = root_dir / "cli.log" + rate_limited = False + if cli_log.exists(): + try: + stat = cli_log.stat() + if (time.time() - stat.st_mtime) < 18000: + with cli_log.open("r", encoding="utf-8", errors="ignore") as f: + content = f.read()[-8192:] + if "RESOURCE_EXHAUSTED" in content or "rate limit exceeded" in content.lower(): + rate_limited = True + except Exception: + pass + + session_pct = 1.0 if rate_limited else round(min(1.0, prompts_5h / 50.0), 2) + weekly_pct = round(min(1.0, prompts_7d / 250.0), 2) + + return [ + { + "label": "Session (5-hour)", + "title": "Session", + "percent": session_pct, + "resetsAt": reset_5h.strftime("%Y-%m-%dT%H:%M:%S+00:00"), + }, + { + "label": "Weekly (7-day)", + "title": "Weekly", + "percent": weekly_pct, + "resetsAt": (now + timedelta(days=7)).strftime("%Y-%m-%dT%H:%M:%S+00:00"), + }, + ] + + +def collect_limits( + root_dir: Path, + now: datetime, + auth_ok: bool, + access_token: str, + expires_at_ms: int, + auth_method: str, + force: bool, +) -> tuple[list[dict[str, Any]], str, str, str]: + """Collect limits via API probe, falling back to cached or rolling local stats.""" + limits_cache = cache_root() / f"{AGENT_ID}-limits.json" + cached = read_fresh_json(limits_cache, float("inf")) or {} + fallback = usable_cached_limits(cached) + cached_tier = str(cached.get("tierLabel") or "") + + if not auth_ok: + return [], "", "Sign in to see limits", AUTH_HELP + + # If access token is empty, we cannot probe the remote endpoint + if not access_token: + limits = fallback if fallback else calculate_limits(root_dir, now) + return limits, cached_tier, "", AUTH_HELP + + # Check reuse window for limits probe (PROBE_MIN_INTERVAL_SECONDS) + fetched_at = float(cached.get("fetchedAtMs") or 0) / 1000.0 + if fallback and not force and (time.time() - fetched_at) < PROBE_MIN_INTERVAL_SECONDS: + return fallback, cached_tier, "", AUTH_HELP + + # Perform live probe + probe = probe_antigravity(access_token) + if probe["ok"]: + probed_limits = probe["limits"] + probed_tier = probe["tierLabel"] + write_json( + limits_cache, + { + "fetchedAtMs": round(time.time() * 1000), + "limits": probed_limits, + "tierLabel": probed_tier, + }, + ) + return probed_limits, probed_tier, "", AUTH_HELP + + # Probe failed + help_text = probe["helpText"] if probe.get("helpText") else AUTH_HELP + usage_status = "" + if probe.get("rateLimited"): + usage_status = "Antigravity limits rate limited" + elif probe.get("transport"): + usage_status = "" + + if fallback: + return fallback, cached_tier, usage_status, help_text + + # Fallback to local rolling calculation + local_limits = calculate_limits(root_dir, now) + return local_limits, cached_tier, usage_status, help_text + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--force", action="store_true", help="Bypass local caches and force a fresh scan") + parser.add_argument("--limits-only", action="store_true", help="Recompute only rate limits if cache is fresh") + args = parser.parse_args() + + root_dir = app_root() + now = datetime.now(timezone.utc) + today_str = now.strftime("%Y-%m-%d") + + auth_ok, email, auth_method, access_token, expires_at_ms = check_auth() + + cache_file, lock_file = scan_cache_paths() + stats = None + + if not args.force: + max_age = LIMITS_ONLY_REUSE_SECONDS if args.limits_only else SCAN_REUSE_SECONDS + cached = read_fresh_json(cache_file, max_age) + if isinstance(cached, dict) and cached.get("scanDate") == today_str: + stats = cached.get("stats") + + if stats is None: + try: + with lock_file.open("w") as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + stats = collect_local_stats(root_dir, now) + write_json(cache_file, {"schemaVersion": 1, "scanDate": today_str, "stats": stats}) + except Exception as exc: + print(f"omarchy-agent-usage-antigravity: cache lock failed ({exc}), scanning directly", file=sys.stderr) + stats = collect_local_stats(root_dir, now) + + limits, probed_tier, usage_status_text, auth_help_text = collect_limits( + root_dir, + now, + auth_ok, + access_token, + expires_at_ms, + auth_method, + args.force, + ) + + cached_limits_file = cache_root() / f"{AGENT_ID}-limits.json" + cached_payload = read_fresh_json(cached_limits_file, float("inf")) or {} + cached_tier = str(cached_payload.get("tierLabel") or "") + tier_label = detect_tier_label(root_dir, auth_ok, email, auth_method, probed_tier, cached_tier) + + total_prompts = stats.get("totalPrompts", 0) + ready = auth_ok or total_prompts > 0 or len(limits) > 0 + has_local_stats = total_prompts > 0 + + record = { + "schemaVersion": 1, + "id": AGENT_ID, + "name": AGENT_NAME, + "updatedAt": now.strftime("%Y-%m-%dT%H:%M:%S.%f+00:00"), + "ready": ready, + "hasLocalStats": has_local_stats, + "todayPrompts": stats.get("todayPrompts", 0), + "todaySessions": stats.get("todaySessions", 0), + "todayTotalTokens": stats.get("todayTotalTokens", 0), + "todayTokensByModel": stats.get("todayTokensByModel", {}), + "recentDays": stats.get("recentDays", []), + "totalPrompts": total_prompts, + "totalSessions": stats.get("totalSessions", 0), + "activeDays": stats.get("activeDays", 0), + "activeDates": stats.get("activeDates", []), + "modelUsage": stats.get("modelUsage", {}), + "limits": limits, + "tierLabel": tier_label, + "usageStatusText": usage_status_text, + "authHelpText": auth_help_text, + } + + print(json.dumps(record, separators=(",", ":"))) + + +if __name__ == "__main__": + main() diff --git a/shell/plugins/agents/README.md b/shell/plugins/agents/README.md index 69beb6675cd..dccd2f6cdcc 100644 --- a/shell/plugins/agents/README.md +++ b/shell/plugins/agents/README.md @@ -75,6 +75,7 @@ light surfaces — and the bar glyph stands in when there is none. | `grok` | The credits endpoint behind Grok's `/usage` view (the billing period's included usage) | Each session's `usage.json` (the ledger `grok usage` prints: tokens by model per finished turn), plus `summary.json` for sessions | | `fireworks` | Estimated prepaid balance: configured funding minus rated account costs | Fireworks billing API, grouped by day and model for the last 30 days | | `hermes` | None (Hermes has no subscription of its own) | `state.db` in Hermes' home and in each named profile under `~/.hermes/profiles`, read-only | +| `antigravity` | Google Cloud Code API (`loadCodeAssist` and `retrieveUserQuotaSummary`) with rolling session/weekly calculation fallback | Transcripts and history logs under `~/.gemini/antigravity-cli/` (`AGY_DIR`) | When `~/.local/state/omarchy/agents/accounts/.json` registers more than one account, the `claude`, `codex`, and `grok` records @@ -90,7 +91,9 @@ since every account shares the primary home's history. After each run, active account crosses its threshold, and re-collects the record if the active account changed. -Claude limits need a signed-in CLI; without credentials the panel says so and +Antigravity checks authentication via the Secret Service keyring +(`gemini`/`antigravity`), queries Google's Cloud Code endpoints for account +tier and quota limits, and honors a custom app root via `AGY_DIR`. Claude limits need a signed-in CLI; without credentials the panel says so and falls back to local stats only. A non-default Claude directory is honored via `CLAUDE_CONFIG_DIR`, Codex via `CODEX_HOME`, Grok via `GROK_HOME`, Hermes via `HERMES_HOME`. Grok's diff --git a/shell/plugins/agents/assets/antigravity-light.svg b/shell/plugins/agents/assets/antigravity-light.svg new file mode 100644 index 00000000000..67ced694468 --- /dev/null +++ b/shell/plugins/agents/assets/antigravity-light.svg @@ -0,0 +1 @@ +Antigravity diff --git a/shell/plugins/agents/assets/antigravity.svg b/shell/plugins/agents/assets/antigravity.svg new file mode 100644 index 00000000000..3743be8edf4 --- /dev/null +++ b/shell/plugins/agents/assets/antigravity.svg @@ -0,0 +1 @@ +Antigravity diff --git a/shell/plugins/agents/manifest.json b/shell/plugins/agents/manifest.json index 6b3e2262650..6a90aea67ad 100644 --- a/shell/plugins/agents/manifest.json +++ b/shell/plugins/agents/manifest.json @@ -5,7 +5,7 @@ "version": "1.0.0", "author": "Omarchy", "license": "MIT", - "description": "Claude Code, Codex, Grok, Fireworks, and Hermes usage, limits, and pace in a native Omarchy bar panel.", + "description": "Claude Code, Codex, Grok, Fireworks, Hermes, and Antigravity usage, limits, and pace in a native Omarchy bar panel.", "kinds": ["bar-widget"], "activation": "on-demand", "entryPoints": { @@ -23,7 +23,8 @@ "codex": { "enabled": true }, "grok": { "enabled": true }, "fireworks": { "enabled": true }, - "hermes": { "enabled": true } + "hermes": { "enabled": true }, + "antigravity": { "enabled": true } }, "refreshIntervalSec": 900, "syncMode": "Off", diff --git a/test/shell.d/agent-usage-antigravity-limits-test.sh b/test/shell.d/agent-usage-antigravity-limits-test.sh new file mode 100755 index 00000000000..89c0dbdc945 --- /dev/null +++ b/test/shell.d/agent-usage-antigravity-limits-test.sh @@ -0,0 +1,136 @@ +#!/bin/bash + +source "$(dirname "$0")/base-test.sh" + +require_command jq +require_command python3 + +COLLECTOR="$ROOT/bin/omarchy-agent-usage-antigravity" + +# probe_antigravity reaches Google's Cloud Code endpoints, so the readers that +# interpret the answers are exercised with recorded payloads standing in. + +test_tier() { + COLLECTOR="$COLLECTOR" PAYLOAD="$1" python3 - <<'PY' +import importlib.machinery, importlib.util, json, os + +loader = importlib.machinery.SourceFileLoader("collector", os.environ["COLLECTOR"]) +spec = importlib.util.spec_from_loader(loader.name, loader) +collector = importlib.util.module_from_spec(spec) +loader.exec_module(collector) + +payload = json.loads(os.environ["PAYLOAD"]) +print(collector.extract_tier(payload)) +PY +} + +ultra_tier=$(test_tier '{ + "currentTier": { "id": "free-tier", "name": "Antigravity" }, + "paidTier": { "id": "g1-ultra-tier", "name": "Google AI Ultra" } +}') +[[ $ultra_tier == "Ultra" ]] || + fail "Antigravity collector extracts Ultra tier from paidTier" "$ultra_tier" +pass "Antigravity collector extracts Ultra tier from paidTier" + +pro_tier=$(test_tier '{ + "currentTier": { "id": "free-tier", "name": "Antigravity" }, + "paidTier": { "id": "g1-pro-tier", "name": "Google AI Pro" } +}') +[[ $pro_tier == "Pro" ]] || + fail "Antigravity collector extracts Pro tier from paidTier" "$pro_tier" +pass "Antigravity collector extracts Pro tier from paidTier" + +free_tier=$(test_tier '{ + "currentTier": { "id": "free-tier", "name": "Antigravity" } +}') +[[ $free_tier == "Free" ]] || + fail "Antigravity collector extracts Free tier when no paidTier" "$free_tier" +pass "Antigravity collector extracts Free tier when no paidTier" + +# Test quota extraction +test_limits() { + COLLECTOR="$COLLECTOR" PAYLOAD="$1" python3 - <<'PY' +import importlib.machinery, importlib.util, json, os + +loader = importlib.machinery.SourceFileLoader("collector", os.environ["COLLECTOR"]) +spec = importlib.util.spec_from_loader(loader.name, loader) +collector = importlib.util.module_from_spec(spec) +loader.exec_module(collector) + +payload = json.loads(os.environ["PAYLOAD"]) +print(json.dumps(collector.extract_limits(payload))) +PY +} + +quota_payload='{ + "groups": [ + { + "displayName": "Gemini Models", + "buckets": [ + { + "bucketId": "gemini-weekly", + "displayName": "Weekly Limit Remaining", + "window": "weekly", + "resetTime": "2026-09-17T18:19:42Z", + "remainingFraction": 0.95 + }, + { + "bucketId": "gemini-5h", + "displayName": "Five Hour Limit Remaining", + "window": "5h", + "resetTime": "2026-09-15T19:31:46Z", + "remainingFraction": 0.90 + } + ] + }, + { + "displayName": "Claude and GPT models", + "buckets": [ + { + "bucketId": "3p-weekly", + "displayName": "Weekly Limit Remaining", + "window": "weekly", + "resetTime": "2026-09-18T12:56:14Z", + "remainingFraction": 0.98 + }, + { + "bucketId": "3p-5h", + "displayName": "Five Hour Limit Remaining", + "window": "5h", + "resetTime": "2026-09-15T22:16:22Z", + "remainingFraction": 1.0 + } + ] + } + ] +}' + +limits=$(test_limits "$quota_payload") + +[[ $(jq -r '.[0].label' <<<"$limits") == "Session (5-hour)" ]] || + fail "Antigravity collector places Session window first" "$limits" +pass "Antigravity collector places Session window first" + +[[ $(jq -r '.[0].percent' <<<"$limits") == "0.1" ]] || + fail "Antigravity collector calculates used percentage from remainingFraction" "$limits" +pass "Antigravity collector calculates used percentage from remainingFraction" + +[[ $(jq -r '.[1].label' <<<"$limits") == "Weekly (7-day)" ]] || + fail "Antigravity collector places Weekly window second" "$limits" +pass "Antigravity collector places Weekly window second" + +[[ $(jq -r '.[1].percent' <<<"$limits") == "0.05" ]] || + fail "Antigravity collector calculates weekly used percentage" "$limits" +pass "Antigravity collector calculates weekly used percentage" + +[[ $(jq -r '.[2].title' <<<"$limits") == "Claude/GPT Weekly" ]] || + fail "Antigravity collector includes used scoped 3P limits" "$limits" +pass "Antigravity collector includes used scoped 3P limits" + +[[ $(jq -r '.[2].percent' <<<"$limits") == "0.02" ]] || + fail "Antigravity collector calculates 3P used percentage" "$limits" +pass "Antigravity collector calculates 3P used percentage" + +[[ $(jq -r 'length' <<<"$limits") == "3" ]] || + fail "Antigravity collector omits unused 3P session limit" "$limits" +pass "Antigravity collector omits unused 3P session limit" diff --git a/test/shell.d/agent-usage-antigravity-scanner-test.sh b/test/shell.d/agent-usage-antigravity-scanner-test.sh new file mode 100755 index 00000000000..48f6fed5119 --- /dev/null +++ b/test/shell.d/agent-usage-antigravity-scanner-test.sh @@ -0,0 +1,115 @@ +#!/bin/bash + +source "$(dirname "$0")/base-test.sh" + +require_command jq +require_command python3 + +TEST_HOME=$(mktemp -d) +trap 'rm -rf "$TEST_HOME"' EXIT + +AGY_DIR="$TEST_HOME/.gemini/antigravity-cli" +mkdir -p "$AGY_DIR/brain/test-conv/.system_generated/logs" "$TEST_HOME/bin" + +# Mock secret-tool to simulate authenticated state +cat >"$TEST_HOME/bin/secret-tool" <<'EOF' +#!/bin/bash +if [[ "$*" == *"lookup service gemini username antigravity"* ]]; then + echo '{"auth_method":"consumer","id_token":"header.eyJlbWFpbCI6InRlc3RAZ21haWwuY29tIn0.signature"}' + exit 0 +fi +exit 1 +EOF +chmod +x "$TEST_HOME/bin/secret-tool" + +# Write mock history and transcript +now_date=$(date +%Y-%m-%d) +timestamp="$(date -u +%Y-%m-%dT%H:%M:%SZ)" + +cat >"$AGY_DIR/history.jsonl" <"$AGY_DIR/brain/test-conv/.system_generated/logs/transcript.jsonl" <"$TEST_HOME/.config/omarchy/agents/antigravity.json" <<'EOF' +{ + "tier": "Enterprise" +} +EOF + +cfg_tier_result=$(HOME="$TEST_HOME" AGY_DIR="$AGY_DIR" XDG_CACHE_HOME="$TEST_HOME/.cache" PATH="$TEST_HOME/bin:$PATH" \ + "$ROOT/bin/omarchy-agent-usage-antigravity" --force) + +[[ $(jq -r '.tierLabel' <<<"$cfg_tier_result") == "Enterprise" ]] || + fail "Antigravity collector honors config file tier override" "$cfg_tier_result" +pass "Antigravity collector honors config file tier override" + +# Test unauthenticated state +cat >"$TEST_HOME/bin/secret-tool" <<'EOF' +#!/bin/bash +exit 1 +EOF + +unauth_result=$(HOME="$TEST_HOME" AGY_DIR="$AGY_DIR" XDG_CACHE_HOME="$TEST_HOME/.cache" PATH="$TEST_HOME/bin:$PATH" \ + "$ROOT/bin/omarchy-agent-usage-antigravity" --force) + +[[ $(jq -r '.tierLabel' <<<"$unauth_result") == "" ]] || + fail "Antigravity collector has empty tier when unauthenticated" "$unauth_result" +pass "Antigravity collector has empty tier when unauthenticated" + +[[ $(jq -r '.usageStatusText' <<<"$unauth_result") == "Sign in to see limits" ]] || + fail "Antigravity collector indicates sign in status when unauthenticated" "$unauth_result" +pass "Antigravity collector indicates sign in status when unauthenticated" From af92a06a35f867fbba87fd14a819df922865eb83 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Thu, 1 Oct 2026 21:57:36 -0400 Subject: [PATCH 04/19] Key Antigravity usage as agy and drop its made-up limits The record shares the id omarchy-default-agent and omarchy-agent use for the CLI. Limits now come only from Google's quota answer, kept from the last good check and marked stale (limitsStale, limitsFetchedAt) after a failed one, instead of being estimated from prompt counts. The plan is no longer assumed to be Pro, and transcript steps without reported usage no longer get estimated tokens. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-agy | 399 ++++++++++ bin/omarchy-agent-usage-antigravity | 737 ------------------ manual/17-ai.md | 2 +- shell/plugins/agents/README.md | 11 +- .../{antigravity-light.svg => agy-light.svg} | 0 .../assets/{antigravity.svg => agy.svg} | 0 shell/plugins/agents/manifest.json | 2 +- ...test.sh => agent-usage-agy-limits-test.sh} | 71 +- test/shell.d/agent-usage-agy-scanner-test.sh | 83 ++ .../agent-usage-antigravity-scanner-test.sh | 115 --- 10 files changed, 560 insertions(+), 860 deletions(-) create mode 100755 bin/omarchy-agent-usage-agy delete mode 100755 bin/omarchy-agent-usage-antigravity rename shell/plugins/agents/assets/{antigravity-light.svg => agy-light.svg} (100%) rename shell/plugins/agents/assets/{antigravity.svg => agy.svg} (100%) rename test/shell.d/{agent-usage-antigravity-limits-test.sh => agent-usage-agy-limits-test.sh} (55%) create mode 100755 test/shell.d/agent-usage-agy-scanner-test.sh delete mode 100755 test/shell.d/agent-usage-antigravity-scanner-test.sh diff --git a/bin/omarchy-agent-usage-agy b/bin/omarchy-agent-usage-agy new file mode 100755 index 00000000000..32432dc7b94 --- /dev/null +++ b/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() diff --git a/bin/omarchy-agent-usage-antigravity b/bin/omarchy-agent-usage-antigravity deleted file mode 100755 index 4f205e1dd76..00000000000 --- a/bin/omarchy-agent-usage-antigravity +++ /dev/null @@ -1,737 +0,0 @@ -#!/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. - -Local stats come from Antigravity conversation transcripts and history logs -under ~/.gemini/antigravity-cli/. Limits, tier, and session window calculations -come from the Google Cloud Code API (v1internal:loadCodeAssist and -v1internal:retrieveUserQuotaSummary) with local rolling calculations as fallback. -The agents panel only ever reads the JSON this prints. -""" - -from __future__ import annotations - -import argparse -import base64 -import fcntl -import json -import os -import shutil -import sqlite3 -import subprocess -import sys -import tempfile -import time -import urllib.error -import urllib.request -from datetime import datetime, timedelta, timezone -from pathlib import Path -from typing import Any - -AGENT_ID = "antigravity" -AGENT_NAME = "Antigravity" -AUTH_HELP = "Run `agy` to authenticate." - -SCAN_REUSE_SECONDS = 20 -LIMITS_ONLY_REUSE_SECONDS = 900 -PROBE_MIN_INTERVAL_SECONDS = 15 - - -def app_root() -> Path: - return Path(os.environ.get("AGY_DIR", Path.home() / ".gemini/antigravity-cli")).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 scan_cache_paths() -> tuple[Path, Path]: - root = cache_root() - return root / f"{AGENT_ID}-stats-cache.json", root / f"{AGENT_ID}-stats-cache.lock" - - -def read_fresh_json(path: Path, max_age_seconds: float) -> Any: - try: - stat = path.stat() - except OSError: - return None - age = time.time() - stat.st_mtime - if age > max_age_seconds: - return None - try: - with path.open("r", encoding="utf-8") as f: - return json.load(f) - except Exception: - return None - - -def write_json(path: Path, payload: Any) -> None: - path.parent.mkdir(parents=True, exist_ok=True) - fd, tmp_path = tempfile.mkstemp(prefix=f".{path.name}.", dir=str(path.parent)) - tmp = Path(tmp_path) - try: - with os.fdopen(fd, "w", encoding="utf-8") as f: - json.dump(payload, f, separators=(",", ":")) - f.write("\n") - tmp.chmod(0o644) - tmp.replace(path) - except BaseException: - tmp.unlink(missing_ok=True) - raise - - -def check_auth() -> tuple[bool, str, str, str, int]: - """Check authentication state via Secret Service keyring.""" - try: - raw = subprocess.check_output( - ["secret-tool", "lookup", "service", "gemini", "username", "antigravity"], - text=True, - stderr=subprocess.DEVNULL, - timeout=3, - ) - if raw.strip(): - data = json.loads(raw) - auth_method = data.get("auth_method", "consumer") - token_dict = data.get("token") or {} - access_token = str(token_dict.get("access_token") or data.get("access_token") or "") - expiry_raw = str(token_dict.get("expiry") or data.get("expiry") or "") - expires_at_ms = 0 - if expiry_raw: - try: - if expiry_raw.isdigit(): - expires_at_ms = int(expiry_raw) - if expires_at_ms < 1e12: - expires_at_ms *= 1000 - else: - dt = datetime.fromisoformat(expiry_raw.replace("Z", "+00:00")) - expires_at_ms = int(dt.timestamp() * 1000) - except Exception: - expires_at_ms = 0 - - id_token = data.get("id_token") - email = "" - if id_token: - parts = id_token.split(".") - if len(parts) >= 2: - padded = parts[1] + "=" * ((4 - len(parts[1]) % 4) % 4) - claims = json.loads(base64.urlsafe_b64decode(padded)) - email = claims.get("email", "") - return True, email, auth_method, access_token, expires_at_ms - except Exception: - pass - return False, "", "consumer", "", 0 - - -def normalize_tier_label(tier_name: str, tier_id: str = "") -> str: - tier_id_lower = tier_id.lower() - tier_name_lower = tier_name.lower() - if "ultra" in tier_id_lower or "ultra" in tier_name_lower: - return "Ultra" - if "pro" in tier_id_lower or "pro" in tier_name_lower: - return "Pro" - if "free" in tier_id_lower or "free" in tier_name_lower: - return "Free" - if tier_name.lower().startswith("google ai "): - return tier_name[10:].strip() - if tier_name: - return tier_name - if tier_id: - return tier_id.replace("-tier", "").capitalize() - return "Free" - - -def extract_tier(data: dict[str, Any], auth_method: str = "consumer") -> str: - paid_tier = data.get("paidTier") - if isinstance(paid_tier, dict): - t_name = str(paid_tier.get("name") or "").strip() - t_id = str(paid_tier.get("id") or "").strip() - if t_name or t_id: - return normalize_tier_label(t_name, t_id) - - current_tier = data.get("currentTier") - if isinstance(current_tier, dict): - c_id = str(current_tier.get("id") or "").strip() - c_name = str(current_tier.get("name") or "").strip() - if c_id or c_name: - return normalize_tier_label(c_name, c_id) - - if auth_method == "enterprise": - return "Enterprise" - if auth_method in ("cloud", "vertex"): - return "Cloud" - return "Free" - - -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 datetime.fromtimestamp(ts / 1000, timezone.utc).isoformat() - except Exception: - return raw - try: - parsed = datetime.fromisoformat(raw.replace("Z", "+00:00")) - return parsed.isoformat() - except Exception: - return raw - - -def extract_limits(data: dict[str, Any]) -> list[dict[str, Any]]: - groups = data.get("groups") - if not isinstance(groups, list): - return [] - - gemini_group = None - other_groups = [] - for g in groups: - if not isinstance(g, dict): - continue - name = str(g.get("displayName") or "").lower() - if "gemini" in name or gemini_group is None: - gemini_group = g - else: - other_groups.append(g) - - limits = [] - # Primary group: gemini models - if gemini_group: - buckets = gemini_group.get("buckets") or [] - bucket_5h = None - bucket_weekly = None - for b in buckets: - if not isinstance(b, dict): - continue - w = str(b.get("window") or "").lower() - bid = str(b.get("bucketId") or "").lower() - dname = str(b.get("displayName") or "").lower() - if w == "5h" or "5h" in bid or "five" in dname: - bucket_5h = b - elif w == "weekly" or "weekly" in bid or "week" in dname: - bucket_weekly = b - - # Enforce Session (5-hour) first, then Weekly (7-day) - if bucket_5h: - rem = float(bucket_5h.get("remainingFraction", 1.0)) - used = round(max(0.0, min(1.0, 1.0 - rem)), 4) - limits.append({ - "label": "Session (5-hour)", - "title": "Session", - "percent": used, - "resetsAt": normalize_reset_at(bucket_5h.get("resetTime")), - }) - if bucket_weekly: - rem = float(bucket_weekly.get("remainingFraction", 1.0)) - used = round(max(0.0, min(1.0, 1.0 - rem)), 4) - limits.append({ - "label": "Weekly (7-day)", - "title": "Weekly", - "percent": used, - "resetsAt": normalize_reset_at(bucket_weekly.get("resetTime")), - }) - - # Scoped limits for other groups (e.g. Claude and GPT models) - for g in other_groups: - g_name = str(g.get("displayName") or "").strip() - short_name = "Claude/GPT" if "claude" in g_name.lower() else (g_name or "3P") - for b in (g.get("buckets") or []): - if not isinstance(b, dict): - continue - rem = float(b.get("remainingFraction", 1.0)) - used = round(max(0.0, min(1.0, 1.0 - rem)), 4) - # Only show secondary model limits if there has been usage - if used > 0: - w = str(b.get("window") or "").lower() - w_title = "Weekly" if "week" in w else "Session" - title = f"{short_name} {w_title}" - limits.append({ - "label": title, - "title": title, - "percent": used, - "resetsAt": normalize_reset_at(b.get("resetTime")), - }) - - return limits - - -def probe_antigravity(access_token: str, endpoint_base: str = "") -> dict[str, Any]: - if not endpoint_base: - endpoint_base = os.environ.get("CLOUD_CODE_URL", "https://daily-cloudcode-pa.googleapis.com").rstrip("/") - - headers = { - "Authorization": f"Bearer {access_token}", - "Content-Type": "application/json", - "User-Agent": "antigravity", - } - - result: dict[str, Any] = { - "ok": False, - "tierLabel": "", - "limits": [], - "helpText": "", - "rateLimited": False, - "transport": False, - } - - tier = "" - try: - req = urllib.request.Request(f"{endpoint_base}/v1internal:loadCodeAssist", data=b"{}", headers=headers) - with urllib.request.urlopen(req, timeout=8) as resp: - payload = json.loads(resp.read().decode("utf-8", errors="replace")) - tier = extract_tier(payload) - except urllib.error.HTTPError as err: - if err.code == 429: - result["rateLimited"] = True - elif err.code == 401: - result["helpText"] = "Antigravity authentication expired. Run `agy` to authenticate." - return result - except Exception: - result["transport"] = True - result["helpText"] = "Could not reach Antigravity service." - - limits = [] - try: - req = urllib.request.Request(f"{endpoint_base}/v1internal:retrieveUserQuotaSummary", data=b"{}", headers=headers) - with urllib.request.urlopen(req, timeout=8) as resp: - payload = json.loads(resp.read().decode("utf-8", errors="replace")) - limits = extract_limits(payload) - except urllib.error.HTTPError as err: - if err.code == 429: - result["rateLimited"] = True - elif err.code == 401: - result["helpText"] = "Antigravity authentication expired. Run `agy` to authenticate." - return result - except Exception: - result["transport"] = True - - result["ok"] = bool(tier or limits) - result["tierLabel"] = tier - result["limits"] = limits - return result - - -def limit_window_open(entry: dict[str, Any], now: datetime) -> bool: - raw = str(entry.get("resetsAt") or "") - if raw == "": - return True - try: - resets_at = datetime.fromisoformat(raw.replace("Z", "+00:00")) - except Exception: - return True - if resets_at.tzinfo is None: - resets_at = resets_at.replace(tzinfo=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 = datetime.now(timezone.utc) - return [entry for entry in entries if isinstance(entry, dict) and limit_window_open(entry, now)] - - -def detect_tier_label( - root_dir: Path, - auth_ok: bool, - email: str, - auth_method: str = "consumer", - probed_tier: str = "", - cached_tier: str = "", -) -> str: - if not auth_ok: - return "" - - if os.environ.get("AGY_TIER"): - return os.environ["AGY_TIER"] - - cfg_file = Path.home() / ".config" / "omarchy" / "agents" / "antigravity.json" - if cfg_file.exists(): - try: - with cfg_file.open("r", encoding="utf-8") as f: - data = json.load(f) - if data.get("tier"): - return str(data["tier"]) - except Exception: - pass - - if probed_tier: - return probed_tier - - if cached_tier: - return cached_tier - - if auth_method == "enterprise": - return "Enterprise" - if auth_method in ("cloud", "vertex"): - return "Cloud" - - return "Pro" - - -def collect_local_stats(root_dir: Path, now: datetime) -> dict[str, Any]: - today_str = now.strftime("%Y-%m-%d") - recent_dates = [(now - timedelta(days=i)).strftime("%Y-%m-%d") for i in range(6, -1, -1)] - recent = {d: {"date": d, "messageCount": 0} for d in recent_dates} - - history_file = root_dir / "history.jsonl" - brain_dir = root_dir / "brain" - - today_prompts = 0 - total_prompts = 0 - active_dates: set[str] = set() - today_sessions: set[str] = set() - total_sessions: set[str] = set() - - if history_file.exists(): - try: - with history_file.open("r", encoding="utf-8", errors="ignore") as f: - for line in f: - line = line.strip() - if not line: - continue - try: - entry = json.loads(line) - ts = entry.get("timestamp") - if not ts: - continue - try: - if isinstance(ts, (int, float)): - dt = datetime.fromtimestamp(ts / 1000.0, timezone.utc) - else: - dt = datetime.fromisoformat(str(ts).replace("Z", "+00:00")) - except Exception: - continue - d_str = dt.strftime("%Y-%m-%d") - active_dates.add(d_str) - total_prompts += 1 - conv_id = entry.get("conversationId") or entry.get("conversation_id") - if conv_id: - total_sessions.add(conv_id) - if d_str == today_str: - today_prompts += 1 - if conv_id: - today_sessions.add(conv_id) - except Exception: - continue - except Exception: - pass - - # Fallback to conversation summaries if history.jsonl is sparse - summaries_db = root_dir / "conversation_summaries.db" - if summaries_db.exists(): - try: - conn = sqlite3.connect(f"file:{summaries_db}?mode=ro", uri=True) - cur = conn.cursor() - for row in cur.execute("SELECT conversation_id, last_modified_time FROM conversation_summaries"): - cid, lmt = row[0], str(row[1] or "") - if cid: - total_sessions.add(cid) - if lmt.startswith(today_str): - today_sessions.add(cid) - conn.close() - except Exception: - pass - - # Compute token counts across conversation transcripts - model_usage: dict[str, dict[str, int]] = {} - today_tokens = 0 - today_tokens_by_model: dict[str, int] = {} - - if brain_dir.exists(): - try: - for conv_folder in brain_dir.iterdir(): - if not conv_folder.is_dir(): - continue - tr_file = conv_folder / ".system_generated/logs/transcript.jsonl" - if not tr_file.exists(): - continue - try: - with tr_file.open("r", encoding="utf-8", errors="ignore") as f: - for line in f: - line = line.strip() - if not line: - continue - try: - step = json.loads(line) - m_name = step.get("model") or "gemini-3.8-flash-high" - if m_name not in model_usage: - model_usage[m_name] = { - "inputTokens": 0, - "outputTokens": 0, - "cacheReadInputTokens": 0, - "cacheCreationInputTokens": 0, - } - - step_usage = step.get("usage") - is_model = step.get("source") == "MODEL" or step.get("type") == "PLANNER_RESPONSE" - if isinstance(step_usage, dict): - inp = int(step_usage.get("prompt_token_count") or step_usage.get("input_tokens") or 0) - out = int(step_usage.get("candidates_token_count") or step_usage.get("output_tokens") or 0) - tok = inp + out - model_usage[m_name]["inputTokens"] += inp - model_usage[m_name]["outputTokens"] += out - else: - c_text = step.get("content") or "" - t_text = step.get("thinking") or "" - tok = max(1, (len(c_text) + len(t_text)) // 4) - if is_model: - model_usage[m_name]["outputTokens"] += tok - else: - model_usage[m_name]["inputTokens"] += tok - - created = step.get("created_at") - is_today = False - if created: - step_d = created[:10] - if step_d in recent: - recent[step_d]["messageCount"] += tok - if step_d == today_str: - is_today = True - else: - is_today = True - - if is_today: - today_tokens += tok - today_tokens_by_model[m_name] = today_tokens_by_model.get(m_name, 0) + tok - except Exception: - continue - except Exception: - continue - except Exception: - pass - - return { - "todayPrompts": today_prompts, - "todaySessions": len(today_sessions), - "todayTotalTokens": today_tokens, - "todayTokensByModel": today_tokens_by_model, - "recentDays": [recent[d] for d in recent_dates], - "totalPrompts": total_prompts, - "totalSessions": len(total_sessions), - "activeDays": len(active_dates), - "activeDates": sorted(list(active_dates)), - "modelUsage": model_usage, - } - - -def calculate_limits(root_dir: Path, now: datetime) -> list[dict[str, Any]]: - """Fallback calculation of rolling limits from local history files.""" - history_file = root_dir / "history.jsonl" - five_hours_ago_ts = int((now - timedelta(hours=5)).timestamp() * 1000) - - prompts_5h = 0 - oldest_5h_ts = None - prompts_7d = 0 - - seven_days_ago_ts = int((now - timedelta(days=7)).timestamp() * 1000) - - if history_file.exists(): - try: - with history_file.open("r", encoding="utf-8", errors="ignore") as f: - for line in f: - line = line.strip() - if not line: - continue - try: - entry = json.loads(line) - ts = entry.get("timestamp") - if not ts: - continue - if ts >= seven_days_ago_ts: - prompts_7d += 1 - if ts >= five_hours_ago_ts: - prompts_5h += 1 - if oldest_5h_ts is None or ts < oldest_5h_ts: - oldest_5h_ts = ts - except Exception: - continue - except Exception: - pass - - # Rolling 5h reset - if oldest_5h_ts: - reset_5h = datetime.fromtimestamp(oldest_5h_ts / 1000.0, timezone.utc) + timedelta(hours=5) - else: - reset_5h = now + timedelta(hours=5) - - # Check if a 429 / resource exhausted error happened recently - cli_log = root_dir / "cli.log" - rate_limited = False - if cli_log.exists(): - try: - stat = cli_log.stat() - if (time.time() - stat.st_mtime) < 18000: - with cli_log.open("r", encoding="utf-8", errors="ignore") as f: - content = f.read()[-8192:] - if "RESOURCE_EXHAUSTED" in content or "rate limit exceeded" in content.lower(): - rate_limited = True - except Exception: - pass - - session_pct = 1.0 if rate_limited else round(min(1.0, prompts_5h / 50.0), 2) - weekly_pct = round(min(1.0, prompts_7d / 250.0), 2) - - return [ - { - "label": "Session (5-hour)", - "title": "Session", - "percent": session_pct, - "resetsAt": reset_5h.strftime("%Y-%m-%dT%H:%M:%S+00:00"), - }, - { - "label": "Weekly (7-day)", - "title": "Weekly", - "percent": weekly_pct, - "resetsAt": (now + timedelta(days=7)).strftime("%Y-%m-%dT%H:%M:%S+00:00"), - }, - ] - - -def collect_limits( - root_dir: Path, - now: datetime, - auth_ok: bool, - access_token: str, - expires_at_ms: int, - auth_method: str, - force: bool, -) -> tuple[list[dict[str, Any]], str, str, str]: - """Collect limits via API probe, falling back to cached or rolling local stats.""" - limits_cache = cache_root() / f"{AGENT_ID}-limits.json" - cached = read_fresh_json(limits_cache, float("inf")) or {} - fallback = usable_cached_limits(cached) - cached_tier = str(cached.get("tierLabel") or "") - - if not auth_ok: - return [], "", "Sign in to see limits", AUTH_HELP - - # If access token is empty, we cannot probe the remote endpoint - if not access_token: - limits = fallback if fallback else calculate_limits(root_dir, now) - return limits, cached_tier, "", AUTH_HELP - - # Check reuse window for limits probe (PROBE_MIN_INTERVAL_SECONDS) - fetched_at = float(cached.get("fetchedAtMs") or 0) / 1000.0 - if fallback and not force and (time.time() - fetched_at) < PROBE_MIN_INTERVAL_SECONDS: - return fallback, cached_tier, "", AUTH_HELP - - # Perform live probe - probe = probe_antigravity(access_token) - if probe["ok"]: - probed_limits = probe["limits"] - probed_tier = probe["tierLabel"] - write_json( - limits_cache, - { - "fetchedAtMs": round(time.time() * 1000), - "limits": probed_limits, - "tierLabel": probed_tier, - }, - ) - return probed_limits, probed_tier, "", AUTH_HELP - - # Probe failed - help_text = probe["helpText"] if probe.get("helpText") else AUTH_HELP - usage_status = "" - if probe.get("rateLimited"): - usage_status = "Antigravity limits rate limited" - elif probe.get("transport"): - usage_status = "" - - if fallback: - return fallback, cached_tier, usage_status, help_text - - # Fallback to local rolling calculation - local_limits = calculate_limits(root_dir, now) - return local_limits, cached_tier, usage_status, help_text - - -def main() -> None: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--force", action="store_true", help="Bypass local caches and force a fresh scan") - parser.add_argument("--limits-only", action="store_true", help="Recompute only rate limits if cache is fresh") - args = parser.parse_args() - - root_dir = app_root() - now = datetime.now(timezone.utc) - today_str = now.strftime("%Y-%m-%d") - - auth_ok, email, auth_method, access_token, expires_at_ms = check_auth() - - cache_file, lock_file = scan_cache_paths() - stats = None - - if not args.force: - max_age = LIMITS_ONLY_REUSE_SECONDS if args.limits_only else SCAN_REUSE_SECONDS - cached = read_fresh_json(cache_file, max_age) - if isinstance(cached, dict) and cached.get("scanDate") == today_str: - stats = cached.get("stats") - - if stats is None: - try: - with lock_file.open("w") as lock: - fcntl.flock(lock, fcntl.LOCK_EX) - stats = collect_local_stats(root_dir, now) - write_json(cache_file, {"schemaVersion": 1, "scanDate": today_str, "stats": stats}) - except Exception as exc: - print(f"omarchy-agent-usage-antigravity: cache lock failed ({exc}), scanning directly", file=sys.stderr) - stats = collect_local_stats(root_dir, now) - - limits, probed_tier, usage_status_text, auth_help_text = collect_limits( - root_dir, - now, - auth_ok, - access_token, - expires_at_ms, - auth_method, - args.force, - ) - - cached_limits_file = cache_root() / f"{AGENT_ID}-limits.json" - cached_payload = read_fresh_json(cached_limits_file, float("inf")) or {} - cached_tier = str(cached_payload.get("tierLabel") or "") - tier_label = detect_tier_label(root_dir, auth_ok, email, auth_method, probed_tier, cached_tier) - - total_prompts = stats.get("totalPrompts", 0) - ready = auth_ok or total_prompts > 0 or len(limits) > 0 - has_local_stats = total_prompts > 0 - - record = { - "schemaVersion": 1, - "id": AGENT_ID, - "name": AGENT_NAME, - "updatedAt": now.strftime("%Y-%m-%dT%H:%M:%S.%f+00:00"), - "ready": ready, - "hasLocalStats": has_local_stats, - "todayPrompts": stats.get("todayPrompts", 0), - "todaySessions": stats.get("todaySessions", 0), - "todayTotalTokens": stats.get("todayTotalTokens", 0), - "todayTokensByModel": stats.get("todayTokensByModel", {}), - "recentDays": stats.get("recentDays", []), - "totalPrompts": total_prompts, - "totalSessions": stats.get("totalSessions", 0), - "activeDays": stats.get("activeDays", 0), - "activeDates": stats.get("activeDates", []), - "modelUsage": stats.get("modelUsage", {}), - "limits": limits, - "tierLabel": tier_label, - "usageStatusText": usage_status_text, - "authHelpText": auth_help_text, - } - - print(json.dumps(record, separators=(",", ":"))) - - -if __name__ == "__main__": - main() diff --git a/manual/17-ai.md b/manual/17-ai.md index f9df5ba956a..fb2f6882256 100644 --- a/manual/17-ai.md +++ b/manual/17-ai.md @@ -38,7 +38,7 @@ There are terminal shortcuts too: `a` runs the default agent inline in the curre ### The agents panel -The top bar has an agents icon. Before you've set up an agent, its panel walks you through signing in to Claude, Codex, or Grok, and the first one becomes your default agent. After that, the panel puts every subscription on one page: each agent and account with its plan and the percentage used of its 5-hour session and weekly limits (or the remaining prepaid balance), under a line that rotates through what you've used across all of them, like tokens this week and your most used model. Under _Make something_ it offers starter prompts for a new theme, plugin, or app that hand the task to your default agent. Apps are built the way Omarchy's own Hype, Monologue, and Omacut are, in C++ and Qt Quick, following the omarchy-app agent skill, and the compiler and Qt pieces they need come installed. Claude Code, Codex, Grok, and Fireworks are covered out of the box, and Hermes shows its sessions and tokens, from every profile, once you've used it. +The top bar has an agents icon. Before you've set up an agent, its panel walks you through signing in to Claude, Codex, or Grok, and the first one becomes your default agent. After that, the panel puts every subscription on one page: each agent and account with its plan and the percentage used of its 5-hour session and weekly limits (or the remaining prepaid balance), under a line that rotates through what you've used across all of them, like tokens this week and your most used model. Under _Make something_ it offers starter prompts for a new theme, plugin, or app that hand the task to your default agent. Apps are built the way Omarchy's own Hype, Monologue, and Omacut are, in C++ and Qt Quick, following the omarchy-app agent skill, and the compiler and Qt pieces they need come installed. Claude Code, Codex, Grok, and Fireworks are covered out of the box, as is Antigravity with its 5-hour and weekly limits, and Hermes shows its sessions and tokens, from every profile, once you've used it. Left-click the bar icon for the panel, right-click to launch your default agent. The usage records behind it are regenerated every 15 minutes by `omarchy agent usage-update`, and the panel can even merge usage from your other machines via a synced folder. See the README under `$OMARCHY_PATH/shell/plugins/agents/` for the full settings. diff --git a/shell/plugins/agents/README.md b/shell/plugins/agents/README.md index dccd2f6cdcc..34cb0f39344 100644 --- a/shell/plugins/agents/README.md +++ b/shell/plugins/agents/README.md @@ -75,7 +75,7 @@ light surfaces — and the bar glyph stands in when there is none. | `grok` | The credits endpoint behind Grok's `/usage` view (the billing period's included usage) | Each session's `usage.json` (the ledger `grok usage` prints: tokens by model per finished turn), plus `summary.json` for sessions | | `fireworks` | Estimated prepaid balance: configured funding minus rated account costs | Fireworks billing API, grouped by day and model for the last 30 days | | `hermes` | None (Hermes has no subscription of its own) | `state.db` in Hermes' home and in each named profile under `~/.hermes/profiles`, read-only | -| `antigravity` | Google Cloud Code API (`loadCodeAssist` and `retrieveUserQuotaSummary`) with rolling session/weekly calculation fallback | Transcripts and history logs under `~/.gemini/antigravity-cli/` (`AGY_DIR`) | +| `agy` | Google's Cloud Code API (`loadCodeAssist` for the plan, `retrieveUserQuotaSummary` for the 5-hour and weekly windows) | `history.jsonl`, `conversation_summaries.db`, and each conversation's `transcript.jsonl` under `~/.gemini/antigravity-cli/` (`AGY_DIR`) | When `~/.local/state/omarchy/agents/accounts/.json` registers more than one account, the `claude`, `codex`, and `grok` records @@ -91,12 +91,13 @@ since every account shares the primary home's history. After each run, active account crosses its threshold, and re-collects the record if the active account changed. -Antigravity checks authentication via the Secret Service keyring -(`gemini`/`antigravity`), queries Google's Cloud Code endpoints for account -tier and quota limits, and honors a custom app root via `AGY_DIR`. Claude limits need a signed-in CLI; without credentials the panel says so and +Claude limits need a signed-in CLI; without credentials the panel says so and falls back to local stats only. A non-default Claude directory is honored via `CLAUDE_CONFIG_DIR`, Codex via `CODEX_HOME`, Grok via `GROK_HOME`, Hermes via -`HERMES_HOME`. Grok's +`HERMES_HOME`, Antigravity via `AGY_DIR`. Antigravity's plan and limits are +asked with the sign-in `agy` keeps in the Secret Service keyring (service +`gemini`, username `antigravity`); the record's id is `agy`, the name the +default agent knows it by. Grok's plan comes from the settings it caches in its home, and its limit from the credits endpoint its own `/usage` view reads, asked with each account's sign-in; a sign-in left to lapse shows the last credits until Grok runs diff --git a/shell/plugins/agents/assets/antigravity-light.svg b/shell/plugins/agents/assets/agy-light.svg similarity index 100% rename from shell/plugins/agents/assets/antigravity-light.svg rename to shell/plugins/agents/assets/agy-light.svg diff --git a/shell/plugins/agents/assets/antigravity.svg b/shell/plugins/agents/assets/agy.svg similarity index 100% rename from shell/plugins/agents/assets/antigravity.svg rename to shell/plugins/agents/assets/agy.svg diff --git a/shell/plugins/agents/manifest.json b/shell/plugins/agents/manifest.json index 6a90aea67ad..a62fe2564cb 100644 --- a/shell/plugins/agents/manifest.json +++ b/shell/plugins/agents/manifest.json @@ -24,7 +24,7 @@ "grok": { "enabled": true }, "fireworks": { "enabled": true }, "hermes": { "enabled": true }, - "antigravity": { "enabled": true } + "agy": { "enabled": true } }, "refreshIntervalSec": 900, "syncMode": "Off", diff --git a/test/shell.d/agent-usage-antigravity-limits-test.sh b/test/shell.d/agent-usage-agy-limits-test.sh similarity index 55% rename from test/shell.d/agent-usage-antigravity-limits-test.sh rename to test/shell.d/agent-usage-agy-limits-test.sh index 89c0dbdc945..1198a4356e8 100755 --- a/test/shell.d/agent-usage-antigravity-limits-test.sh +++ b/test/shell.d/agent-usage-agy-limits-test.sh @@ -5,7 +5,7 @@ source "$(dirname "$0")/base-test.sh" require_command jq require_command python3 -COLLECTOR="$ROOT/bin/omarchy-agent-usage-antigravity" +COLLECTOR="$ROOT/bin/omarchy-agent-usage-agy" # probe_antigravity reaches Google's Cloud Code endpoints, so the readers that # interpret the answers are exercised with recorded payloads standing in. @@ -134,3 +134,72 @@ pass "Antigravity collector calculates 3P used percentage" [[ $(jq -r 'length' <<<"$limits") == "3" ]] || fail "Antigravity collector omits unused 3P session limit" "$limits" pass "Antigravity collector omits unused 3P session limit" + +reordered=$(test_limits '{"groups": [ + {"displayName": "Claude and GPT models", "buckets": [{"window": "5h", "remainingFraction": 1.0}]}, + {"displayName": "Gemini Models", "buckets": [{"window": "weekly", "remainingFraction": 0.5}]} +]}') +[[ $(jq -c 'map({label, percent})' <<<"$reordered") == '[{"label":"Weekly (7-day)","percent":0.5}]' ]] || + fail "Antigravity collector finds the Gemini group wherever it is listed" "$reordered" +pass "Antigravity collector finds the Gemini group wherever it is listed" + +# The whole record with Google's endpoints and the keyring standing in: a +# good answer is cached, and a failed check after it keeps those limits, +# marked stale, rather than making any up. +test_tmp=$(mktemp -d) +trap 'rm -rf "$test_tmp"' EXIT +mkdir -p "$test_tmp/bin" +cat >"$test_tmp/bin/secret-tool" <<'EOF' +#!/bin/bash +echo '{"auth_method":"consumer","token":{"access_token":"token-main","expiry":"2999-01-01T00:00:00Z"}}' +EOF +chmod +x "$test_tmp/bin/secret-tool" +reset_at=$(python3 -c 'import datetime as dt; print((dt.datetime.now(dt.timezone.utc) + dt.timedelta(hours=3)).isoformat())') + +collect() { + HOME="$test_tmp" XDG_CACHE_HOME="$test_tmp/cache" PATH="$test_tmp/bin:$PATH" COLLECTOR="$COLLECTOR" RESET_AT="$reset_at" python3 - "$@" <<'PY' +import importlib.machinery, importlib.util, io, json, os, sys, urllib.error + +loader = importlib.machinery.SourceFileLoader("collector", os.environ["COLLECTOR"]) +spec = importlib.util.spec_from_loader(loader.name, loader) +collector = importlib.util.module_from_spec(spec) +loader.exec_module(collector) + +answers = { + "loadCodeAssist": {"paidTier": {"id": "g1-pro-tier", "name": "Google AI Pro"}}, + "retrieveUserQuotaSummary": {"groups": [{"displayName": "Gemini Models", "buckets": [ + {"window": "5h", "remainingFraction": 0.75, "resetTime": os.environ["RESET_AT"]} + ]}]}, +} +state = sys.argv[1] + +def urlopen(request, timeout=None): + if state == "down": + raise urllib.error.URLError("offline") + assert request.get_header("Authorization") == "Bearer token-main" + return io.BytesIO(json.dumps(answers[request.full_url.rsplit(":", 1)[1]]).encode()) + +collector.urllib.request.urlopen = urlopen +sys.argv = ["omarchy-agent-usage-agy", "--force"] +collector.main() +PY +} + +record=$(collect up) +[[ $(jq -c '{id, ready, tierLabel, stale: .limitsStale, percent: .limits[0].percent}' <<<"$record") == '{"id":"agy","ready":true,"tierLabel":"Pro","stale":false,"percent":0.25}' ]] || + fail "Antigravity collector reports its plan and live limits" "$record" +(( $(jq '.limitsFetchedAt' <<<"$record") > 0 )) || fail "Antigravity collector says when the limits were fetched" "$record" +pass "Antigravity collector reports its plan and live limits" + +fetched=$(jq '.limitsFetchedAt' <<<"$record") +record=$(collect down) +[[ $(jq -c '{tierLabel, stale: .limitsStale, percent: .limits[0].percent}' <<<"$record") == '{"tierLabel":"Pro","stale":true,"percent":0.25}' ]] || + fail "Antigravity collector keeps the last limits, marked stale, after a failed check" "$record" +[[ $(jq '.limitsFetchedAt' <<<"$record") == "$fetched" ]] || fail "stale Antigravity limits keep when they were fetched" "$record" +pass "Antigravity collector keeps the last limits, marked stale, after a failed check" + +rm -rf "$test_tmp/cache" +record=$(collect down) +[[ $(jq -c '{limits, usageStatusText}' <<<"$record") == '{"limits":[],"usageStatusText":"Antigravity limits unavailable"}' ]] || + fail "Antigravity collector invents no limits when it has never reached Google" "$record" +pass "Antigravity collector invents no limits when it has never reached Google" diff --git a/test/shell.d/agent-usage-agy-scanner-test.sh b/test/shell.d/agent-usage-agy-scanner-test.sh new file mode 100755 index 00000000000..a7e6d7539d5 --- /dev/null +++ b/test/shell.d/agent-usage-agy-scanner-test.sh @@ -0,0 +1,83 @@ +#!/bin/bash + +source "$(dirname "$0")/base-test.sh" + +require_command jq +require_command python3 + +TEST_HOME=$(mktemp -d) +trap 'rm -rf "$TEST_HOME"' EXIT + +AGY_DIR="$TEST_HOME/.gemini/antigravity-cli" +mkdir -p "$AGY_DIR/brain/test-conv/.system_generated/logs" "$TEST_HOME/bin" + +# Mock secret-tool to simulate a signed-in keyring entry without a token, so +# nothing is asked of Google. +cat >"$TEST_HOME/bin/secret-tool" <<'EOF' +#!/bin/bash +if [[ "$*" == *"lookup service gemini username antigravity"* ]]; then + echo '{"auth_method":"consumer","id_token":"header.eyJlbWFpbCI6InRlc3RAZ21haWwuY29tIn0.signature"}' + exit 0 +fi +exit 1 +EOF +chmod +x "$TEST_HOME/bin/secret-tool" + +# Write mock history and transcript +timestamp="$(date -u +%Y-%m-%dT%H:%M:%SZ)" + +cat >"$AGY_DIR/history.jsonl" <"$AGY_DIR/brain/test-conv/.system_generated/logs/transcript.jsonl" <"$TEST_HOME/bin/secret-tool" <<'EOF' +#!/bin/bash +exit 1 +EOF +rm -rf "$AGY_DIR" +result=$(collect) + +[[ $(jq -c '{ready, tierLabel, limits, totalPrompts, totalSessions, activeDays, modelUsage}' <<<"$result") == '{"ready":false,"tierLabel":"","limits":[],"totalPrompts":0,"totalSessions":0,"activeDays":0,"modelUsage":{}}' ]] || + fail "Antigravity collector prints an empty record without Antigravity" "$result" +pass "Antigravity collector prints an empty record without Antigravity" diff --git a/test/shell.d/agent-usage-antigravity-scanner-test.sh b/test/shell.d/agent-usage-antigravity-scanner-test.sh deleted file mode 100755 index 48f6fed5119..00000000000 --- a/test/shell.d/agent-usage-antigravity-scanner-test.sh +++ /dev/null @@ -1,115 +0,0 @@ -#!/bin/bash - -source "$(dirname "$0")/base-test.sh" - -require_command jq -require_command python3 - -TEST_HOME=$(mktemp -d) -trap 'rm -rf "$TEST_HOME"' EXIT - -AGY_DIR="$TEST_HOME/.gemini/antigravity-cli" -mkdir -p "$AGY_DIR/brain/test-conv/.system_generated/logs" "$TEST_HOME/bin" - -# Mock secret-tool to simulate authenticated state -cat >"$TEST_HOME/bin/secret-tool" <<'EOF' -#!/bin/bash -if [[ "$*" == *"lookup service gemini username antigravity"* ]]; then - echo '{"auth_method":"consumer","id_token":"header.eyJlbWFpbCI6InRlc3RAZ21haWwuY29tIn0.signature"}' - exit 0 -fi -exit 1 -EOF -chmod +x "$TEST_HOME/bin/secret-tool" - -# Write mock history and transcript -now_date=$(date +%Y-%m-%d) -timestamp="$(date -u +%Y-%m-%dT%H:%M:%SZ)" - -cat >"$AGY_DIR/history.jsonl" <"$AGY_DIR/brain/test-conv/.system_generated/logs/transcript.jsonl" <"$TEST_HOME/.config/omarchy/agents/antigravity.json" <<'EOF' -{ - "tier": "Enterprise" -} -EOF - -cfg_tier_result=$(HOME="$TEST_HOME" AGY_DIR="$AGY_DIR" XDG_CACHE_HOME="$TEST_HOME/.cache" PATH="$TEST_HOME/bin:$PATH" \ - "$ROOT/bin/omarchy-agent-usage-antigravity" --force) - -[[ $(jq -r '.tierLabel' <<<"$cfg_tier_result") == "Enterprise" ]] || - fail "Antigravity collector honors config file tier override" "$cfg_tier_result" -pass "Antigravity collector honors config file tier override" - -# Test unauthenticated state -cat >"$TEST_HOME/bin/secret-tool" <<'EOF' -#!/bin/bash -exit 1 -EOF - -unauth_result=$(HOME="$TEST_HOME" AGY_DIR="$AGY_DIR" XDG_CACHE_HOME="$TEST_HOME/.cache" PATH="$TEST_HOME/bin:$PATH" \ - "$ROOT/bin/omarchy-agent-usage-antigravity" --force) - -[[ $(jq -r '.tierLabel' <<<"$unauth_result") == "" ]] || - fail "Antigravity collector has empty tier when unauthenticated" "$unauth_result" -pass "Antigravity collector has empty tier when unauthenticated" - -[[ $(jq -r '.usageStatusText' <<<"$unauth_result") == "Sign in to see limits" ]] || - fail "Antigravity collector indicates sign in status when unauthenticated" "$unauth_result" -pass "Antigravity collector indicates sign in status when unauthenticated" From 60e940dced7a334fce10d6b2c82ed1e27afd306a Mon Sep 17 00:00:00 2001 From: FelipeMayerDev Date: Thu, 1 Oct 2026 21:58:58 -0400 Subject: [PATCH 05/19] Add Z.ai GLM Coding Plan agent usage collector Co-authored-by: Felipe Mayer --- bin/omarchy-agent-usage-zai | 667 +++++++++++++++++++ shell/plugins/agents/README.md | 7 + shell/plugins/agents/assets/zai-light.svg | 1 + shell/plugins/agents/assets/zai.svg | 1 + shell/plugins/agents/manifest.json | 3 +- test/shell.d/agent-usage-zai-limits-test.sh | 151 +++++ test/shell.d/agent-usage-zai-scanner-test.sh | 90 +++ 7 files changed, 919 insertions(+), 1 deletion(-) create mode 100755 bin/omarchy-agent-usage-zai create mode 100644 shell/plugins/agents/assets/zai-light.svg create mode 100644 shell/plugins/agents/assets/zai.svg create mode 100755 test/shell.d/agent-usage-zai-limits-test.sh create mode 100755 test/shell.d/agent-usage-zai-scanner-test.sh diff --git a/bin/omarchy-agent-usage-zai b/bin/omarchy-agent-usage-zai new file mode 100755 index 00000000000..ed684aef791 --- /dev/null +++ b/bin/omarchy-agent-usage-zai @@ -0,0 +1,667 @@ +#!/usr/bin/python3 +# omarchy:summary=Print the Z.ai GLM Coding Plan usage record as JSON +# omarchy:args=[--force] [--limits-only] +# omarchy:hidden=true +"""Collect Z.ai (GLM Coding Plan) usage into one display-ready JSON record. + +Rate limits and the plan tier come from Z.ai's coding-plan monitor endpoint +(…/api/monitor/usage/quota/limit), which reports the rolling 5-hour and weekly +token windows. That endpoint only answers for keys that carry a Coding Plan; a +generic pay-as-you-go key gets "no coding plan" and the record falls back to +local stats alone, exactly like the claude collector without a signed-in CLI. + +Local stats come from pi/omp sessions that ran through the Z.ai provider (the +`glm-*` models). Both the local scan and the network probe are cached: the +panel fires `--limits-only` on every open, so a repeated open must reuse a +recent scan and must not re-probe Z.ai more than once every few seconds. The +agents panel only ever reads the JSON this prints. +""" + +from __future__ import annotations + +import argparse +import contextlib +import fcntl +import hashlib +import json +import os +import stat +import sys +import tempfile +import time +import urllib.error +import urllib.request +from datetime import datetime, timedelta, timezone +from pathlib import Path +from typing import Any, Iterator + +AGENT_ID = "zai" +AGENT_NAME = "Z.ai" +AUTH_HELP_NO_KEY = ( + "Set ZAI_API_KEY (or add ~/.config/omarchy/agents/zai.json) to show Coding Plan limits." +) +AUTH_HELP_BAD_KEY = "Z.ai rejected the saved API key; check ~/.config/omarchy/agents/zai.json." + +# Global (Z.ai) vs China (Zhipu / BigModel) split the same API under two hosts. +PLATFORMS = { + "zai": "https://api.z.ai", + "zhipu": "https://open.bigmodel.cn", +} +QUOTA_PATH = "/api/monitor/usage/quota/limit" + +# Coding-plan limit windows are tagged by (unit, number); the monitor endpoint +# reports one CREDIT_LIMIT row per window (older responses used TOKENS_LIMIT). +LIMIT_TYPES = {"CREDIT_LIMIT", "TOKENS_LIMIT"} +LIMIT_WINDOWS = { + (3, 5): "Session (5-hour)", + (6, 1): "Weekly (7-day)", +} + +# A scan this recent is only reused to dedup concurrent collector runs (the +# update command backgrounds one per agent while the panel refreshes on its +# own); every periodic widget refresh lands a real rescan. --limits-only +# promises only fresh limits, so it may reuse a scan for up to 15 minutes. +SCAN_REUSE_SECONDS = 20 +LIMITS_ONLY_REUSE_SECONDS = 900 +# The panel opens and shuts fire a probe each; a recent probe result is reused +# for this window rather than turning every flick into a request. +PROBE_MIN_INTERVAL_SECONDS = 15 + +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)] + + +def number(value: Any) -> int: + try: + return max(0, int(value or 0)) + except (TypeError, ValueError): + return 0 + + +def local_day(value: Any) -> str: + if value is None: + return datetime.now().strftime("%Y-%m-%d") + if isinstance(value, (int, float)): + # pi/omp message timestamps are milliseconds. + 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 ValueError: + return datetime.now().strftime("%Y-%m-%d") + + +def empty_bucket() -> dict[str, int]: + return { + "inputTokens": 0, + "outputTokens": 0, + "cacheReadInputTokens": 0, + "cacheCreationInputTokens": 0, + } + + +def empty_stats() -> dict[str, Any]: + return { + "todayPrompts": 0, + "todaySessions": 0, + "todayTotalTokens": 0, + "todayTokensByModel": {}, + "recentDays": [], + "totalPrompts": 0, + "totalSessions": 0, + "activeDays": 0, + "activeDates": [], + "modelUsage": {}, + } + + +def base_record(**overrides: Any) -> dict[str, Any]: + record: dict[str, Any] = { + "schemaVersion": 1, + "id": AGENT_ID, + "name": AGENT_NAME, + "updatedAt": datetime.now(timezone.utc).isoformat(), + "ready": False, + "hasLocalStats": False, + "hasPromptStats": True, + "tierLabel": "", + "usageStatusText": "", + "authHelpText": "", + "limits": [], + } + record.update(empty_stats()) + record.update(overrides) + return record + + +# -------------------------------------------------------------------------- +# Cache primitives (shared by the local-stats scan and the limits probe) +# -------------------------------------------------------------------------- + +def cache_root() -> Path: + root = Path(os.environ.get("XDG_CACHE_HOME") or (Path.home() / ".cache")) / "omarchy" / "agent-usage" + root.mkdir(parents=True, exist_ok=True) + return root + + +def read_fresh_json(path: Path, max_age_seconds: float) -> Any: + if max_age_seconds <= 0 or not path.exists(): + return None + try: + # A negative age means the mtime is in the future: the clock moved + # backwards since the write, so the cache's freshness cannot be trusted. + age = time.time() - path.stat().st_mtime + if 0 <= age <= max_age_seconds: + return json.loads(path.read_text(encoding="utf-8")) + except Exception: + return None + return None + + +def write_json(path: Path, payload: Any) -> None: + # A temp name unique to this writer, not derived from the target: several + # collectors can run at once, 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=(",", ":")) + "\n") + tmp.chmod(0o644) + tmp.replace(path) + except BaseException: + tmp.unlink(missing_ok=True) + raise + + +@contextlib.contextmanager +def exclusive_lock(path: Path) -> Iterator[None]: + """Hold flock on `path` for the block, or run the block unlocked.""" + fd: int | None = None + try: + flags = os.O_WRONLY | os.O_CREAT | os.O_NOFOLLOW | os.O_NONBLOCK | os.O_CLOEXEC + try: + fd = os.open(path, flags | os.O_EXCL, 0o600) + except FileExistsError: + fd = os.open(path, flags) + info = os.fstat(fd) + if ( + not stat.S_ISREG(info.st_mode) + or info.st_nlink == 0 + or info.st_uid != os.geteuid() + or info.st_mode & 0o022 + ): + os.close(fd) + fd = None + except OSError: + if fd is not None: + os.close(fd) + fd = None + if fd is not None: + try: + fcntl.flock(fd, fcntl.LOCK_EX) + except OSError: + os.close(fd) + fd = None + try: + yield + finally: + if fd is not None: + os.close(fd) + + +# -------------------------------------------------------------------------- +# Credentials +# -------------------------------------------------------------------------- + +def read_env_file(path: Path) -> dict[str, str]: + values: dict[str, str] = {} + try: + for line in path.read_text().splitlines(): + line = line.strip() + if line.startswith("export "): + line = line[len("export "):].strip() + if not line or line.startswith("#") or "=" not in line: + continue + key, _, val = line.partition("=") + values[key.strip()] = val.strip().strip("'\"") + except OSError: + return {} + return values + + +def read_config() -> dict[str, Any]: + config_home = Path(os.environ.get("XDG_CONFIG_HOME") or (Path.home() / ".config")) + try: + parsed = json.loads((config_home / "omarchy" / "agents" / "zai.json").read_text()) + return parsed if isinstance(parsed, dict) else {} + except (OSError, json.JSONDecodeError): + return {} + + +def credentials() -> tuple[str, str]: + """Return (api_key, platform). Platform is a PLATFORMS key.""" + config = read_config() + + # An explicit config file wins, then live env, then whatever omp/pi stored. + key = str(config.get("apiKey") or "").strip() + platform = str(config.get("platform") or "").strip().lower() + + if not key: + if os.environ.get("ZAI_API_KEY"): + key, platform = os.environ["ZAI_API_KEY"].strip(), platform or "zai" + elif os.environ.get("ZHIPU_API_KEY") or os.environ.get("ZHIPUAI_API_KEY"): + key = (os.environ.get("ZHIPU_API_KEY") or os.environ.get("ZHIPUAI_API_KEY")).strip() + platform = platform or "zhipu" + + if not key: + for env_path in ( + Path.home() / ".omp" / "agent" / ".env", + Path.home() / ".pi" / "agent" / ".env", + ): + values = read_env_file(env_path) + if values.get("ZAI_API_KEY"): + key, platform = values["ZAI_API_KEY"], platform or "zai" + break + if values.get("ZHIPU_API_KEY") or values.get("ZHIPUAI_API_KEY"): + key = values.get("ZHIPU_API_KEY") or values["ZHIPUAI_API_KEY"] + platform = platform or "zhipu" + break + + if platform not in PLATFORMS: + platform = "zai" + return key, platform + + +# -------------------------------------------------------------------------- +# Coding-plan limits +# -------------------------------------------------------------------------- + +def iso_from_ms(value: Any) -> str: + try: + ts = float(value) + except (TypeError, ValueError): + return "" + if ts <= 0: + return "" + # The field is milliseconds, but a plain-seconds epoch (< year 33658) is + # accepted too rather than rendered as a date 1000x in the future. + if ts < 1e12: + ts *= 1000 + try: + return datetime.fromtimestamp(ts / 1000, timezone.utc).isoformat() + except (OverflowError, OSError, ValueError): + return "" + + +def fetch_limits(api_key: str, platform: str) -> dict[str, Any]: + """Probe the monitor endpoint. + + On success returns {"limits": [...], "tierLabel": str}. On a miss returns + {"status": str} plus, when applicable, "transport" (reached no server, so + worth an early retry) or "keyRejected" (the key itself is bad). The endpoint + answers HTTP 200 even when the key has no plan, carrying + {"success": false, "msg": "…coding plan…"} instead of a data payload. + """ + url = PLATFORMS[platform] + QUOTA_PATH + request = urllib.request.Request( + url, + method="GET", + headers={ + # This endpoint takes the raw token, not a Bearer prefix. + "Authorization": api_key, + "Accept-Language": "en-US,en", + "Content-Type": "application/json", + }, + ) + try: + with urllib.request.urlopen(request, timeout=15) as response: + body = json.loads(response.read(1_048_576)) + except urllib.error.HTTPError as error: + if error.code in (401, 403): + return {"status": "Z.ai rejected the API key", "keyRejected": True} + return {"status": f"Z.ai monitor returned HTTP {error.code}"} + except (urllib.error.URLError, TimeoutError): + return {"status": "Could not reach the Z.ai monitor endpoint", "transport": True} + except json.JSONDecodeError: + return {"status": "Z.ai returned an invalid monitor response"} + + data = body.get("data") if isinstance(body, dict) else None + if not isinstance(data, dict) or not isinstance(data.get("limits"), list): + msg = str(body.get("msg") or "").strip()[:200] if isinstance(body, dict) else "" + status = "No GLM Coding Plan on this key" if "coding plan" in msg.lower() else ( + msg or "Coding Plan limits unavailable" + ) + return {"status": status} + + limits = [] + for item in data["limits"]: + if not isinstance(item, dict) or item.get("type") not in LIMIT_TYPES: + continue + label = LIMIT_WINDOWS.get((item.get("unit"), item.get("number"))) + if not label: + # A window Z.ai adds later would vanish silently otherwise. + print( + f"omarchy-agent-usage-zai: unmapped limit window (unit={item.get('unit')}, " + f"number={item.get('number')})", + file=sys.stderr, + ) + continue + try: + percent = max(0.0, min(1.0, float(item.get("percentage")) / 100.0)) + except (TypeError, ValueError): + # A missing/unparseable percentage is skipped: a false 0% meter reads + # as a fresh allowance, which is worse than showing one fewer window. + continue + limits.append({ + "label": label, + "percent": percent, + "resetsAt": iso_from_ms(item.get("nextResetTime")), + }) + + tier = str(data.get("level") or "").strip()[:64] + return {"limits": limits, "tierLabel": tier[:1].upper() + tier[1:].lower() if tier else ""} + + +# 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. 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_utc: datetime) -> bool: + raw = str(entry.get("resetsAt") or "") + if raw == "": + return True + try: + resets_at = datetime.fromisoformat(raw.replace("Z", "+00:00")) + except ValueError: + return True + if resets_at.tzinfo is None: + resets_at = resets_at.replace(tzinfo=timezone.utc) + return resets_at > now_utc + + +def usable_cached_limits(cached: dict[str, Any]) -> list[dict[str, Any]]: + entries = cached.get("limits") + if not isinstance(entries, list): + return [] + now_utc = datetime.now(timezone.utc) + return [entry for entry in entries if isinstance(entry, dict) and limit_window_open(entry, now_utc)] + + +def collect_limits(api_key: str, platform: str, force: bool) -> dict[str, Any]: + """Fresh limits when a probe is due, else the last known ones. + + Reuses a recent probe for PROBE_MIN_INTERVAL_SECONDS so repeated panel opens + do not spam the endpoint, and keeps the last known limits as the answer of + record when a later probe fails but its window has not yet reset. + """ + result = {"limits": [], "tierLabel": "", "usageStatusText": "", "authHelpText": "", "retryAdvised": False} + + probe_cache = cache_root() / "zai-limits.json" + cached = read_fresh_json(probe_cache, float("inf")) or {} + fallback = usable_cached_limits(cached) if isinstance(cached, dict) else [] + cached_tier = str(cached.get("tierLabel") or "") if isinstance(cached, dict) else "" + + fetched_at = number(cached.get("fetchedAtMs")) / 1000 if isinstance(cached, dict) else 0 + if fallback and not force and time.time() - fetched_at < PROBE_MIN_INTERVAL_SECONDS: + result["limits"], result["tierLabel"] = fallback, cached_tier + return result + + with exclusive_lock(scan_cache_paths()[2]): + if not force: + cached = read_fresh_json(probe_cache, float("inf")) or {} + fetched_at = number(cached.get("fetchedAtMs")) / 1000 if isinstance(cached, dict) else 0 + fresh = usable_cached_limits(cached) if isinstance(cached, dict) else [] + if fresh and time.time() - fetched_at < PROBE_MIN_INTERVAL_SECONDS: + result["limits"], result["tierLabel"] = fresh, str(cached.get("tierLabel") or "") + return result + + probe = fetch_limits(api_key, platform) + if probe.get("limits"): + result["limits"] = probe["limits"] + result["tierLabel"] = probe.get("tierLabel", "") + try: + write_json(probe_cache, { + "fetchedAtMs": round(time.time() * 1000), + "limits": probe["limits"], + "tierLabel": result["tierLabel"], + }) + except Exception as exc: + print(f"omarchy-agent-usage-zai: could not write limits cache ({exc})", file=sys.stderr) + return result + + # No fresh limits. A transport failure reached no server — typically the + # seconds after login before the network is up — so ask the shell to retry + # sooner than its regular interval. + if probe.get("transport"): + result["retryAdvised"] = True + if fallback: + result["limits"], result["tierLabel"] = fallback, cached_tier + result["usageStatusText"] = probe.get("status", "") + # Only blame the key when the key is the problem; a working key with no plan, + # or an endpoint hiccup, must not tell the user to go set a key they have. + if probe.get("keyRejected"): + result["authHelpText"] = AUTH_HELP_BAD_KEY + return result + + +# -------------------------------------------------------------------------- +# Local session stats +# -------------------------------------------------------------------------- + +def scan_sessions() -> dict[str, Any]: + recent = {day: {"date": day, "messageCount": 0} for day in recent_dates} + + model_usage: dict[str, dict[str, int]] = {} + today_by_model: dict[str, int] = {} + active_dates: set[str] = set() + total_sessions: set[str] = set() + today_sessions: set[str] = set() + seen: set[str] = set() + totals = {"today_prompts": 0, "today_tokens": 0, "total_prompts": 0} + + def add_usage(day, session_key, model, inp, out, cread, cwrite): + total = inp + out + cread + cwrite + totals["total_prompts"] += 1 + total_sessions.add(session_key) + active_dates.add(day) + bucket = model_usage.setdefault(model, empty_bucket()) + bucket["inputTokens"] += inp + bucket["outputTokens"] += out + bucket["cacheReadInputTokens"] += cread + bucket["cacheCreationInputTokens"] += cwrite + if day in recent: + recent[day]["messageCount"] += total + if day == today: + totals["today_prompts"] += 1 + today_sessions.add(session_key) + totals["today_tokens"] += total + today_by_model[model] = today_by_model.get(model, 0) + total + + roots = [ + Path.home() / ".pi" / "agent" / "sessions", + Path.home() / ".omp" / "agent" / "sessions", + ] + for root in roots: + if not root.exists(): + continue + for path in root.rglob("*.jsonl"): + try: + handle = path.open(errors="replace") + except OSError: + continue + with handle: + for line in handle: + # Cheap substring gate before the JSON parse: an assistant turn on a + # Z.ai provider always carries both markers. + if '"usage"' not in line or ('"zai"' not in line and '"zhipu"' not in line): + continue + try: + entry = json.loads(line) + except json.JSONDecodeError: + continue + if not isinstance(entry, dict) or entry.get("type") != "message": + continue + # One malformed line (a non-dict where a dict belongs) must not take + # every real stat down with it, so each shape is checked explicitly. + message = entry.get("message") + if not isinstance(message, dict) or message.get("role") != "assistant": + continue + # Exact provider match: a custom `zai-proxy` gateway is not this + # subscription, the same discipline the stock collectors keep. + provider = str(message.get("provider") or "").lower() + if provider not in ("zai", "zhipu"): + continue + usage = message.get("usage") + if not isinstance(usage, dict): + continue + inp = number(usage.get("input")) + # pi reports reasoning tokens apart from output; both are generated. + out = number(usage.get("output")) + number(usage.get("reasoning")) + cread = number(usage.get("cacheRead")) + cwrite = number(usage.get("cacheWrite")) + if not (inp or out or cread or cwrite): + inp = number(usage.get("totalTokens")) + if not (inp or out or cread or cwrite): + continue + message_key = str(path) + ":" + str(entry.get("id") or "") + if message_key in seen: + continue + seen.add(message_key) + day = local_day(entry.get("timestamp") or message.get("timestamp")) + model = str(message.get("model") or "glm") + add_usage(day, str(path), model, inp, out, cread, cwrite) + + return { + "hasLocalStats": bool(model_usage), + "todayPrompts": totals["today_prompts"], + "todaySessions": len(today_sessions), + "todayTotalTokens": totals["today_tokens"], + "todayTokensByModel": today_by_model, + "recentDays": [recent[day] for day in recent_dates], + "totalPrompts": totals["total_prompts"], + "totalSessions": len(total_sessions), + "activeDays": len(active_dates), + "activeDates": sorted(active_dates), + "modelUsage": model_usage, + } + + +def scan_cache_paths() -> tuple[Path, Path, Path]: + # The digest covers every data path the scan reads (both session roots live + # under the home directory). + digest = hashlib.sha1(str(Path.home()).encode("utf-8")).hexdigest()[:16] + root = cache_root() + return root / f"zai-scan-{digest}.json", root / f"zai-scan-{digest}.lock", root / "zai-limits.lock" + + +def read_cached_stats(cache_file: Path, max_age_seconds: float) -> dict[str, Any] | None: + cached = read_fresh_json(cache_file, max_age_seconds) + if not isinstance(cached, dict) or cached.get("schemaVersion") != 1: + return None + # today* fields only mean "today" on the day they were scanned. A cache from + # another local date (midnight passed, or the clock moved) is a miss. + if cached.get("scanDate") != today: + return None + stats = cached.get("stats") + if not isinstance(stats, dict): + return None + if not all(key in stats for key in ("todayPrompts", "todayTotalTokens", "recentDays", "activeDates", "modelUsage")): + return None + return stats + + +def write_cached_stats(cache_file: Path, stats: dict[str, Any]) -> None: + try: + write_json(cache_file, {"schemaVersion": 1, "scanDate": today, "stats": stats}) + except Exception as exc: + print(f"omarchy-agent-usage-zai: could not write usage cache ({exc})", file=sys.stderr) + + +def cached_local_stats(max_age: float) -> dict[str, Any]: + """Local stats, with the cache as a pure optimization. + + Any cache-layer failure degrades to a direct scan and a warning: the JSON + record is the contract, the cache is not. + """ + try: + return _cached_local_stats(max_age) + except Exception as exc: + print(f"omarchy-agent-usage-zai: cache unavailable ({exc}); scanning directly", file=sys.stderr) + return scan_sessions() + + +def _cached_local_stats(max_age: float) -> dict[str, Any]: + cache_file, lock_file, _ = scan_cache_paths() + + cached = read_cached_stats(cache_file, max_age) + if cached is not None: + return cached + + with exclusive_lock(lock_file): + cached = read_cached_stats(cache_file, max_age) + if cached is not None: + return cached + stats = scan_sessions() + write_cached_stats(cache_file, stats) + return stats + + +# -------------------------------------------------------------------------- +# Assembly +# -------------------------------------------------------------------------- + +def scan(force: bool, limits_only: bool) -> dict[str, Any]: + record = base_record() + max_age = 0 if force else (LIMITS_ONLY_REUSE_SECONDS if limits_only else SCAN_REUSE_SECONDS) + record.update(cached_local_stats(max_age)) + + api_key, platform = credentials() + if not api_key: + record["authHelpText"] = AUTH_HELP_NO_KEY + record["ready"] = record["hasLocalStats"] + return record + + limits = collect_limits(api_key, platform, force) + record["limits"] = limits["limits"] + if limits["tierLabel"]: + record["tierLabel"] = limits["tierLabel"] + record["usageStatusText"] = limits["usageStatusText"] + if limits["authHelpText"]: + record["authHelpText"] = limits["authHelpText"] + if limits["retryAdvised"]: + record["retryAdvised"] = True + + record["ready"] = bool(record["limits"] or record["hasLocalStats"]) + return record + + +def main() -> int: + parser = argparse.ArgumentParser(description="Print the Z.ai GLM usage record as JSON") + # --force rescans and re-probes, ignoring both caches. --limits-only reuses a + # recent local scan (only the limits probe must be fresh); a bare run reuses a + # scan only long enough to dedup concurrent collector runs. + parser.add_argument("--force", action="store_true") + parser.add_argument("--limits-only", action="store_true") + args = parser.parse_args() + + try: + record = scan(args.force, args.limits_only) + except Exception as error: # noqa: BLE001 - a collector must always emit a record + record = base_record(usageStatusText="Z.ai unavailable", authHelpText="Z.ai usage scan failed") + print(f"omarchy-agent-usage-zai: {type(error).__name__}", file=sys.stderr) + print(json.dumps(record, separators=(",", ":"))) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/shell/plugins/agents/README.md b/shell/plugins/agents/README.md index 34cb0f39344..05ca49816e5 100644 --- a/shell/plugins/agents/README.md +++ b/shell/plugins/agents/README.md @@ -76,6 +76,7 @@ light surfaces — and the bar glyph stands in when there is none. | `fireworks` | Estimated prepaid balance: configured funding minus rated account costs | Fireworks billing API, grouped by day and model for the last 30 days | | `hermes` | None (Hermes has no subscription of its own) | `state.db` in Hermes' home and in each named profile under `~/.hermes/profiles`, read-only | | `agy` | Google's Cloud Code API (`loadCodeAssist` for the plan, `retrieveUserQuotaSummary` for the 5-hour and weekly windows) | `history.jsonl`, `conversation_summaries.db`, and each conversation's `transcript.jsonl` under `~/.gemini/antigravity-cli/` (`AGY_DIR`) | +| `zai` | Z.ai's coding-plan monitor endpoint (5-hour session + 7-day weekly) | pi/omp sessions on the Z.ai provider (`glm-*` models) | When `~/.local/state/omarchy/agents/accounts/.json` registers more than one account, the `claude`, `codex`, and `grok` records @@ -107,6 +108,12 @@ again. Fireworks reads opencode stores in `~/.local/share/opencode/auth.json` when Fireworks is signed in there. +Z.ai limits need a key that carries a GLM Coding Plan; a pay-as-you-go key +reports "no coding plan" and the panel falls back to local stats only. The key +comes from `~/.config/omarchy/agents/zai.json`, then `ZAI_API_KEY` / +`ZHIPU_API_KEY` in the environment, then pi/omp's `.env`. Set `platform` to +`"zhipu"` in the config for the China (`open.bigmodel.cn`) host. + ### Fireworks balance The collector first asks the account's `:getBalance` endpoint for the real diff --git a/shell/plugins/agents/assets/zai-light.svg b/shell/plugins/agents/assets/zai-light.svg new file mode 100644 index 00000000000..040a6d88830 --- /dev/null +++ b/shell/plugins/agents/assets/zai-light.svg @@ -0,0 +1 @@ +Z.ai diff --git a/shell/plugins/agents/assets/zai.svg b/shell/plugins/agents/assets/zai.svg new file mode 100644 index 00000000000..48c8d1d7397 --- /dev/null +++ b/shell/plugins/agents/assets/zai.svg @@ -0,0 +1 @@ +Z.ai diff --git a/shell/plugins/agents/manifest.json b/shell/plugins/agents/manifest.json index a62fe2564cb..02160885108 100644 --- a/shell/plugins/agents/manifest.json +++ b/shell/plugins/agents/manifest.json @@ -24,7 +24,8 @@ "grok": { "enabled": true }, "fireworks": { "enabled": true }, "hermes": { "enabled": true }, - "agy": { "enabled": true } + "agy": { "enabled": true }, + "zai": { "enabled": true } }, "refreshIntervalSec": 900, "syncMode": "Off", diff --git a/test/shell.d/agent-usage-zai-limits-test.sh b/test/shell.d/agent-usage-zai-limits-test.sh new file mode 100755 index 00000000000..277c03e329a --- /dev/null +++ b/test/shell.d/agent-usage-zai-limits-test.sh @@ -0,0 +1,151 @@ +#!/bin/bash + +source "$(dirname "$0")/base-test.sh" + +require_command jq +require_command python3 + +# fetch_limits/collect_limits reach Z.ai, so the readers that interpret the +# response are exercised on their own: the collector loads as a module and a +# recorded payload (or a raised error) stands in for the network. +CACHE_HOME=$(mktemp -d) +trap 'rm -rf "$CACHE_HOME"' EXIT + +read_limits() { + COLLECTOR="$ROOT/bin/omarchy-agent-usage-zai" PAYLOAD="$1" python3 - <<'PY' +import importlib.machinery, importlib.util, io, json, os + +loader = importlib.machinery.SourceFileLoader("collector", os.environ["COLLECTOR"]) +spec = importlib.util.spec_from_loader(loader.name, loader) +collector = importlib.util.module_from_spec(spec) +loader.exec_module(collector) + +collector.urllib.request.urlopen = lambda request, timeout=None: io.BytesIO(os.environ["PAYLOAD"].encode()) +print(json.dumps(collector.fetch_limits("token", "zai"))) +PY +} + +# force, cached-json, mode (success|transport|http401), payload -> {result, probes, cached} +drive() { + COLLECTOR="$ROOT/bin/omarchy-agent-usage-zai" FORCE="$1" CACHED="$2" MODE="$3" PAYLOAD="$4" \ + XDG_CACHE_HOME="$CACHE_HOME" python3 - <<'PY' +import importlib.machinery, importlib.util, io, json, os, urllib.error + +loader = importlib.machinery.SourceFileLoader("collector", os.environ["COLLECTOR"]) +spec = importlib.util.spec_from_loader(loader.name, loader) +collector = importlib.util.module_from_spec(spec) +loader.exec_module(collector) + +cache = collector.cache_root() / "zai-limits.json" +cached = os.environ["CACHED"] +if cached: + cache.write_text(cached, encoding="utf-8") +elif cache.exists(): + cache.unlink() + +probes = [] +mode = os.environ["MODE"] + +def urlopen(request, timeout=None): + probes.append(1) + if mode == "transport": + raise urllib.error.URLError("no route to host") + if mode == "http401": + raise urllib.error.HTTPError("u", 401, "no", {}, None) + return io.BytesIO(os.environ["PAYLOAD"].encode()) + +collector.urllib.request.urlopen = urlopen +result = collector.collect_limits("token", "zai", os.environ["FORCE"] == "true") +out = {"result": result, "probes": len(probes)} +if cache.exists(): + out["cached"] = json.loads(cache.read_text(encoding="utf-8")) +print(json.dumps(out)) +PY +} + +soon_ms=$(python3 -c 'import time; print(round((time.time() + 3*3600) * 1000))') +week_sec=$(python3 -c 'import time; print(round(time.time() + 5*86400))') +week_date=$(python3 -c "import datetime; print(datetime.datetime.fromtimestamp($week_sec, datetime.timezone.utc).date().isoformat())") +open_at=$(python3 -c 'import datetime as dt; print((dt.datetime.now(dt.timezone.utc) + dt.timedelta(hours=3)).isoformat())') +past_at=$(python3 -c 'import datetime as dt; print((dt.datetime.now(dt.timezone.utc) - dt.timedelta(hours=3)).isoformat())') + +# The mapped session + weekly windows, then the entries that must drop: a window +# tagged with an unmapped (unit, number) and one whose percentage will not parse. +# nextResetTime is milliseconds for the first and plain seconds for the second, +# to prove a seconds epoch is not rendered 1000x in the future. +payload=$(jq -nc --argjson sm "$soon_ms" --argjson ws "$week_sec" '{ + data: { level: "pro", limits: [ + { type: "CREDIT_LIMIT", unit: 3, number: 5, percentage: 42.0, nextResetTime: $sm }, + { type: "TOKENS_LIMIT", unit: 6, number: 1, percentage: 10, nextResetTime: $ws }, + { type: "CREDIT_LIMIT", unit: 9, number: 9, percentage: 5, nextResetTime: $sm }, + { type: "CREDIT_LIMIT", unit: 3, number: 5, percentage: null, nextResetTime: $sm } + ] } +}') +limits=$(read_limits "$payload") + +[[ $(jq -c '[.limits[] | {label, percent}]' <<<"$limits") == '[{"label":"Session (5-hour)","percent":0.42},{"label":"Weekly (7-day)","percent":0.1}]' ]] || + fail "Z.ai collector maps both windows and drops the unmapped and unparseable rows" "$limits" +[[ $(jq -r '.tierLabel' <<<"$limits") == "Pro" ]] || + fail "Z.ai collector titles the plan tier" "$limits" +[[ $(jq -r '.limits[1].resetsAt' <<<"$limits") == "$week_date"* ]] || + fail "Z.ai collector reads a plain-seconds reset time, not milliseconds" "$limits" +pass "Z.ai collector reads windows, tier, seconds/ms resets, and drops bad rows" + +# A key with no plan explains the empty section without blaming the key. +noplan=$(drive false "" success '{"success":false,"msg":"This API key has no coding plan"}') +[[ $(jq -r '.result.usageStatusText' <<<"$noplan") == "No GLM Coding Plan on this key" ]] || + fail "Z.ai collector reports a key without a coding plan" "$noplan" +[[ $(jq -r '.result.authHelpText' <<<"$noplan") == "" && $(jq -r '.result.retryAdvised' <<<"$noplan") == "false" ]] || + fail "Z.ai collector does not tell a working key to set a key" "$noplan" +pass "Z.ai collector reports a no-plan key without blaming it" + +# A rejected key is the one case that should point back at the key. +rejected=$(drive false "" http401 "") +[[ $(jq -r '.result.authHelpText' <<<"$rejected") == *"rejected the saved API key"* ]] || + fail "Z.ai collector blames the key only when the key is rejected" "$rejected" +pass "Z.ai collector points at the key only on a 401" + +# A transport failure with an open cached window keeps the last known figures, +# drops the window that already reset, and asks the shell to retry sooner. +cache=$(jq -nc --arg open "$open_at" --arg past "$past_at" '{ + fetchedAtMs: 1, tierLabel: "Pro", limits: [ + { label: "Session (5-hour)", percent: 0.31, resetsAt: $past }, + { label: "Weekly (7-day)", percent: 0.11, resetsAt: $open } + ] +}') +unreachable=$(drive false "$cache" transport "") +[[ $(jq -c '[.result.limits[].label]' <<<"$unreachable") == '["Weekly (7-day)"]' ]] || + fail "Z.ai collector keeps only cached windows that have not reset" "$unreachable" +[[ $(jq -r '.result.tierLabel' <<<"$unreachable") == "Pro" && $(jq -r '.result.retryAdvised' <<<"$unreachable") == "true" ]] || + fail "Z.ai collector keeps the cached tier and advises a retry after a transport failure" "$unreachable" +pass "Z.ai collector serves the last known limits and advises a retry when offline" + +# Repeated panel opens inside the probe interval share one answer; --force does not. +fresh=$(jq -nc --arg open "$open_at" --argjson now "$(python3 -c 'import time; print(round(time.time() * 1000))')" '{ + fetchedAtMs: $now, tierLabel: "Pro", limits: [{ label: "Weekly (7-day)", percent: 0.11, resetsAt: $open }] +}') +success_payload=$(jq -nc --argjson sm "$soon_ms" '{ data: { level: "max", limits: [ + { type: "CREDIT_LIMIT", unit: 3, number: 5, percentage: 44.0, nextResetTime: $sm } +] } }') + +reused=$(drive false "$fresh" success "$success_payload") +[[ $(jq -r '.probes' <<<"$reused") == "0" && $(jq -c '[.result.limits[].percent]' <<<"$reused") == "[0.11]" ]] || + fail "Z.ai collector reuses a cache younger than the probe interval" "$reused" +pass "Z.ai collector reuses a probe younger than the interval" + +forced=$(drive true "$fresh" success "$success_payload") +[[ $(jq -r '.probes' <<<"$forced") == "1" && $(jq -c '[.result.limits[].percent]' <<<"$forced") == "[0.44]" ]] || + fail "Z.ai collector re-probes on --force despite a fresh cache" "$forced" +[[ $(jq -c '[.cached.limits[].percent]' <<<"$forced") == "[0.44]" && $(jq -r '.cached.tierLabel' <<<"$forced") == "Max" ]] || + fail "Z.ai collector caches a successful probe for the next run" "$forced" +pass "Z.ai collector re-probes on --force and caches the result" + +# A predictable lock path must not follow and truncate a planted symlink. +sentinel="$CACHE_HOME/sentinel" +printf intact >"$sentinel" +rm -f "$CACHE_HOME/omarchy/agent-usage/zai-limits.lock" +ln -s "$sentinel" "$CACHE_HOME/omarchy/agent-usage/zai-limits.lock" +drive true "$fresh" success "$success_payload" >/dev/null +[[ $(<"$sentinel") == "intact" ]] || + fail "Z.ai collector does not follow a planted lock symlink" +pass "Z.ai collector does not follow a planted lock symlink" diff --git a/test/shell.d/agent-usage-zai-scanner-test.sh b/test/shell.d/agent-usage-zai-scanner-test.sh new file mode 100755 index 00000000000..a83cab33fab --- /dev/null +++ b/test/shell.d/agent-usage-zai-scanner-test.sh @@ -0,0 +1,90 @@ +#!/bin/bash + +source "$(dirname "$0")/base-test.sh" + +require_command jq +require_command python3 + +TEST_HOME=$(mktemp -d) +trap 'rm -rf "$TEST_HOME"' EXIT + +mkdir -p "$TEST_HOME/.pi/agent/sessions/project" + +today="$(date +%Y-%m-%d)" +timestamp="${today}T12:00:00Z" + +# One session file exercises the whole scan contract: +# a zai, today, with reasoning + cache tokens -> counted, today +# b zhipu, today -> counted, today +# c openai provider -> ignored +# d zai but a user turn -> ignored +# a' a repeated message id -> deduped +# f zai, no top-level timestamp, message.timestamp -> counted, NOT today +# e malformed JSON -> skipped, rest survive +session="$TEST_HOME/.pi/agent/sessions/project/pi.jsonl" +cat >"$session" </dev/null | head -n 1) +[[ -n $cache_file && $(jq -r '.schemaVersion' "$cache_file") == "1" && $(jq -r '.stats.todayTotalTokens' "$cache_file") == "50" ]] || + fail "Z.ai collector writes a versioned local-stats cache on first scan" "$result" +pass "Z.ai collector writes a local-stats cache on first scan" + +# A new turn arrives; --limits-only must reuse the cache instead of rescanning. +cat >>"$session" < Date: Thu, 1 Oct 2026 22:00:09 -0400 Subject: [PATCH 06/19] Mark stale Z.ai limits and mention Z.ai in the manual Limits kept from an earlier check after a failed one now carry limitsStale and limitsFetchedAt, so the panel dims them. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-zai | 15 ++++++++++++--- manual/17-ai.md | 2 +- test/shell.d/agent-usage-zai-limits-test.sh | 4 ++++ 3 files changed, 17 insertions(+), 4 deletions(-) diff --git a/bin/omarchy-agent-usage-zai b/bin/omarchy-agent-usage-zai index ed684aef791..989c1da33c1 100755 --- a/bin/omarchy-agent-usage-zai +++ b/bin/omarchy-agent-usage-zai @@ -402,9 +402,12 @@ def collect_limits(api_key: str, platform: str, force: bool) -> dict[str, Any]: Reuses a recent probe for PROBE_MIN_INTERVAL_SECONDS so repeated panel opens do not spam the endpoint, and keeps the last known limits as the answer of - record when a later probe fails but its window has not yet reset. + record when a later probe fails but its window has not yet reset. `live` + says whether the limits are this check's (or a probe recent enough to stand + in for one), and `fetchedAtMs` when they were measured, so the panel can dim + limits kept from an earlier check. """ - result = {"limits": [], "tierLabel": "", "usageStatusText": "", "authHelpText": "", "retryAdvised": False} + result = {"limits": [], "tierLabel": "", "usageStatusText": "", "authHelpText": "", "retryAdvised": False, "live": False, "fetchedAtMs": 0} probe_cache = cache_root() / "zai-limits.json" cached = read_fresh_json(probe_cache, float("inf")) or {} @@ -414,6 +417,7 @@ def collect_limits(api_key: str, platform: str, force: bool) -> dict[str, Any]: fetched_at = number(cached.get("fetchedAtMs")) / 1000 if isinstance(cached, dict) else 0 if fallback and not force and time.time() - fetched_at < PROBE_MIN_INTERVAL_SECONDS: result["limits"], result["tierLabel"] = fallback, cached_tier + result.update(live=True, fetchedAtMs=round(fetched_at * 1000)) return result with exclusive_lock(scan_cache_paths()[2]): @@ -423,15 +427,17 @@ def collect_limits(api_key: str, platform: str, force: bool) -> dict[str, Any]: fresh = usable_cached_limits(cached) if isinstance(cached, dict) else [] if fresh and time.time() - fetched_at < PROBE_MIN_INTERVAL_SECONDS: result["limits"], result["tierLabel"] = fresh, str(cached.get("tierLabel") or "") + result.update(live=True, fetchedAtMs=round(fetched_at * 1000)) return result probe = fetch_limits(api_key, platform) if probe.get("limits"): result["limits"] = probe["limits"] result["tierLabel"] = probe.get("tierLabel", "") + result.update(live=True, fetchedAtMs=round(time.time() * 1000)) try: write_json(probe_cache, { - "fetchedAtMs": round(time.time() * 1000), + "fetchedAtMs": result["fetchedAtMs"], "limits": probe["limits"], "tierLabel": result["tierLabel"], }) @@ -446,6 +452,7 @@ def collect_limits(api_key: str, platform: str, force: bool) -> dict[str, Any]: result["retryAdvised"] = True if fallback: result["limits"], result["tierLabel"] = fallback, cached_tier + result["fetchedAtMs"] = round(fetched_at * 1000) result["usageStatusText"] = probe.get("status", "") # Only blame the key when the key is the problem; a working key with no plan, # or an endpoint hiccup, must not tell the user to go set a key they have. @@ -633,6 +640,8 @@ def scan(force: bool, limits_only: bool) -> dict[str, Any]: limits = collect_limits(api_key, platform, force) record["limits"] = limits["limits"] + record["limitsStale"] = not limits["live"] + record["limitsFetchedAt"] = limits["fetchedAtMs"] if limits["tierLabel"]: record["tierLabel"] = limits["tierLabel"] record["usageStatusText"] = limits["usageStatusText"] diff --git a/manual/17-ai.md b/manual/17-ai.md index fb2f6882256..655ebfb90f6 100644 --- a/manual/17-ai.md +++ b/manual/17-ai.md @@ -38,7 +38,7 @@ There are terminal shortcuts too: `a` runs the default agent inline in the curre ### The agents panel -The top bar has an agents icon. Before you've set up an agent, its panel walks you through signing in to Claude, Codex, or Grok, and the first one becomes your default agent. After that, the panel puts every subscription on one page: each agent and account with its plan and the percentage used of its 5-hour session and weekly limits (or the remaining prepaid balance), under a line that rotates through what you've used across all of them, like tokens this week and your most used model. Under _Make something_ it offers starter prompts for a new theme, plugin, or app that hand the task to your default agent. Apps are built the way Omarchy's own Hype, Monologue, and Omacut are, in C++ and Qt Quick, following the omarchy-app agent skill, and the compiler and Qt pieces they need come installed. Claude Code, Codex, Grok, and Fireworks are covered out of the box, as is Antigravity with its 5-hour and weekly limits, and Hermes shows its sessions and tokens, from every profile, once you've used it. +The top bar has an agents icon. Before you've set up an agent, its panel walks you through signing in to Claude, Codex, or Grok, and the first one becomes your default agent. After that, the panel puts every subscription on one page: each agent and account with its plan and the percentage used of its 5-hour session and weekly limits (or the remaining prepaid balance), under a line that rotates through what you've used across all of them, like tokens this week and your most used model. Under _Make something_ it offers starter prompts for a new theme, plugin, or app that hand the task to your default agent. Apps are built the way Omarchy's own Hype, Monologue, and Omacut are, in C++ and Qt Quick, following the omarchy-app agent skill, and the compiler and Qt pieces they need come installed. Claude Code, Codex, Grok, and Fireworks are covered out of the box, as is Antigravity with its 5-hour and weekly limits, and Hermes shows its sessions and tokens, from every profile, once you've used it. Provider plans and keys you use through pi, omp, or opencode show up as well: the Z.ai GLM Coding Plan with its 5-hour and weekly limits. Left-click the bar icon for the panel, right-click to launch your default agent. The usage records behind it are regenerated every 15 minutes by `omarchy agent usage-update`, and the panel can even merge usage from your other machines via a synced folder. See the README under `$OMARCHY_PATH/shell/plugins/agents/` for the full settings. diff --git a/test/shell.d/agent-usage-zai-limits-test.sh b/test/shell.d/agent-usage-zai-limits-test.sh index 277c03e329a..36915183fac 100755 --- a/test/shell.d/agent-usage-zai-limits-test.sh +++ b/test/shell.d/agent-usage-zai-limits-test.sh @@ -118,6 +118,8 @@ unreachable=$(drive false "$cache" transport "") fail "Z.ai collector keeps only cached windows that have not reset" "$unreachable" [[ $(jq -r '.result.tierLabel' <<<"$unreachable") == "Pro" && $(jq -r '.result.retryAdvised' <<<"$unreachable") == "true" ]] || fail "Z.ai collector keeps the cached tier and advises a retry after a transport failure" "$unreachable" +[[ $(jq -c '{live: .result.live, fetchedAtMs: .result.fetchedAtMs}' <<<"$unreachable") == '{"live":false,"fetchedAtMs":1}' ]] || + fail "Z.ai collector marks kept limits stale with when they were fetched" "$unreachable" pass "Z.ai collector serves the last known limits and advises a retry when offline" # Repeated panel opens inside the probe interval share one answer; --force does not. @@ -138,6 +140,8 @@ forced=$(drive true "$fresh" success "$success_payload") fail "Z.ai collector re-probes on --force despite a fresh cache" "$forced" [[ $(jq -c '[.cached.limits[].percent]' <<<"$forced") == "[0.44]" && $(jq -r '.cached.tierLabel' <<<"$forced") == "Max" ]] || fail "Z.ai collector caches a successful probe for the next run" "$forced" +[[ $(jq -r '.result.live' <<<"$forced") == "true" && $(jq -r '.result.fetchedAtMs == .cached.fetchedAtMs' <<<"$forced") == "true" ]] || + fail "Z.ai collector marks a fresh probe live" "$forced" pass "Z.ai collector re-probes on --force and caches the result" # A predictable lock path must not follow and truncate a planted symlink. From 34d02a94dc54dd45e57b4652f1d4085e04f37371 Mon Sep 17 00:00:00 2001 From: Dante Loi Date: Thu, 1 Oct 2026 22:00:55 -0400 Subject: [PATCH 07/19] Add Ollama Cloud usage collector to the agents panel Ship omarchy-agent-usage-ollama, which reads the account's 5-hour session and 7-day weekly limits from Ollama's usage endpoint and local token stats from pi and omp sessions on the ollama-cloud provider. --- bin/omarchy-agent-usage-ollama | 434 ++++++++++++++++++ shell/plugins/agents/README.md | 25 + shell/plugins/agents/assets/ollama-light.svg | 4 + shell/plugins/agents/assets/ollama.svg | 4 + shell/plugins/agents/manifest.json | 3 +- .../agent-usage-ollama-scanner-test.sh | 211 +++++++++ 6 files changed, 680 insertions(+), 1 deletion(-) create mode 100755 bin/omarchy-agent-usage-ollama create mode 100644 shell/plugins/agents/assets/ollama-light.svg create mode 100644 shell/plugins/agents/assets/ollama.svg create mode 100755 test/shell.d/agent-usage-ollama-scanner-test.sh diff --git a/bin/omarchy-agent-usage-ollama b/bin/omarchy-agent-usage-ollama new file mode 100755 index 00000000000..da5477abf0f --- /dev/null +++ b/bin/omarchy-agent-usage-ollama @@ -0,0 +1,434 @@ +#!/usr/bin/python3 +# omarchy:summary=Print the Ollama Cloud usage record as JSON +# omarchy:args=[--force] [--limits-only] +# omarchy:hidden=true +"""Collect Ollama Cloud usage into one display-ready JSON record. + +Everything the agents panel shows for Ollama Cloud comes from this one +command: local token stats from pi and omp sessions that ran on the +`ollama-cloud` provider, and the account's 5-hour session and 7-day weekly +limits from Ollama's 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 json +import os +import sys +import tempfile +import time +import urllib.error +import urllib.request +from pathlib import Path +from typing import Any + +AGENT_ID = "ollama" +AGENT_NAME = "Ollama Cloud" +USAGE_ENDPOINT = "https://ollama.com/api/usage" +ME_ENDPOINT = "https://ollama.com/api/me" +PROBE_MIN_INTERVAL_SECONDS = 15 +SESSION_WINDOW_SECONDS = 5 * 3600 +WEEKLY_WINDOW_SECONDS = 7 * 86400 +# Ollama's weekly windows end Monday 00:00 UTC, which is 4 days after the +# epoch (a Thursday) — hence the offset in the reset computation. Session +# windows are epoch-aligned 5-hour blocks. The usage endpoint reports the +# fractions but not the reset times, and the resets are the same for every +# account, so both are computed rather than fetched. +WEEKLY_EPOCH_OFFSET_SECONDS = 4 * 86400 + + +def config_dir() -> Path: + return Path(os.environ.get("XDG_CONFIG_HOME", Path.home() / ".config")) / "omarchy" / "agents" + + +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 read_config() -> dict[str, Any]: + try: + data = json.loads((config_dir() / "ollama.json").read_text(encoding="utf-8")) + except Exception: + return {} + return data if isinstance(data, dict) else {} + + +def api_key() -> str: + env = os.environ.get("OLLAMA_API_KEY", "").strip() + if env: + return env + configured = str(read_config().get("apiKey") or "").strip() + if configured: + return configured + # pi and omp store the Ollama Cloud key they signed in with; either one + # proves the account, and the usage endpoint accepts the same key. + for auth_path in (Path.home() / ".pi" / "agent" / "auth.json", Path.home() / ".omp" / "agent" / "auth.json"): + try: + data = json.loads(auth_path.read_text(encoding="utf-8")) + entry = data.get("ollama-cloud") + if isinstance(entry, dict) and str(entry.get("key") or "").strip(): + return str(entry["key"]).strip() + except Exception: + continue + return "" + + +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() + + 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 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, + } + + +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 + + +# ---------------------------------------------------------------- local scan +# +# pi and omp record every assistant message with its provider, model, and +# token usage in JSONL session files. The `ollama-cloud` provider is the +# Ollama Cloud API, which serves Claude models alongside the open ones. + + +def scan_sessions(max_age_seconds: float, model_prefix: str) -> dict[str, Any] | None: + roots = [ + Path.home() / ".pi" / "agent" / "sessions", + Path.home() / ".omp" / "agent" / "sessions", + ] + cache_file = cache_root() / "ollama-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): + # Cheap pre-filter before JSON parsing keeps files with unrelated + # lines inexpensive. + if '"usage"' not in line or '"ollama-cloud"' 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 + if str(message.get("provider") or "") != "ollama-cloud": + 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 = number(usage.get("input")) + output_tokens = number(usage.get("output")) + cache_read = number(usage.get("cacheRead")) + cache_write = number(usage.get("cacheWrite")) + total = input_tokens + output_tokens + cache_read + cache_write + if total <= 0: + total = number(usage.get("totalTokens")) + input_tokens = total + if total <= 0: + continue + # Ollama model ids carry a version tag ("deepseek-v4-pro:0813"); + # merge tags so the chart groups by model, not by release. + model = str(message.get("model") or "ollama").split(":", 1)[0] + if model_prefix and not model.startswith(model_prefix): + continue + 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: + # 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 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 + + +# ------------------------------------------------------------------- limits + + +def clamp_fraction(value: Any) -> float: + try: + n = float(value) + except Exception: + return -1.0 + if not (n >= 0): + return -1.0 + return min(1.0, n) + + +def reset_iso(window_seconds: int, epoch_offset_seconds: int = 0) -> str: + now = time.time() + nxt = now - ((now - epoch_offset_seconds) % window_seconds) + window_seconds + return dt.datetime.fromtimestamp(nxt, dt.timezone.utc).isoformat() + + +def probe_limits(access_key: str) -> dict[str, Any]: + request = urllib.request.Request( + USAGE_ENDPOINT, + headers={"Authorization": "Bearer " + access_key, "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: + if error.code in (401, 403): + return {"ok": False, "auth": True, "helpText": "The Ollama Cloud API key was rejected. Local stats are still shown."} + return {"ok": False, "helpText": f"Ollama's usage endpoint returned status {error.code}. Local stats are still shown."} + 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 Ollama's usage endpoint. Retrying shortly. Local stats are still shown.", + } + + limits = payload.get("limits") if isinstance(payload, dict) else None + if not isinstance(limits, dict): + return {"ok": False, "helpText": "Ollama's usage endpoint returned no limits. Local stats are still shown."} + + out = [] + session = limits.get("session") + weekly = limits.get("weekly") + if isinstance(session, dict): + percent = clamp_fraction(session.get("usage")) + if percent >= 0: + out.append({"label": "Session (5-hour)", "percent": percent, "resetsAt": reset_iso(SESSION_WINDOW_SECONDS)}) + if isinstance(weekly, dict): + percent = clamp_fraction(weekly.get("usage")) + if percent >= 0: + out.append({"label": "Weekly (7-day)", "percent": percent, "resetsAt": reset_iso(WEEKLY_WINDOW_SECONDS, WEEKLY_EPOCH_OFFSET_SECONDS)}) + if not out: + return {"ok": False, "helpText": "Ollama's usage endpoint returned no limits. Local stats are still shown."} + return {"ok": True, "limits": out} + + +def probe_plan(access_key: str) -> str: + request = urllib.request.Request( + ME_ENDPOINT, + headers={"Authorization": "Bearer " + access_key, "Accept": "application/json"}, + method="POST", + ) + try: + with urllib.request.urlopen(request, timeout=10) as response: + payload = json.loads(response.read().decode("utf-8", errors="replace")) + except Exception: + return "" + plan = str(payload.get("Plan") or "").strip() + return plan[0].upper() + plan[1:] if plan else "" + + +def collect_limits(access_key: str, force: bool) -> dict[str, Any]: + result = {"limits": [], "tierLabel": "", "usageStatusText": "", "authHelpText": ""} + + # 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() / "ollama-limits.json" + cached = read_fresh_json(probe_cache, float("inf")) or {} + fallback = cached.get("limits") if isinstance(cached.get("limits"), list) else [] + result["tierLabel"] = str(cached.get("plan") or "") + + if access_key == "": + result["limits"] = fallback + result["usageStatusText"] = "Waiting for auth" + result["authHelpText"] = "Set OLLAMA_API_KEY, or sign in to pi with an Ollama Cloud key, to show the account limits." + 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 + return result + + probe = probe_limits(access_key) + if probe["ok"]: + plan = probe_plan(access_key) + result["limits"] = probe["limits"] + result["tierLabel"] = plan + write_json(probe_cache, {"fetchedAtMs": round(time.time() * 1000), "limits": probe["limits"], "plan": plan}) + 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 + else: + result["usageStatusText"] = "Ollama limits unavailable" + result["authHelpText"] = probe["helpText"] + return result + + +# -------------------------------------------------------------------- record + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--force", action="store_true", help="rescan sessions and re-probe limits, ignoring caches") + parser.add_argument("--limits-only", action="store_true", help="reuse any recent session scan; only the limits probe must be fresh") + parser.add_argument("--cache-seconds", type=float, default=20) + args = parser.parse_args() + + config = read_config() + model_prefix = str(config.get("modelPrefix") or "").strip().lower() + scan_age = 0 if args.force else (900 if args.limits_only else args.cache_seconds) + stats = scan_sessions(scan_age, model_prefix) + + limits = collect_limits(api_key(), args.force) + + record = { + "schemaVersion": 1, + "id": AGENT_ID, + "name": AGENT_NAME, + "updatedAt": dt.datetime.now(dt.timezone.utc).isoformat(), + "ready": number((stats or {}).get("totalPrompts")) > 0 or len(limits["limits"]) > 0, + "hasLocalStats": True, + "tierLabel": limits["tierLabel"], + "usageStatusText": limits["usageStatusText"], + "authHelpText": limits["authHelpText"], + "limits": limits["limits"], + } + if limits.get("retryAdvised"): + record["retryAdvised"] = True + record.update(stats or {}) + print(json.dumps(record, separators=(",", ":"), sort_keys=True)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/shell/plugins/agents/README.md b/shell/plugins/agents/README.md index 05ca49816e5..5cbb1c849a0 100644 --- a/shell/plugins/agents/README.md +++ b/shell/plugins/agents/README.md @@ -77,6 +77,7 @@ light surfaces — and the bar glyph stands in when there is none. | `hermes` | None (Hermes has no subscription of its own) | `state.db` in Hermes' home and in each named profile under `~/.hermes/profiles`, read-only | | `agy` | Google's Cloud Code API (`loadCodeAssist` for the plan, `retrieveUserQuotaSummary` for the 5-hour and weekly windows) | `history.jsonl`, `conversation_summaries.db`, and each conversation's `transcript.jsonl` under `~/.gemini/antigravity-cli/` (`AGY_DIR`) | | `zai` | Z.ai's coding-plan monitor endpoint (5-hour session + 7-day weekly) | pi/omp sessions on the Z.ai provider (`glm-*` models) | +| `ollama` | Ollama Cloud's usage endpoint (5-hour session + 7-day weekly) | pi and omp sessions on the `ollama-cloud` provider | When `~/.local/state/omarchy/agents/accounts/.json` registers more than one account, the `claude`, `codex`, and `grok` records @@ -114,6 +115,30 @@ comes from `~/.config/omarchy/agents/zai.json`, then `ZAI_API_KEY` / `ZHIPU_API_KEY` in the environment, then pi/omp's `.env`. Set `platform` to `"zhipu"` in the config for the China (`open.bigmodel.cn`) host. +### Ollama Cloud limits + +Ollama Cloud reads `OLLAMA_API_KEY` first, then the key pi or omp signed in +with in `~/.pi/agent/auth.json` or `~/.omp/agent/auth.json`. + +Ollama's usage endpoint reports the account's session and weekly usage as +0..1 fractions but not the reset times; the resets are the same for every +account (epoch-aligned 5-hour session windows, weekly windows ending Monday +00:00 UTC), so the collector computes them. The limits cover the whole +account — every model the key can reach — while the local stats count pi and +omp sessions on the `ollama-cloud` provider. To count only the Claude models +Ollama Cloud serves, set a model prefix in +`~/.config/omarchy/agents/ollama.json`: + +```json +{ + "modelPrefix": "claude", + "apiKey": "" +} +``` + +`modelPrefix` restricts the local stats to models whose name starts with the +prefix (empty means every model); `apiKey` overrides the key lookup order. + ### Fireworks balance The collector first asks the account's `:getBalance` endpoint for the real diff --git a/shell/plugins/agents/assets/ollama-light.svg b/shell/plugins/agents/assets/ollama-light.svg new file mode 100644 index 00000000000..9c267376af3 --- /dev/null +++ b/shell/plugins/agents/assets/ollama-light.svg @@ -0,0 +1,4 @@ + + Ollama + + diff --git a/shell/plugins/agents/assets/ollama.svg b/shell/plugins/agents/assets/ollama.svg new file mode 100644 index 00000000000..3729abfd1d6 --- /dev/null +++ b/shell/plugins/agents/assets/ollama.svg @@ -0,0 +1,4 @@ + + Ollama + + diff --git a/shell/plugins/agents/manifest.json b/shell/plugins/agents/manifest.json index 02160885108..07e4880ca86 100644 --- a/shell/plugins/agents/manifest.json +++ b/shell/plugins/agents/manifest.json @@ -25,7 +25,8 @@ "fireworks": { "enabled": true }, "hermes": { "enabled": true }, "agy": { "enabled": true }, - "zai": { "enabled": true } + "zai": { "enabled": true }, + "ollama": { "enabled": true } }, "refreshIntervalSec": 900, "syncMode": "Off", diff --git a/test/shell.d/agent-usage-ollama-scanner-test.sh b/test/shell.d/agent-usage-ollama-scanner-test.sh new file mode 100755 index 00000000000..d3cbad832d4 --- /dev/null +++ b/test/shell.d/agent-usage-ollama-scanner-test.sh @@ -0,0 +1,211 @@ +#!/bin/bash + +source "$(dirname "$0")/base-test.sh" + +require_command jq +require_command python3 + +TEST_HOME=$(mktemp -d) +trap 'rm -rf "$TEST_HOME"' EXIT + +# Without credentials the collector must still print a full, hidden-by-default +# record: the update runner writes whatever valid JSON appears on stdout. +no_key=$(HOME="$TEST_HOME" XDG_CONFIG_HOME="$TEST_HOME/.config" XDG_CACHE_HOME="$TEST_HOME/.cache" \ + OLLAMA_API_KEY="" "$ROOT/bin/omarchy-agent-usage-ollama" --force) + +[[ $(jq -r '.id + ":" + (.ready | tostring) + ":" + .usageStatusText' <<<"$no_key") == "ollama:false:Waiting for auth" ]] || + fail "Ollama collector prints a valid record without credentials" "$no_key" +pass "Ollama collector prints a valid record without credentials" + +# Pi and omp sessions on the ollama-cloud provider are the local stats; other +# providers in the same files must not leak in, and model version tags merge. +PI_HOME=$(mktemp -d) +trap 'rm -rf "$TEST_HOME" "$PI_HOME"' EXIT +mkdir -p "$PI_HOME/.pi/agent/sessions/project" "$PI_HOME/.omp/agent/sessions/project" + +timestamp="$(date +%Y-%m-%d)T12:00:00Z" +cat >"$PI_HOME/.pi/agent/sessions/project/pi.jsonl" <"$PI_HOME/.omp/agent/sessions/project/omp.jsonl" <"$PI_HOME/.config/omarchy/agents/ollama.json" <<'EOF' +{ "modelPrefix": "kimi" } +EOF + +result=$(HOME="$PI_HOME" XDG_CONFIG_HOME="$PI_HOME/.config" XDG_CACHE_HOME="$PI_HOME/.cache" \ + OLLAMA_API_KEY="" "$ROOT/bin/omarchy-agent-usage-ollama" --force) + +[[ $(jq -c '.modelUsage' <<<"$result") == '{"kimi-k3":{"cacheCreationInputTokens":1,"cacheReadInputTokens":4,"inputTokens":20,"outputTokens":5}}' ]] || + fail "Ollama collector honors the configured model prefix" "$result" +pass "Ollama collector honors the configured model prefix" + +# probe_limits reaches Ollama, so the reader that interprets its answer is +# exercised on its own: the collector loads as a module, and a recorded +# payload stands in for the response. +read_limits() { + COLLECTOR="$ROOT/bin/omarchy-agent-usage-ollama" PAYLOAD="$1" python3 - <<'PY' +import importlib.machinery, importlib.util, io, json, os + +loader = importlib.machinery.SourceFileLoader("collector", os.environ["COLLECTOR"]) +spec = importlib.util.spec_from_loader(loader.name, loader) +collector = importlib.util.module_from_spec(spec) +loader.exec_module(collector) + +collector.urllib.request.urlopen = lambda request, timeout=None: io.BytesIO(os.environ["PAYLOAD"].encode()) + +print(json.dumps(collector.probe_limits("key"))) +PY +} + +limits=$(read_limits '{ + "activity": { "cost": "0.00" }, + "limits": { + "session": { "usage": 0.434, "models": [{ "name": "deepseek-v4-pro:0813", "request_count": 340 }] }, + "weekly": { "usage": 0.298, "models": [{ "name": "kimi-k3", "request_count": 377 }] } + } +}') + +[[ $(jq -c '[.limits[].label]' <<<"$limits") == '["Session (5-hour)","Weekly (7-day)"]' ]] || + fail "Ollama collector reads the session and weekly windows" "$limits" +[[ $(jq -c '[.limits[].percent]' <<<"$limits") == "[0.434,0.298]" ]] || + fail "Ollama collector keeps the endpoint's 0..1 fractions" "$limits" +pass "Ollama collector reads the session and weekly fractions" + +# Reset times are computed, not fetched: session windows are epoch-aligned +# 5-hour blocks and weekly windows end Monday 00:00 UTC. +resets=$(python3 - "$ROOT/bin/omarchy-agent-usage-ollama" <<'PY' +import datetime as dt +import importlib.machinery, importlib.util, sys + +loader = importlib.machinery.SourceFileLoader("collector", sys.argv[1]) +spec = importlib.util.spec_from_loader(loader.name, loader) +collector = importlib.util.module_from_spec(spec) +loader.exec_module(collector) + +session = dt.datetime.fromisoformat(collector.reset_iso(collector.SESSION_WINDOW_SECONDS)) +weekly = dt.datetime.fromisoformat(collector.reset_iso(collector.WEEKLY_WINDOW_SECONDS, collector.WEEKLY_EPOCH_OFFSET_SECONDS)) +print(f"{session.timestamp() % collector.SESSION_WINDOW_SECONDS == 0}:{weekly.strftime('%A')}:{weekly.hour}:{weekly.minute}") +PY +) + +[[ "$resets" == "True:Monday:0:0" ]] || + fail "Ollama collector computes epoch-aligned session and Monday-midnight weekly resets" "$resets" +pass "Ollama collector computes epoch-aligned session and Monday-midnight weekly resets" + +# A payload with no usable windows explains itself instead of printing none. +empty=$(read_limits '{"limits": {"session": {"usage": "unknown"}, "weekly": null}}') +[[ $(jq -r '.ok' <<<"$empty") == "false" && $(jq -r '.helpText' <<<"$empty") == *"no limits"* ]] || + fail "Ollama collector reports a payload without usable limits" "$empty" +pass "Ollama collector reports a payload without usable limits" + +# collect_limits drives the probe, the plan lookup, and the cache. A mocked +# urlopen answers both endpoints, and the cache decides reuse and fallback. +CACHE_HOME=$(mktemp -d) +trap 'rm -rf "$TEST_HOME" "$PI_HOME" "$CACHE_HOME"' EXIT + +collect_limits() { + COLLECTOR="$ROOT/bin/omarchy-agent-usage-ollama" FORCE="$1" CACHED="$2" MODE="$3" \ + XDG_CACHE_HOME="$CACHE_HOME" python3 - <<'PY' +import importlib.machinery, importlib.util, io, json, os + +loader = importlib.machinery.SourceFileLoader("collector", os.environ["COLLECTOR"]) +spec = importlib.util.spec_from_loader(loader.name, loader) +collector = importlib.util.module_from_spec(spec) +loader.exec_module(collector) + +cache = collector.cache_root() / "ollama-limits.json" +cached = os.environ["CACHED"] +if cached: + cache.write_text(cached, encoding="utf-8") +elif cache.exists(): + cache.unlink() + +probes = [] + +def urlopen(request, timeout=None): + probes.append(request.full_url) + mode = os.environ["MODE"] + if mode == "transport": + raise OSError("no route to host") + if mode == "rejected": + raise collector.urllib.error.HTTPError(request.full_url, 401, "Unauthorized", {}, io.BytesIO(b"")) + if request.full_url.endswith("/api/me"): + return io.BytesIO(b'{"Plan":"pro"}') + return io.BytesIO(b'{"limits":{"session":{"usage":0.44},"weekly":{"usage":0.11}}}') + +collector.urllib.request.urlopen = urlopen +result = collector.collect_limits("key", os.environ["FORCE"] == "true") +print(json.dumps({ + "result": result, + "probes": probes, + "cached": json.loads(cache.read_text(encoding="utf-8")) if cache.exists() else None, +})) +PY +} + +fresh=$(jq -nc --argjson now "$(python3 -c 'import time; print(round(time.time() * 1000))')" '{ + fetchedAtMs: $now, + limits: [{ label: "Weekly (7-day)", percent: 0.11, resetsAt: "" }], + plan: "Pro" +}') +stale=$(jq -nc '{ + fetchedAtMs: 1, + limits: [{ label: "Weekly (7-day)", percent: 0.11, resetsAt: "" }], + plan: "Pro" +}') + +# A live key that cannot reach the endpoint keeps the cached window and asks +# the shell to retry sooner than its regular interval. +unreachable=$(collect_limits false "$stale" transport) +[[ $(jq -c '[.result.limits[].label]' <<<"$unreachable") == '["Weekly (7-day)"]' ]] || + fail "Ollama collector falls back to cache when the probe cannot connect" "$unreachable" +[[ $(jq -r '.result.retryAdvised' <<<"$unreachable") == "true" ]] || + fail "Ollama collector advises a retry after a transport failure" "$unreachable" +pass "Ollama collector falls back to cache when the probe cannot connect" + +# A rejected key says so instead of pretending the endpoint is down. +rejected=$(collect_limits false "" rejected) +[[ $(jq -r '.result.usageStatusText' <<<"$rejected") == "Ollama limits unavailable" ]] || + fail "Ollama collector reports a rejected key" "$rejected" +[[ $(jq -r '.result.authHelpText' <<<"$rejected") == *"rejected"* ]] || + fail "Ollama collector explains a rejected key" "$rejected" +pass "Ollama collector reports a rejected key" + +# A fresh cache answers repeated panel opens without a request apiece. +reused=$(collect_limits false "$fresh" normal) +[[ $(jq -r '.probes | length' <<<"$reused") == "0" && $(jq -c '[.result.limits[].percent]' <<<"$reused") == "[0.11]" ]] || + fail "Ollama collector reuses a cache younger than the probe interval" "$reused" +pass "Ollama collector reuses a cache younger than the probe interval" + +# --force is someone pressing refresh, so the reuse window must not outrank it. +forced=$(collect_limits true "$fresh" normal) +[[ $(jq -r '.probes | length' <<<"$forced") == "2" ]] || + fail "Ollama collector re-probes limits and plan on --force despite a fresh cache" "$forced" +[[ $(jq -c '[.result.limits[].percent]' <<<"$forced") == "[0.44,0.11]" && $(jq -r '.result.tierLabel' <<<"$forced") == "Pro" ]] || + fail "Ollama collector returns the forced probe's numbers and plan" "$forced" +pass "Ollama collector re-probes on --force despite a fresh cache" + +# A probe that lands becomes the next run's fallback. +[[ $(jq -c '[.cached.limits[].percent]' <<<"$forced") == "[0.44,0.11]" && $(jq -r '.cached.plan' <<<"$forced") == "Pro" ]] || + fail "Ollama collector caches a successful probe" "$forced" +pass "Ollama collector caches a successful probe" From 636bd96b7130b729ba9a134a627b7a377e6ca165 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Thu, 1 Oct 2026 22:01:47 -0400 Subject: [PATCH 08/19] Drop Ollama Cloud's guessed reset times and mark stale limits Ollama's usage endpoint names no reset times, and the epoch-aligned windows the collector computed were an unproven guess, so the limits now carry none. Limits kept after a failed check carry limitsStale and limitsFetchedAt, a record without local stats no longer claims some, and the manual mentions Ollama Cloud. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-ollama | 34 ++++++++----------- manual/17-ai.md | 2 +- shell/plugins/agents/README.md | 5 ++- .../agent-usage-ollama-scanner-test.sh | 30 +++++----------- 4 files changed, 27 insertions(+), 44 deletions(-) diff --git a/bin/omarchy-agent-usage-ollama b/bin/omarchy-agent-usage-ollama index da5477abf0f..9ed1503da0c 100755 --- a/bin/omarchy-agent-usage-ollama +++ b/bin/omarchy-agent-usage-ollama @@ -31,14 +31,8 @@ AGENT_NAME = "Ollama Cloud" USAGE_ENDPOINT = "https://ollama.com/api/usage" ME_ENDPOINT = "https://ollama.com/api/me" PROBE_MIN_INTERVAL_SECONDS = 15 -SESSION_WINDOW_SECONDS = 5 * 3600 -WEEKLY_WINDOW_SECONDS = 7 * 86400 -# Ollama's weekly windows end Monday 00:00 UTC, which is 4 days after the -# epoch (a Thursday) — hence the offset in the reset computation. Session -# windows are epoch-aligned 5-hour blocks. The usage endpoint reports the -# fractions but not the reset times, and the resets are the same for every -# account, so both are computed rather than fetched. -WEEKLY_EPOCH_OFFSET_SECONDS = 4 * 86400 +# The usage endpoint reports how much of each window is used but not when it +# resets, so the limits carry no reset time rather than a guessed one. def config_dir() -> Path: @@ -286,12 +280,6 @@ def clamp_fraction(value: Any) -> float: return min(1.0, n) -def reset_iso(window_seconds: int, epoch_offset_seconds: int = 0) -> str: - now = time.time() - nxt = now - ((now - epoch_offset_seconds) % window_seconds) + window_seconds - return dt.datetime.fromtimestamp(nxt, dt.timezone.utc).isoformat() - - def probe_limits(access_key: str) -> dict[str, Any]: request = urllib.request.Request( USAGE_ENDPOINT, @@ -324,11 +312,11 @@ def probe_limits(access_key: str) -> dict[str, Any]: if isinstance(session, dict): percent = clamp_fraction(session.get("usage")) if percent >= 0: - out.append({"label": "Session (5-hour)", "percent": percent, "resetsAt": reset_iso(SESSION_WINDOW_SECONDS)}) + out.append({"label": "Session (5-hour)", "percent": percent, "resetsAt": ""}) if isinstance(weekly, dict): percent = clamp_fraction(weekly.get("usage")) if percent >= 0: - out.append({"label": "Weekly (7-day)", "percent": percent, "resetsAt": reset_iso(WEEKLY_WINDOW_SECONDS, WEEKLY_EPOCH_OFFSET_SECONDS)}) + out.append({"label": "Weekly (7-day)", "percent": percent, "resetsAt": ""}) if not out: return {"ok": False, "helpText": "Ollama's usage endpoint returned no limits. Local stats are still shown."} return {"ok": True, "limits": out} @@ -350,7 +338,10 @@ def probe_plan(access_key: str) -> str: def collect_limits(access_key: str, force: bool) -> dict[str, Any]: - result = {"limits": [], "tierLabel": "", "usageStatusText": "", "authHelpText": ""} + # `live` says whether the limits are this check's (or a probe recent enough + # to stand in for one), `fetchedAtMs` when they were measured, so the panel + # can dim limits kept from an earlier check. + result = {"limits": [], "tierLabel": "", "usageStatusText": "", "authHelpText": "", "live": False, "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 @@ -359,6 +350,7 @@ def collect_limits(access_key: str, force: bool) -> dict[str, Any]: cached = read_fresh_json(probe_cache, float("inf")) or {} fallback = cached.get("limits") if isinstance(cached.get("limits"), list) else [] result["tierLabel"] = str(cached.get("plan") or "") + result["fetchedAtMs"] = number(cached.get("fetchedAtMs")) if access_key == "": result["limits"] = fallback @@ -372,6 +364,7 @@ def collect_limits(access_key: str, force: bool) -> dict[str, Any]: 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_key) @@ -379,7 +372,8 @@ def collect_limits(access_key: str, force: bool) -> dict[str, Any]: plan = probe_plan(access_key) result["limits"] = probe["limits"] result["tierLabel"] = plan - write_json(probe_cache, {"fetchedAtMs": round(time.time() * 1000), "limits": probe["limits"], "plan": plan}) + result.update(live=True, fetchedAtMs=round(time.time() * 1000)) + write_json(probe_cache, {"fetchedAtMs": result["fetchedAtMs"], "limits": probe["limits"], "plan": plan}) return result # The first probe after login often fires before DHCP has handed out a @@ -417,11 +411,13 @@ def main() -> int: "name": AGENT_NAME, "updatedAt": dt.datetime.now(dt.timezone.utc).isoformat(), "ready": number((stats or {}).get("totalPrompts")) > 0 or len(limits["limits"]) > 0, - "hasLocalStats": True, + "hasLocalStats": stats is not None, "tierLabel": limits["tierLabel"], "usageStatusText": limits["usageStatusText"], "authHelpText": limits["authHelpText"], "limits": limits["limits"], + "limitsStale": not limits["live"], + "limitsFetchedAt": limits["fetchedAtMs"], } if limits.get("retryAdvised"): record["retryAdvised"] = True diff --git a/manual/17-ai.md b/manual/17-ai.md index 655ebfb90f6..b1f1ed6d69e 100644 --- a/manual/17-ai.md +++ b/manual/17-ai.md @@ -38,7 +38,7 @@ There are terminal shortcuts too: `a` runs the default agent inline in the curre ### The agents panel -The top bar has an agents icon. Before you've set up an agent, its panel walks you through signing in to Claude, Codex, or Grok, and the first one becomes your default agent. After that, the panel puts every subscription on one page: each agent and account with its plan and the percentage used of its 5-hour session and weekly limits (or the remaining prepaid balance), under a line that rotates through what you've used across all of them, like tokens this week and your most used model. Under _Make something_ it offers starter prompts for a new theme, plugin, or app that hand the task to your default agent. Apps are built the way Omarchy's own Hype, Monologue, and Omacut are, in C++ and Qt Quick, following the omarchy-app agent skill, and the compiler and Qt pieces they need come installed. Claude Code, Codex, Grok, and Fireworks are covered out of the box, as is Antigravity with its 5-hour and weekly limits, and Hermes shows its sessions and tokens, from every profile, once you've used it. Provider plans and keys you use through pi, omp, or opencode show up as well: the Z.ai GLM Coding Plan with its 5-hour and weekly limits. +The top bar has an agents icon. Before you've set up an agent, its panel walks you through signing in to Claude, Codex, or Grok, and the first one becomes your default agent. After that, the panel puts every subscription on one page: each agent and account with its plan and the percentage used of its 5-hour session and weekly limits (or the remaining prepaid balance), under a line that rotates through what you've used across all of them, like tokens this week and your most used model. Under _Make something_ it offers starter prompts for a new theme, plugin, or app that hand the task to your default agent. Apps are built the way Omarchy's own Hype, Monologue, and Omacut are, in C++ and Qt Quick, following the omarchy-app agent skill, and the compiler and Qt pieces they need come installed. Claude Code, Codex, Grok, and Fireworks are covered out of the box, as is Antigravity with its 5-hour and weekly limits, and Hermes shows its sessions and tokens, from every profile, once you've used it. Provider plans and keys you use through pi, omp, or opencode show up as well: the Z.ai GLM Coding Plan and Ollama Cloud with their 5-hour and weekly limits. Left-click the bar icon for the panel, right-click to launch your default agent. The usage records behind it are regenerated every 15 minutes by `omarchy agent usage-update`, and the panel can even merge usage from your other machines via a synced folder. See the README under `$OMARCHY_PATH/shell/plugins/agents/` for the full settings. diff --git a/shell/plugins/agents/README.md b/shell/plugins/agents/README.md index 5cbb1c849a0..d977d2cb0c6 100644 --- a/shell/plugins/agents/README.md +++ b/shell/plugins/agents/README.md @@ -121,9 +121,8 @@ Ollama Cloud reads `OLLAMA_API_KEY` first, then the key pi or omp signed in with in `~/.pi/agent/auth.json` or `~/.omp/agent/auth.json`. Ollama's usage endpoint reports the account's session and weekly usage as -0..1 fractions but not the reset times; the resets are the same for every -account (epoch-aligned 5-hour session windows, weekly windows ending Monday -00:00 UTC), so the collector computes them. The limits cover the whole +0..1 fractions but not the reset times, so those windows show no reset time +rather than a guessed one. The limits cover the whole account — every model the key can reach — while the local stats count pi and omp sessions on the `ollama-cloud` provider. To count only the Claude models Ollama Cloud serves, set a model prefix in diff --git a/test/shell.d/agent-usage-ollama-scanner-test.sh b/test/shell.d/agent-usage-ollama-scanner-test.sh index d3cbad832d4..feaa7ce0d22 100755 --- a/test/shell.d/agent-usage-ollama-scanner-test.sh +++ b/test/shell.d/agent-usage-ollama-scanner-test.sh @@ -13,7 +13,7 @@ trap 'rm -rf "$TEST_HOME"' EXIT no_key=$(HOME="$TEST_HOME" XDG_CONFIG_HOME="$TEST_HOME/.config" XDG_CACHE_HOME="$TEST_HOME/.cache" \ OLLAMA_API_KEY="" "$ROOT/bin/omarchy-agent-usage-ollama" --force) -[[ $(jq -r '.id + ":" + (.ready | tostring) + ":" + .usageStatusText' <<<"$no_key") == "ollama:false:Waiting for auth" ]] || +[[ $(jq -r '.id + ":" + (.ready | tostring) + ":" + .usageStatusText + ":" + .tierLabel + ":" + (.limits | length | tostring) + ":" + (.hasLocalStats | tostring)' <<<"$no_key") == "ollama:false:Waiting for auth::0:false" ]] || fail "Ollama collector prints a valid record without credentials" "$no_key" pass "Ollama collector prints a valid record without credentials" @@ -91,26 +91,10 @@ limits=$(read_limits '{ fail "Ollama collector keeps the endpoint's 0..1 fractions" "$limits" pass "Ollama collector reads the session and weekly fractions" -# Reset times are computed, not fetched: session windows are epoch-aligned -# 5-hour blocks and weekly windows end Monday 00:00 UTC. -resets=$(python3 - "$ROOT/bin/omarchy-agent-usage-ollama" <<'PY' -import datetime as dt -import importlib.machinery, importlib.util, sys - -loader = importlib.machinery.SourceFileLoader("collector", sys.argv[1]) -spec = importlib.util.spec_from_loader(loader.name, loader) -collector = importlib.util.module_from_spec(spec) -loader.exec_module(collector) - -session = dt.datetime.fromisoformat(collector.reset_iso(collector.SESSION_WINDOW_SECONDS)) -weekly = dt.datetime.fromisoformat(collector.reset_iso(collector.WEEKLY_WINDOW_SECONDS, collector.WEEKLY_EPOCH_OFFSET_SECONDS)) -print(f"{session.timestamp() % collector.SESSION_WINDOW_SECONDS == 0}:{weekly.strftime('%A')}:{weekly.hour}:{weekly.minute}") -PY -) - -[[ "$resets" == "True:Monday:0:0" ]] || - fail "Ollama collector computes epoch-aligned session and Monday-midnight weekly resets" "$resets" -pass "Ollama collector computes epoch-aligned session and Monday-midnight weekly resets" +# The endpoint names no reset times, and none are made up. +[[ $(jq -c '[.limits[].resetsAt]' <<<"$limits") == '["",""]' ]] || + fail "Ollama collector leaves reset times it isn't told empty" "$limits" +pass "Ollama collector leaves reset times it isn't told empty" # A payload with no usable windows explains itself instead of printing none. empty=$(read_limits '{"limits": {"session": {"usage": "unknown"}, "weekly": null}}') @@ -181,6 +165,8 @@ unreachable=$(collect_limits false "$stale" transport) fail "Ollama collector falls back to cache when the probe cannot connect" "$unreachable" [[ $(jq -r '.result.retryAdvised' <<<"$unreachable") == "true" ]] || fail "Ollama collector advises a retry after a transport failure" "$unreachable" +[[ $(jq -c '{live: .result.live, fetchedAtMs: .result.fetchedAtMs}' <<<"$unreachable") == '{"live":false,"fetchedAtMs":1}' ]] || + fail "Ollama collector marks kept limits stale with when they were fetched" "$unreachable" pass "Ollama collector falls back to cache when the probe cannot connect" # A rejected key says so instead of pretending the endpoint is down. @@ -203,6 +189,8 @@ forced=$(collect_limits true "$fresh" normal) fail "Ollama collector re-probes limits and plan on --force despite a fresh cache" "$forced" [[ $(jq -c '[.result.limits[].percent]' <<<"$forced") == "[0.44,0.11]" && $(jq -r '.result.tierLabel' <<<"$forced") == "Pro" ]] || fail "Ollama collector returns the forced probe's numbers and plan" "$forced" +[[ $(jq -r '.result.live and .result.fetchedAtMs == .cached.fetchedAtMs' <<<"$forced") == "true" ]] || + fail "Ollama collector marks a fresh probe live" "$forced" pass "Ollama collector re-probes on --force despite a fresh cache" # A probe that lands becomes the next run's fallback. From 53d4a6491a5962cf0194fccfbdf774b2a0ef2263 Mon Sep 17 00:00:00 2001 From: Michael D Date: Thu, 1 Oct 2026 22:02:18 -0400 Subject: [PATCH 09/19] Add MiniMax Token Plan usage collector The collector part of #9597: the Token Plan's 5-hour and weekly quota from MiniMax's API, with the key from MINIMAX_API_KEY, mmx, or opencode. --- bin/omarchy-agent-usage-minimax | 257 +++++++++++++++++++++++ shell/plugins/agents/README.md | 3 + shell/plugins/agents/manifest.json | 3 +- test/shell.d/agent-usage-minimax-test.sh | 117 +++++++++++ 4 files changed, 379 insertions(+), 1 deletion(-) create mode 100755 bin/omarchy-agent-usage-minimax create mode 100755 test/shell.d/agent-usage-minimax-test.sh diff --git a/bin/omarchy-agent-usage-minimax b/bin/omarchy-agent-usage-minimax new file mode 100755 index 00000000000..539abc2341b --- /dev/null +++ b/bin/omarchy-agent-usage-minimax @@ -0,0 +1,257 @@ +#!/usr/bin/python3 +# omarchy:summary=Print the MiniMax Token Plan usage record as JSON +# omarchy:args=[--force] [--limits-only] +# omarchy:hidden=true +"""Collect MiniMax Token Plan limits into a display-ready JSON record.""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +import urllib.error +import urllib.request +from datetime import datetime, timedelta, timezone +from pathlib import Path +from typing import Any + +AGENT_ID = "minimax" +AGENT_NAME = "MiniMax" +DEFAULT_BASE_URL = "https://api.minimax.io" +AUTH_HELP = "Set MINIMAX_API_KEY, run `mmx auth login`, or sign in to MiniMax in opencode." + + +class MiniMaxError(Exception): + pass + + +def empty_stats() -> dict[str, Any]: + return { + "todayPrompts": 0, + "todaySessions": 0, + "todayTotalTokens": 0, + "todayTokensByModel": {}, + "recentDays": [], + "totalPrompts": 0, + "totalSessions": 0, + "activeDays": 0, + "activeDates": [], + "modelUsage": {}, + } + + +def base_record(**overrides: Any) -> dict[str, Any]: + record: dict[str, Any] = { + "schemaVersion": 1, + "id": AGENT_ID, + "name": AGENT_NAME, + "updatedAt": datetime.now(timezone.utc).isoformat(), + "ready": False, + "hasLocalStats": False, + "hasPromptStats": False, + "scope": "account", + "tierLabel": "Token Plan", + "usageStatusText": "", + "authHelpText": "", + "limits": [], + } + record.update(empty_stats()) + record.update(overrides) + return record + + +def read_json(path: Path) -> Any: + try: + return json.loads(path.read_text(encoding="utf-8")) + except Exception: + return None + + +def mmx_config_dir() -> Path: + return Path(os.path.expandvars(os.path.expanduser(os.environ.get("MMX_CONFIG_DIR", "~/.mmx")))) + + +def mmx_config() -> dict[str, Any]: + config_dir = mmx_config_dir() + parsed = read_json(config_dir / "config.json") + if isinstance(parsed, dict): + return parsed + + # Older mmx releases documented YAML configuration. Only read the one scalar + # we need, avoiding a runtime YAML dependency in the base installation. + try: + for line in (config_dir / "config.yaml").read_text(encoding="utf-8").splitlines(): + key, separator, value = line.partition(":") + if separator and key.strip() == "api_key": + return {"api_key": value.strip().strip("'\"")} + except Exception: + pass + return {} + + +def config_key(config: dict[str, Any]) -> str: + value = config.get("api_key") or config.get("apiKey") + if isinstance(value, str) and value.strip(): + return value.strip() + + oauth = config.get("oauth") + if isinstance(oauth, dict): + token = oauth.get("access_token") + if isinstance(token, str) and token.strip(): + return token.strip() + return "" + + +def config_base_url(config: dict[str, Any]) -> str: + oauth = config.get("oauth") + if isinstance(oauth, dict): + resource_url = oauth.get("resource_url") + if isinstance(resource_url, str) and resource_url.strip(): + return resource_url.strip() + if oauth.get("region") == "cn": + return "https://api.minimaxi.com" + value = config.get("base_url") + if isinstance(value, str) and value.strip(): + return value.strip() + if config.get("region") == "cn" or os.environ.get("MINIMAX_REGION") == "cn": + return "https://api.minimaxi.com" + return DEFAULT_BASE_URL + + +def opencode_key() -> str: + data_home = Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local/share")) + parsed = read_json(data_home / "opencode" / "auth.json") + if not isinstance(parsed, dict): + return "" + for provider in ("minimax", "minimax-coding", "minimax-token-plan"): + entry = parsed.get(provider) + if isinstance(entry, dict): + value = entry.get("key") or entry.get("apiKey") or entry.get("token") + if isinstance(value, str) and value.strip(): + return value.strip() + return "" + + +def credentials() -> tuple[str, str]: + config = mmx_config() + api_key = os.environ.get("MINIMAX_API_KEY", "").strip() or config_key(config) or opencode_key() + base_url = os.environ.get("MINIMAX_BASE_URL", "").strip() or config_base_url(config) + return api_key, base_url + + +def number(value: Any) -> float | None: + try: + result = float(value) + return result if result == result else None + except Exception: + return None + + +def reset_at(entry: dict[str, Any], end_key: str, remains_key: str) -> str: + end_ms = number(entry.get(end_key)) + if end_ms and end_ms > 0: + return datetime.fromtimestamp(end_ms / 1000, timezone.utc).isoformat() + remains_ms = number(entry.get(remains_key)) + if remains_ms and remains_ms > 0: + return (datetime.now(timezone.utc) + timedelta(milliseconds=remains_ms)).isoformat() + return "" + + +def used_fraction(remaining: Any) -> float | None: + percent = number(remaining) + if percent is None: + return None + return max(0.0, min(1.0, 1.0 - percent / 100.0)) + + +def quota_limits(payload: dict[str, Any]) -> list[dict[str, Any]]: + rows = payload.get("model_remains") + if not isinstance(rows, list): + raise MiniMaxError("MiniMax returned no Token Plan quotas") + + # "general" is the language/agent quota. Other rows describe modalities + # whose daily limits do not map to the panel's coding-subscription meters. + general = next((row for row in rows if isinstance(row, dict) and row.get("model_name") == "general"), None) + if general is None: + general = next((row for row in rows if isinstance(row, dict)), None) + if not isinstance(general, dict): + raise MiniMaxError("MiniMax returned no usable Token Plan quota") + + limits = [] + interval = used_fraction(general.get("current_interval_remaining_percent")) + weekly = used_fraction(general.get("current_weekly_remaining_percent")) + if interval is not None: + limits.append({ + "label": "Session (5-hour)", + "percent": interval, + "resetsAt": reset_at(general, "end_time", "remains_time"), + }) + if weekly is not None: + limits.append({ + "label": "Weekly (7-day)", + "percent": weekly, + "resetsAt": reset_at(general, "weekly_end_time", "weekly_remains_time"), + }) + if not limits: + raise MiniMaxError("MiniMax returned no percentage-based Token Plan quotas") + return limits + + +def fetch_quota(api_key: str, base_url: str) -> dict[str, Any]: + endpoint = base_url.rstrip("/") + "/v1/token_plan/remains" + request = urllib.request.Request( + endpoint, + headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}, + ) + try: + with urllib.request.urlopen(request, timeout=12) as response: + payload = json.load(response) + except urllib.error.HTTPError as error: + if error.code in (401, 403): + raise MiniMaxError("MiniMax rejected the saved credentials; run `mmx auth login` or `mmx auth refresh`") from error + raise MiniMaxError(f"MiniMax API returned HTTP {error.code}") from error + except urllib.error.URLError as error: + raise MiniMaxError("Could not reach the MiniMax API") from error + except Exception as error: + raise MiniMaxError("MiniMax returned an invalid quota response") from error + + if not isinstance(payload, dict): + raise MiniMaxError("MiniMax returned an invalid quota response") + base_response = payload.get("base_resp") + if isinstance(base_response, dict) and number(base_response.get("status_code")) not in (None, 0): + message = str(base_response.get("status_msg") or "MiniMax rejected the quota request") + raise MiniMaxError(message) + return payload + + +def collect(api_key: str, base_url: str) -> dict[str, Any]: + if not api_key: + return base_record(usageStatusText="MiniMax unavailable", authHelpText=AUTH_HELP) + limits = quota_limits(fetch_quota(api_key, base_url)) + return base_record(ready=True, limits=limits) + + +def main() -> int: + parser = argparse.ArgumentParser(description="Print the MiniMax Token Plan usage record as JSON") + # The API is point-in-time, so both common collector flags intentionally + # take the same path. Accepting them keeps usage-update provider-agnostic. + parser.add_argument("--force", action="store_true") + parser.add_argument("--limits-only", action="store_true") + parser.add_argument("--base-url") + args = parser.parse_args() + + try: + api_key, saved_base_url = credentials() + record = collect(api_key, args.base_url or saved_base_url) + except MiniMaxError as error: + record = base_record(usageStatusText="MiniMax limits unavailable", authHelpText=str(error)) + except Exception as error: + record = base_record(usageStatusText="MiniMax limits unavailable", authHelpText="MiniMax usage scan failed") + print(f"omarchy-agent-usage-minimax: {type(error).__name__}", file=sys.stderr) + print(json.dumps(record, separators=(",", ":"), sort_keys=True)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/shell/plugins/agents/README.md b/shell/plugins/agents/README.md index d977d2cb0c6..9ff78400d50 100644 --- a/shell/plugins/agents/README.md +++ b/shell/plugins/agents/README.md @@ -78,6 +78,7 @@ light surfaces — and the bar glyph stands in when there is none. | `agy` | Google's Cloud Code API (`loadCodeAssist` for the plan, `retrieveUserQuotaSummary` for the 5-hour and weekly windows) | `history.jsonl`, `conversation_summaries.db`, and each conversation's `transcript.jsonl` under `~/.gemini/antigravity-cli/` (`AGY_DIR`) | | `zai` | Z.ai's coding-plan monitor endpoint (5-hour session + 7-day weekly) | pi/omp sessions on the Z.ai provider (`glm-*` models) | | `ollama` | Ollama Cloud's usage endpoint (5-hour session + 7-day weekly) | pi and omp sessions on the `ollama-cloud` provider | +| `minimax` | MiniMax Token Plan API (5-hour session + 7-day weekly) | Not available from the quota API | When `~/.local/state/omarchy/agents/accounts/.json` registers more than one account, the `claude`, `codex`, and `grok` records @@ -138,6 +139,8 @@ Ollama Cloud serves, set a model prefix in `modelPrefix` restricts the local stats to models whose name starts with the prefix (empty means every model); `apiKey` overrides the key lookup order. +MiniMax reads `MINIMAX_API_KEY` first, then the API key or OAuth session saved by `mmx auth login`, then the key opencode stores when MiniMax is signed in there. It uses the region and resource URL saved by `mmx`, or the global `https://api.minimax.io` endpoint by default; `MINIMAX_BASE_URL` overrides both. Run `mmx auth refresh` if the panel reports that a saved OAuth session expired. + ### Fireworks balance The collector first asks the account's `:getBalance` endpoint for the real diff --git a/shell/plugins/agents/manifest.json b/shell/plugins/agents/manifest.json index 07e4880ca86..bc42f51feef 100644 --- a/shell/plugins/agents/manifest.json +++ b/shell/plugins/agents/manifest.json @@ -26,7 +26,8 @@ "hermes": { "enabled": true }, "agy": { "enabled": true }, "zai": { "enabled": true }, - "ollama": { "enabled": true } + "ollama": { "enabled": true }, + "minimax": { "enabled": true } }, "refreshIntervalSec": 900, "syncMode": "Off", diff --git a/test/shell.d/agent-usage-minimax-test.sh b/test/shell.d/agent-usage-minimax-test.sh new file mode 100755 index 00000000000..1f4805966ef --- /dev/null +++ b/test/shell.d/agent-usage-minimax-test.sh @@ -0,0 +1,117 @@ +#!/bin/bash + +source "$(dirname "$0")/base-test.sh" + +require_command jq +require_command python3 + +TEST_HOME=$(mktemp -d) +trap 'rm -rf "$TEST_HOME"' EXIT + +no_key=$(HOME="$TEST_HOME" XDG_DATA_HOME="$TEST_HOME/.local/share" MINIMAX_API_KEY="" \ + MMX_CONFIG_DIR="$TEST_HOME/missing" "$ROOT/bin/omarchy-agent-usage-minimax") + +[[ $(jq -r '.id + ":" + (.ready | tostring) + ":" + .tierLabel' <<<"$no_key") == "minimax:false:Token Plan" ]] || + fail "MiniMax collector prints a valid record without credentials" "$no_key" +pass "MiniMax collector prints a valid record without credentials" + +result=$(python3 - "$ROOT/bin/omarchy-agent-usage-minimax" "$TEST_HOME" <<'PY' +import importlib.machinery +import importlib.util +import json +import os +import sys +from pathlib import Path + +collector_path = sys.argv[1] +home = Path(sys.argv[2]) +os.environ["HOME"] = str(home) +os.environ["MMX_CONFIG_DIR"] = str(home / ".mmx") +os.environ["XDG_DATA_HOME"] = str(home / ".local/share") + +loader = importlib.machinery.SourceFileLoader("minimax_collector", collector_path) +spec = importlib.util.spec_from_loader(loader.name, loader) +scanner = importlib.util.module_from_spec(spec) +loader.exec_module(scanner) + +payload = { + "model_remains": [ + { + "model_name": "general", + "end_time": 1782399600000, + "current_interval_remaining_percent": 10, + "current_interval_status": 1, + "weekly_end_time": 1782691200000, + "current_weekly_remaining_percent": 69, + "current_weekly_status": 1, + }, + { + "model_name": "video", + "current_interval_total_count": 0, + "current_interval_remaining_percent": 100, + "current_interval_status": 3, + }, + ], + "base_resp": {"status_code": 0, "status_msg": "success"}, +} + +limits = scanner.quota_limits(payload) + +config_dir = home / ".mmx" +config_dir.mkdir(parents=True) +(config_dir / "config.json").write_text(json.dumps({"api_key": "mmx_saved"})) +saved_key, saved_base_url = scanner.credentials() + +(config_dir / "config.json").write_text(json.dumps({ + "oauth": { + "access_token": "mmx_oauth", + "refresh_token": "refresh", + "expires_at": "2099-01-01T00:00:00Z", + "region": "cn", + } +})) +oauth_key, oauth_base_url = scanner.credentials() + +(config_dir / "config.json").unlink() +opencode = home / ".local/share/opencode/auth.json" +opencode.parent.mkdir(parents=True) +opencode.write_text(json.dumps({"minimax": {"type": "api", "key": "mmx_opencode"}})) +opencode_key, opencode_base_url = scanner.credentials() + +class WorkingClient: + pass + +scanner.fetch_quota = lambda api_key, base_url: payload +record = scanner.collect("mmx_test", "https://example.invalid") + +print(json.dumps({ + "limits": limits, + "savedKey": saved_key, + "savedBaseUrl": saved_base_url, + "oauthKey": oauth_key, + "oauthBaseUrl": oauth_base_url, + "opencodeKey": opencode_key, + "record": record, +})) +PY +) + +[[ $(jq -c '.limits | map({label, percent})' <<<"$result") == '[{"label":"Session (5-hour)","percent":0.9},{"label":"Weekly (7-day)","percent":0.31000000000000005}]' ]] || + fail "MiniMax remaining percentages become Omarchy used fractions" "$result" +pass "MiniMax remaining percentages become Omarchy used fractions" + +[[ $(jq -r '.limits[0].resetsAt + ":" + .limits[1].resetsAt' <<<"$result") == "2026-06-25T15:00:00+00:00:2026-06-29T00:00:00+00:00" ]] || + fail "MiniMax quota reset timestamps are normalized" "$result" +pass "MiniMax quota reset timestamps are normalized" + +[[ $(jq -r '.savedKey + ":" + .opencodeKey' <<<"$result") == "mmx_saved:mmx_opencode" ]] || + fail "MiniMax credentials use mmx and opencode logins" "$result" +pass "MiniMax credentials use mmx and opencode logins" + +[[ $(jq -r '.oauthKey + ":" + .oauthBaseUrl' <<<"$result") == "mmx_oauth:https://api.minimaxi.com" ]] || + fail "MiniMax credentials use the OAuth session and region saved by mmx" "$result" +pass "MiniMax credentials use the OAuth session and region saved by mmx" + +[[ $(jq -c '.record | {id, ready, hasLocalStats, hasPromptStats, scope, tierLabel}' <<<"$result") == '{"id":"minimax","ready":true,"hasLocalStats":false,"hasPromptStats":false,"scope":"account","tierLabel":"Token Plan"}' ]] || + fail "MiniMax collector prints the display-ready record contract" "$result" +pass "MiniMax collector prints the display-ready record contract" From 1df1ff982a142e2cde38cefb09bc6bb99d7feabc Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Thu, 1 Oct 2026 22:03:07 -0400 Subject: [PATCH 10/19] Keep stale MiniMax limits and skip MiniMax without a key A failed check now keeps the last good limits, marked with limitsStale and limitsFetchedAt, instead of dropping them. Without a key the record has no plan, so the panel skips it. Help text no longer points at the mmx CLI, which Omarchy doesn't install, and the manual mentions MiniMax. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-minimax | 63 +++++++++++++++++++++--- manual/17-ai.md | 2 +- shell/plugins/agents/README.md | 6 ++- test/shell.d/agent-usage-minimax-test.sh | 24 +++++++-- 4 files changed, 80 insertions(+), 15 deletions(-) diff --git a/bin/omarchy-agent-usage-minimax b/bin/omarchy-agent-usage-minimax index 539abc2341b..71117f311c5 100755 --- a/bin/omarchy-agent-usage-minimax +++ b/bin/omarchy-agent-usage-minimax @@ -2,7 +2,11 @@ # omarchy:summary=Print the MiniMax Token Plan usage record as JSON # omarchy:args=[--force] [--limits-only] # omarchy:hidden=true -"""Collect MiniMax Token Plan limits into a display-ready JSON record.""" +"""Collect MiniMax Token Plan limits into a display-ready JSON record. + +The last good limits are cached, so a failed check keeps showing them, marked +stale, rather than dropping the section. +""" from __future__ import annotations @@ -10,6 +14,8 @@ import argparse import json import os import sys +import tempfile +import time import urllib.error import urllib.request from datetime import datetime, timedelta, timezone @@ -19,7 +25,8 @@ from typing import Any AGENT_ID = "minimax" AGENT_NAME = "MiniMax" DEFAULT_BASE_URL = "https://api.minimax.io" -AUTH_HELP = "Set MINIMAX_API_KEY, run `mmx auth login`, or sign in to MiniMax in opencode." +AUTH_HELP = "Set MINIMAX_API_KEY, or sign in to MiniMax in opencode." +TIER_LABEL = "Token Plan" class MiniMaxError(Exception): @@ -51,10 +58,12 @@ def base_record(**overrides: Any) -> dict[str, Any]: "hasLocalStats": False, "hasPromptStats": False, "scope": "account", - "tierLabel": "Token Plan", + "tierLabel": "", "usageStatusText": "", "authHelpText": "", "limits": [], + "limitsStale": False, + "limitsFetchedAt": 0, } record.update(empty_stats()) record.update(overrides) @@ -68,6 +77,24 @@ def read_json(path: Path) -> Any: return None +# A temporary file of its own per write, so overlapping runs can't rename +# each other's away. +def write_json(path: Path, payload: Any) -> None: + 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 limits_cache() -> Path: + return Path(os.environ.get("XDG_CACHE_HOME") or (Path.home() / ".cache")) / "omarchy" / "agent-usage" / "minimax-limits.json" + + def mmx_config_dir() -> Path: return Path(os.path.expandvars(os.path.expanduser(os.environ.get("MMX_CONFIG_DIR", "~/.mmx")))) @@ -209,7 +236,7 @@ def fetch_quota(api_key: str, base_url: str) -> dict[str, Any]: payload = json.load(response) except urllib.error.HTTPError as error: if error.code in (401, 403): - raise MiniMaxError("MiniMax rejected the saved credentials; run `mmx auth login` or `mmx auth refresh`") from error + raise MiniMaxError("MiniMax rejected the saved credentials") from error raise MiniMaxError(f"MiniMax API returned HTTP {error.code}") from error except urllib.error.URLError as error: raise MiniMaxError("Could not reach the MiniMax API") from error @@ -225,11 +252,33 @@ def fetch_quota(api_key: str, base_url: str) -> dict[str, Any]: return payload +# Nobody with a MiniMax key here: an empty record, which the panel skips. +# With one, a failed check keeps the last good limits, marked stale. def collect(api_key: str, base_url: str) -> dict[str, Any]: if not api_key: return base_record(usageStatusText="MiniMax unavailable", authHelpText=AUTH_HELP) - limits = quota_limits(fetch_quota(api_key, base_url)) - return base_record(ready=True, limits=limits) + + record = base_record(ready=True, tierLabel=TIER_LABEL) + try: + limits = quota_limits(fetch_quota(api_key, base_url)) + except MiniMaxError as error: + cached = read_json(limits_cache()) or {} + record.update( + limits=cached.get("limits") or [], + limitsStale=True, + limitsFetchedAt=int(cached.get("fetchedAtMs") or 0), + usageStatusText="MiniMax limits unavailable", + authHelpText=str(error), + ) + return record + + fetched_at = round(time.time() * 1000) + record.update(limits=limits, limitsFetchedAt=fetched_at) + try: + write_json(limits_cache(), {"fetchedAtMs": fetched_at, "limits": limits}) + except OSError: + pass + return record def main() -> int: @@ -244,8 +293,6 @@ def main() -> int: try: api_key, saved_base_url = credentials() record = collect(api_key, args.base_url or saved_base_url) - except MiniMaxError as error: - record = base_record(usageStatusText="MiniMax limits unavailable", authHelpText=str(error)) except Exception as error: record = base_record(usageStatusText="MiniMax limits unavailable", authHelpText="MiniMax usage scan failed") print(f"omarchy-agent-usage-minimax: {type(error).__name__}", file=sys.stderr) diff --git a/manual/17-ai.md b/manual/17-ai.md index b1f1ed6d69e..aa45804c9de 100644 --- a/manual/17-ai.md +++ b/manual/17-ai.md @@ -38,7 +38,7 @@ There are terminal shortcuts too: `a` runs the default agent inline in the curre ### The agents panel -The top bar has an agents icon. Before you've set up an agent, its panel walks you through signing in to Claude, Codex, or Grok, and the first one becomes your default agent. After that, the panel puts every subscription on one page: each agent and account with its plan and the percentage used of its 5-hour session and weekly limits (or the remaining prepaid balance), under a line that rotates through what you've used across all of them, like tokens this week and your most used model. Under _Make something_ it offers starter prompts for a new theme, plugin, or app that hand the task to your default agent. Apps are built the way Omarchy's own Hype, Monologue, and Omacut are, in C++ and Qt Quick, following the omarchy-app agent skill, and the compiler and Qt pieces they need come installed. Claude Code, Codex, Grok, and Fireworks are covered out of the box, as is Antigravity with its 5-hour and weekly limits, and Hermes shows its sessions and tokens, from every profile, once you've used it. Provider plans and keys you use through pi, omp, or opencode show up as well: the Z.ai GLM Coding Plan and Ollama Cloud with their 5-hour and weekly limits. +The top bar has an agents icon. Before you've set up an agent, its panel walks you through signing in to Claude, Codex, or Grok, and the first one becomes your default agent. After that, the panel puts every subscription on one page: each agent and account with its plan and the percentage used of its 5-hour session and weekly limits (or the remaining prepaid balance), under a line that rotates through what you've used across all of them, like tokens this week and your most used model. Under _Make something_ it offers starter prompts for a new theme, plugin, or app that hand the task to your default agent. Apps are built the way Omarchy's own Hype, Monologue, and Omacut are, in C++ and Qt Quick, following the omarchy-app agent skill, and the compiler and Qt pieces they need come installed. Claude Code, Codex, Grok, and Fireworks are covered out of the box, as is Antigravity with its 5-hour and weekly limits, and Hermes shows its sessions and tokens, from every profile, once you've used it. Provider plans and keys you use through pi, omp, or opencode show up as well: the Z.ai GLM Coding Plan, Ollama Cloud, and the MiniMax Token Plan with their 5-hour and weekly limits. Left-click the bar icon for the panel, right-click to launch your default agent. The usage records behind it are regenerated every 15 minutes by `omarchy agent usage-update`, and the panel can even merge usage from your other machines via a synced folder. See the README under `$OMARCHY_PATH/shell/plugins/agents/` for the full settings. diff --git a/shell/plugins/agents/README.md b/shell/plugins/agents/README.md index 9ff78400d50..9daa59da97c 100644 --- a/shell/plugins/agents/README.md +++ b/shell/plugins/agents/README.md @@ -139,7 +139,11 @@ Ollama Cloud serves, set a model prefix in `modelPrefix` restricts the local stats to models whose name starts with the prefix (empty means every model); `apiKey` overrides the key lookup order. -MiniMax reads `MINIMAX_API_KEY` first, then the API key or OAuth session saved by `mmx auth login`, then the key opencode stores when MiniMax is signed in there. It uses the region and resource URL saved by `mmx`, or the global `https://api.minimax.io` endpoint by default; `MINIMAX_BASE_URL` overrides both. Run `mmx auth refresh` if the panel reports that a saved OAuth session expired. +MiniMax reads `MINIMAX_API_KEY` first, then the API key or OAuth session +saved by MiniMax's own `mmx` CLI if you use it, then the key opencode stores +when MiniMax is signed in there. It uses the region and resource URL saved by +`mmx`, or the global `https://api.minimax.io` endpoint by default; +`MINIMAX_BASE_URL` overrides both. ### Fireworks balance diff --git a/test/shell.d/agent-usage-minimax-test.sh b/test/shell.d/agent-usage-minimax-test.sh index 1f4805966ef..a30acb560d8 100755 --- a/test/shell.d/agent-usage-minimax-test.sh +++ b/test/shell.d/agent-usage-minimax-test.sh @@ -8,12 +8,12 @@ require_command python3 TEST_HOME=$(mktemp -d) trap 'rm -rf "$TEST_HOME"' EXIT -no_key=$(HOME="$TEST_HOME" XDG_DATA_HOME="$TEST_HOME/.local/share" MINIMAX_API_KEY="" \ +no_key=$(HOME="$TEST_HOME" XDG_DATA_HOME="$TEST_HOME/.local/share" XDG_CACHE_HOME="$TEST_HOME/.cache" MINIMAX_API_KEY="" \ MMX_CONFIG_DIR="$TEST_HOME/missing" "$ROOT/bin/omarchy-agent-usage-minimax") -[[ $(jq -r '.id + ":" + (.ready | tostring) + ":" + .tierLabel' <<<"$no_key") == "minimax:false:Token Plan" ]] || - fail "MiniMax collector prints a valid record without credentials" "$no_key" -pass "MiniMax collector prints a valid record without credentials" +[[ $(jq -c '{id, ready, tierLabel, limits}' <<<"$no_key") == '{"id":"minimax","ready":false,"tierLabel":"","limits":[]}' ]] || + fail "MiniMax collector prints a record the panel skips without credentials" "$no_key" +pass "MiniMax collector prints a record the panel skips without credentials" result=$(python3 - "$ROOT/bin/omarchy-agent-usage-minimax" "$TEST_HOME" <<'PY' import importlib.machinery @@ -28,6 +28,7 @@ home = Path(sys.argv[2]) os.environ["HOME"] = str(home) os.environ["MMX_CONFIG_DIR"] = str(home / ".mmx") os.environ["XDG_DATA_HOME"] = str(home / ".local/share") +os.environ["XDG_CACHE_HOME"] = str(home / ".cache") loader = importlib.machinery.SourceFileLoader("minimax_collector", collector_path) spec = importlib.util.spec_from_loader(loader.name, loader) @@ -84,6 +85,12 @@ class WorkingClient: scanner.fetch_quota = lambda api_key, base_url: payload record = scanner.collect("mmx_test", "https://example.invalid") +def unreachable(api_key, base_url): + raise scanner.MiniMaxError("Could not reach the MiniMax API") + +scanner.fetch_quota = unreachable +stale = scanner.collect("mmx_test", "https://example.invalid") + print(json.dumps({ "limits": limits, "savedKey": saved_key, @@ -92,6 +99,7 @@ print(json.dumps({ "oauthBaseUrl": oauth_base_url, "opencodeKey": opencode_key, "record": record, + "stale": stale, })) PY ) @@ -112,6 +120,12 @@ pass "MiniMax credentials use mmx and opencode logins" fail "MiniMax credentials use the OAuth session and region saved by mmx" "$result" pass "MiniMax credentials use the OAuth session and region saved by mmx" -[[ $(jq -c '.record | {id, ready, hasLocalStats, hasPromptStats, scope, tierLabel}' <<<"$result") == '{"id":"minimax","ready":true,"hasLocalStats":false,"hasPromptStats":false,"scope":"account","tierLabel":"Token Plan"}' ]] || +[[ $(jq -c '.record | {id, ready, hasLocalStats, hasPromptStats, scope, tierLabel, limitsStale}' <<<"$result") == '{"id":"minimax","ready":true,"hasLocalStats":false,"hasPromptStats":false,"scope":"account","tierLabel":"Token Plan","limitsStale":false}' ]] || fail "MiniMax collector prints the display-ready record contract" "$result" pass "MiniMax collector prints the display-ready record contract" + +[[ $(jq -c '.stale | {ready, tierLabel, limitsStale, percent: .limits[0].percent, usageStatusText, authHelpText}' <<<"$result") == '{"ready":true,"tierLabel":"Token Plan","limitsStale":true,"percent":0.9,"usageStatusText":"MiniMax limits unavailable","authHelpText":"Could not reach the MiniMax API"}' ]] || + fail "MiniMax collector keeps the last limits, marked stale, after a failed check" "$result" +[[ $(jq -r '.stale.limitsFetchedAt == .record.limitsFetchedAt and .record.limitsFetchedAt > 0' <<<"$result") == "true" ]] || + fail "stale MiniMax limits keep when they were fetched" "$result" +pass "MiniMax collector keeps the last limits, marked stale, after a failed check" From c5e5bffe2dbeced52aae9efe7c133b62dddce9f7 Mon Sep 17 00:00:00 2001 From: Azonnali Date: Thu, 1 Oct 2026 22:03:37 -0400 Subject: [PATCH 11/19] Add OpenRouter usage collector to the agents panel --- bin/omarchy-agent-usage-openrouter | 563 ++++++++++++++++++ shell/plugins/agents/README.md | 33 + shell/plugins/agents/manifest.json | 3 +- .../agent-usage-openrouter-scanner-test.sh | 136 +++++ 4 files changed, 734 insertions(+), 1 deletion(-) create mode 100755 bin/omarchy-agent-usage-openrouter create mode 100755 test/shell.d/agent-usage-openrouter-scanner-test.sh diff --git a/bin/omarchy-agent-usage-openrouter b/bin/omarchy-agent-usage-openrouter new file mode 100755 index 00000000000..bd6b908e330 --- /dev/null +++ b/bin/omarchy-agent-usage-openrouter @@ -0,0 +1,563 @@ +#!/usr/bin/python3 +# omarchy:summary=Print the OpenRouter usage record as JSON +# omarchy:args=[--force] [--limits-only] +# omarchy:hidden=true +"""Collect OpenRouter usage into one display-ready JSON record. + +Tier 1 (any API key): balance and key usage come from GET /api/v1/credits +and GET /api/v1/key; token stats come from the local opencode SQLite DB +(providerID == "openrouter", role == "assistant"), read-only. + +Tier 2 (management key): when OPENROUTER_MANAGEMENT_KEY (or the management +key in the config file) is set, the last-30-days token history comes from +GET /api/v1/activity instead of the local DB -- account-global numbers that +cover every machine and client, grouped by day and model. The /key response +carries is_management_key, so a key that cannot do this fails visibly as +403 and the collector falls back to tier 1 automatically. + +Refs: omacom/omarchy#8270 (pi on OpenRouter invisible in panel), +omacom/omarchy#10330 (dedicated pi collector proposal). +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import sqlite3 +import sys +import tempfile +import urllib.request +from datetime import datetime, timedelta, timezone +from pathlib import Path +from typing import Any + +AGENT_ID = "openrouter" +AGENT_NAME = "OpenRouter" +AUTH_HELP = "Set OPENROUTER_API_KEY, or sign in to OpenRouter in opencode. For the 30-day history, add a management key as OPENROUTER_MANAGEMENT_KEY (openrouter.ai/settings/management-keys)." +API_BASE = "https://openrouter.ai/api/v1" + +# A scan this recent is only reused to dedup concurrent collector runs; +# every periodic refresh lands a real rescan. --limits-only promises only +# fresh balance/limits, so it may reuse the local scan for up to 15 minutes. +SCAN_REUSE_SECONDS = 20 +LIMITS_ONLY_REUSE_SECONDS = 900 + + +class OpenRouterError(Exception): + pass + + +def number(value: Any) -> int: + try: + return max(0, int(value or 0)) + except (TypeError, ValueError): + return 0 + + +def local_day(value) -> str: + now = datetime.now() + if value is None: + return now.strftime("%Y-%m-%d") + if isinstance(value, (int, float)): + ts = value / 1000 if value > 10_000_000_000 else value + return datetime.fromtimestamp(ts).astimezone().strftime("%Y-%m-%d") + try: + text = str(value) + dt = datetime.fromisoformat(text.replace("Z", "+00:00")) + if dt.tzinfo is not None: + dt = dt.astimezone() + return dt.strftime("%Y-%m-%d") + except Exception: + return now.strftime("%Y-%m-%d") + + +def model_name(raw: Any) -> str: + value = str(raw or "openrouter").rstrip("/").split("/")[-1] + return value or "openrouter" + + +def empty_bucket() -> dict[str, int]: + return { + "inputTokens": 0, + "outputTokens": 0, + "cacheReadInputTokens": 0, + "cacheCreationInputTokens": 0, + } + + +def empty_stats() -> dict[str, Any]: + return { + "todayPrompts": 0, + "todaySessions": 0, + "todayTotalTokens": 0, + "todayTokensByModel": {}, + "recentDays": [], + "totalPrompts": 0, + "totalSessions": 0, + "activeDays": 0, + "activeDates": [], + "modelUsage": {}, + } + + +# ---------------------------------------------------------------- auth + +def config_path() -> Path: + config_home = Path(os.environ.get("XDG_CONFIG_HOME") or (Path.home() / ".config")) + return config_home / "omarchy" / "agents" / "openrouter.json" + + +def read_config() -> dict[str, Any]: + try: + parsed = json.loads(config_path().read_text()) + return parsed if isinstance(parsed, dict) else {} + except (OSError, json.JSONDecodeError): + return {} + +def opencode_auth_path() -> Path: + data_home = Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) + return data_home / "opencode" / "auth.json" + + +def read_opencode_key(path: Path) -> str: + try: + parsed = json.loads(path.read_text()) + except (OSError, json.JSONDecodeError): + return "" + entry = parsed.get("openrouter") if isinstance(parsed, dict) else None + if not isinstance(entry, dict): + return "" + return str(entry.get("key") or "").strip() + + +def api_key(config: dict[str, Any]) -> str: + # Explicit key or opencode login wins; the config file is the last + # resort for headless/refresh-timer runs without either. + return ( + str(os.environ.get("OPENROUTER_API_KEY", "")).strip() + or read_opencode_key(opencode_auth_path()) + or str(config.get("apiKey") or "").strip() + ) + + +def management_key(config: dict[str, Any]) -> str: + # A management key unlocks /activity (30-day account history). A normal + # inference key answers /key + /credits only and fails /activity with + # 403 -- never send it there speculatively, just skip tier 2. + return ( + str(os.environ.get("OPENROUTER_MANAGEMENT_KEY", "")).strip() + or str(config.get("managementKey") or "").strip() + ) + + +# ---------------------------------------------------------------- API + +class OpenRouterTransportError(OpenRouterError): + pass + + +def api_get(path: str, key: str) -> dict: + req = urllib.request.Request( + API_BASE + path, headers={"Authorization": "Bearer " + key, "Accept": "application/json"} + ) + try: + with urllib.request.urlopen(req, timeout=15) as resp: + payload = json.load(resp) + return payload if isinstance(payload, dict) else {} + except urllib.request.HTTPError as error: + if error.code in (401, 403): + # 403 from /activity with an inference key is the expected + # tier-1 signal ("Only management keys can perform this + # operation"), not a broken setup -- callers decide. + raise OpenRouterError(f"HTTP_{error.code}") + raise OpenRouterError(f"OpenRouter API returned HTTP {error.code}") + except urllib.request.URLError as error: + raise OpenRouterTransportError("Could not reach the OpenRouter API") from error + except (json.JSONDecodeError, ValueError) as error: + raise OpenRouterError("OpenRouter returned an invalid response") from error + + +def summarize_activity(items: list[Any], today: datetime) -> dict[str, Any]: + """Fold /activity rows into the panel's stats dict. + + One row per (date, endpoint): prompt + reasoning tokens count as input + (they are generated), completion as output. The billing API reports + tokens, never prompt or session counts, so hasPromptStats stays False + and the panel keeps those numbers out of today's tooltip. + """ + today_date = today.strftime("%Y-%m-%d") + recent_dates = [(today - timedelta(days=offset)).strftime("%Y-%m-%d") for offset in range(6, -1, -1)] + recent = {day: 0 for day in recent_dates} + today_by_model: dict[str, int] = {} + model_usage: dict[str, dict[str, int]] = {} + active_dates: set[str] = set() + + for raw in items: + if not isinstance(raw, dict): + continue + day = str(raw.get("date") or "") + model = model_name(str(raw.get("model_permaslug") or raw.get("model") or "")) + prompt = number(raw.get("prompt_tokens")) + reasoning = number(raw.get("reasoning_tokens")) + output = number(raw.get("completion_tokens")) + total = prompt + reasoning + output + if total <= 0: + continue + bucket = model_usage.setdefault(model, empty_bucket()) + bucket["inputTokens"] += prompt + bucket["outputTokens"] += output + reasoning + if day: + active_dates.add(day) + if day in recent: + recent[day] += total + if day == today_date: + today_by_model[model] = today_by_model.get(model, 0) + total + + return { + "todayTokensByModel": today_by_model, + "todayTotalTokens": sum(today_by_model.values()), + "recentDays": [{"date": day, "messageCount": recent[day]} for day in recent_dates], + "activeDays": len(active_dates), + "activeDates": sorted(active_dates), + "modelUsage": model_usage, + } + + +def activity_stats(key: str, today: datetime) -> dict[str, Any] | None: + """Tier 2: 30-day account history via management key, or None. + + Returns None when there is no management key (plain tier-1 run) or the + key is rejected -- the caller then falls back to the local DB scan. + """ + if not key: + return None + try: + payload = api_get("/activity", key) + except OpenRouterError: + return None + items = payload.get("data") + if not isinstance(items, list): + return None + stats = empty_stats() + stats.update(summarize_activity(items, today)) + return stats + + +# ---------------------------------------------------------------- local stats (opencode DB) + +class LocalScan: + """Aggregated opencode usage. Module-level in the original prototype; + an instance keeps the scanner testable without global state.""" + + def __init__(self, today: datetime): + self.today = today.strftime("%Y-%m-%d") + self.recent_dates = [(today - timedelta(days=o)).strftime("%Y-%m-%d") for o in range(6, -1, -1)] + self.recent = {day: {"date": day, "messageCount": 0} for day in self.recent_dates} + self.today_tokens_by_model: dict[str, int] = {} + self.model_usage: dict[str, dict[str, int]] = {} + self.today_sessions: set[str] = set() + self.active_days: set[str] = set() + self.today_prompts = 0 + self.today_total_tokens = 0 + self.total_prompts = 0 + self.total_sessions: set[str] = set() + + def add_usage(self, day, session_key, model, in_tok, out_tok, cache_read, cache_write): + total = in_tok + out_tok + cache_read + cache_write + self.total_prompts += 1 + self.total_sessions.add(session_key) + self.active_days.add(day) + bucket = self.model_usage.setdefault(model, empty_bucket()) + bucket["inputTokens"] += in_tok + bucket["outputTokens"] += out_tok + bucket["cacheReadInputTokens"] += cache_read + bucket["cacheCreationInputTokens"] += cache_write + if day in self.recent: + self.recent[day]["messageCount"] += total + if day == self.today: + self.today_prompts += 1 + self.today_sessions.add(session_key) + self.today_total_tokens += total + self.today_tokens_by_model[model] = self.today_tokens_by_model.get(model, 0) + total + + def snapshot(self) -> dict: + return { + "todayPrompts": self.today_prompts, + "todaySessions": len(self.today_sessions), + "todayTotalTokens": self.today_total_tokens, + "todayTokensByModel": self.today_tokens_by_model, + "recentDays": [self.recent[d] for d in self.recent_dates], + "totalPrompts": self.total_prompts, + "totalSessions": len(self.total_sessions), + "activeDays": len(self.active_days), + "activeDates": sorted(self.active_days), + "modelUsage": self.model_usage, + } + + +def opencode_db_path() -> Path: + return Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) / "opencode" / "opencode.db" + + +def scan_opencode(scan: LocalScan, db: Path | None = None) -> bool: + # Returns whether the scan ran to completion: a scan cut short by a + # database error still contributes what it read, but must not be cached + # as if it were the whole story. + db = db or opencode_db_path() + if not db.is_file(): + return True + try: + conn = sqlite3.connect(db.resolve().as_uri() + "?mode=ro", uri=True, timeout=2) + except sqlite3.Error: + return False + try: + conn.execute("PRAGMA query_only = ON") + # The LIKE gates skip rows whose JSON cannot contain the two + # key/value pairs, avoiding the JSON parse of huge message blobs; + # json_extract is the authority for the rows that reach Python. + rows = conn.execute( + "SELECT session_id, data FROM message" + " WHERE data LIKE '%\"role\"%:%\"assistant\"%'" + " AND data LIKE '%\"providerID\"%:%\"openrouter\"%'" + " AND CASE WHEN json_valid(data) THEN json_extract(data, '$.role') END = 'assistant'" + " AND CASE WHEN json_valid(data) THEN json_extract(data, '$.providerID') END = 'openrouter'" + ) + for session_id, raw in rows: + try: + entry = json.loads(raw) + if not isinstance(entry, dict) or entry.get("role") != "assistant": + continue + if str(entry.get("providerID") or "") != "openrouter": + continue + tokens = entry.get("tokens") or {} + cache = tokens.get("cache") or {} + in_tok = number(tokens.get("input")) + # opencode keeps thinking tokens out of output; both are generated. + out_tok = number(tokens.get("output")) + number(tokens.get("reasoning")) + cache_read = number(cache.get("read")) + cache_write = number(cache.get("write")) + if not (in_tok or out_tok or cache_read or cache_write): + continue + day = local_day((entry.get("time") or {}).get("created")) + model = model_name(str(entry.get("modelID") or "")) + except Exception: + continue + scan.add_usage(day, "opencode:" + str(session_id), model, in_tok, out_tok, cache_read, cache_write) + except sqlite3.Error: + return False + finally: + conn.close() + return True + + +def cache_root() -> Path: + root = Path(os.environ.get("XDG_CACHE_HOME") or (Path.home() / ".cache")) / "omarchy" / "agent-usage" + root.mkdir(parents=True, exist_ok=True) + return root + + +def scan_cache_paths() -> tuple[Path, Path]: + db = opencode_db_path() + digest = hashlib.sha1(str(db).encode("utf-8")).hexdigest()[:16] + root = cache_root() + return root / f"openrouter-scan-{digest}.json", root / f"openrouter-scan-{digest}.lock" + + +def read_fresh_json(path: Path, max_age_seconds: float) -> dict | None: + if max_age_seconds <= 0 or not path.exists(): + return None + try: + age = datetime.now().timestamp() - path.stat().st_mtime + if 0 <= age <= max_age_seconds: + parsed = json.loads(path.read_text(encoding="utf-8")) + return parsed if isinstance(parsed, dict) else None + except Exception: + return None + return None + + +def write_json(path: Path, payload: dict) -> None: + 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=(",", ":")) + "\n") + tmp.chmod(0o644) + tmp.replace(path) + except BaseException: + tmp.unlink(missing_ok=True) + raise + + +def cached_local_stats(today: datetime, max_age: float) -> tuple[dict, bool]: + """Local stats with the cache as a pure optimization. + + Any cache-layer failure degrades to a direct scan and a warning on + stderr. The JSON record is the contract; the cache is not. + """ + scan = LocalScan(today) + try: + cache_file, _ = scan_cache_paths() + cached = read_fresh_json(cache_file, max_age) + if ( + isinstance(cached, dict) + and cached.get("schemaVersion") == 1 + and cached.get("scanDate") == scan.today + and isinstance(cached.get("stats"), dict) + ): + return cached["stats"], True + complete = scan_opencode(scan) + stats = scan.snapshot() + if complete: + try: + write_json(cache_file, {"schemaVersion": 1, "scanDate": scan.today, "stats": stats}) + except Exception as exc: + print(f"omarchy-agent-usage-openrouter: could not write usage cache ({exc})", file=sys.stderr) + return stats, complete + except Exception as exc: + print(f"omarchy-agent-usage-openrouter: cache unavailable ({exc}); scanning directly", file=sys.stderr) + complete = scan_opencode(scan) + return scan.snapshot(), complete + + +# ---------------------------------------------------------------- record + +def base_record(**overrides: Any) -> dict: + record: dict[str, Any] = { + "schemaVersion": 1, + "id": AGENT_ID, + "name": AGENT_NAME, + "updatedAt": datetime.now(timezone.utc).isoformat(), + "ready": False, + "hasLocalStats": False, + # Tier-2 numbers are account-global (billing API), like Fireworks: + # every synced device reports the same truth, so aggregation must + # not sum them. Tier 1 (local DB scan) is machine-local instead. + "hasPromptStats": True, + "tierLabel": "Prepaid", + "usageStatusText": "", + "authHelpText": "", + "limits": [], + } + record.update(empty_stats()) + record.update(overrides) + return record + + +def key_limit_meter(key_info: dict) -> list[dict]: + # A key with a credit limit drains toward empty; map limit_reset through + # so the panel does not read a bare "" as "now". + try: + limit = key_info.get("limit") + remaining = key_info.get("limit_remaining") + if limit is None or remaining is None or float(limit) <= 0: + return [] + lim = float(limit) + rem = max(0.0, float(remaining)) + entry: dict[str, Any] = {"label": "Key limit", "percent": (lim - rem) / lim} + reset = str(key_info.get("limit_reset") or "").strip() + if reset: + entry["resetsAt"] = reset + return [entry] + except (TypeError, ValueError): + return [] + + +def scan(args: Any) -> dict[str, Any]: + config = read_config() + key = api_key(config) + if not key: + return base_record(usageStatusText="OpenRouter unavailable", authHelpText=AUTH_HELP) + + today = datetime.now().astimezone() + + # Tier 2 first: with a management key the 30-day account history covers + # every machine and client. Without one (or a rejected one) fall back to + # the local opencode DB scan below. + managed = activity_stats(management_key(config), today) + if managed is not None: + stats = managed + complete = True + has_prompt_stats = False + # Account-global numbers must merge by widest value, not by sum. + stats["scope"] = "account" + else: + max_age = 0 if args.force else (LIMITS_ONLY_REUSE_SECONDS if args.limits_only else SCAN_REUSE_SECONDS) + stats, complete = cached_local_stats(today, max_age) + has_prompt_stats = True + has_data = number(stats.get("totalPrompts")) > 0 or bool(stats.get("modelUsage")) + + try: + key_info = api_get("/key", key).get("data", {}) + if not isinstance(key_info, dict): + key_info = {} + except OpenRouterError as error: + if str(error) == "HTTP_401": + return base_record(usageStatusText="OpenRouter unavailable", + authHelpText="OpenRouter rejected the API key. " + AUTH_HELP) + return base_record(usageStatusText="OpenRouter unavailable", authHelpText=str(error)) + + # Balance: total purchased minus total used. A transport failure here + # must not take the token stats down -- retryAdvised asks the shell for + # one sooner retry instead of waiting out the full refresh interval. + balance = None + retry_advised = False + try: + credits = api_get("/credits", key).get("data", {}) + total_credits = float(credits.get("total_credits") or 0) + total_usage = float(credits.get("total_usage") or 0) + if total_credits > 0: + balance = { + "remaining": round(max(0.0, total_credits - total_usage), 6), + "funded": round(total_credits, 6), + "spent": round(max(0.0, total_usage), 6), + "currency": "USD", + "estimated": False, + } + except OpenRouterError as error: + if str(error) == "Could not reach the OpenRouter API": + retry_advised = True + else: + print(f"omarchy-agent-usage-openrouter: credits lookup failed ({error})", file=sys.stderr) + + record = base_record( + ready=True, + hasLocalStats=complete, + hasPromptStats=has_prompt_stats, + limits=key_limit_meter(key_info), + ) + if retry_advised: + record["retryAdvised"] = True + record.update(stats) + if balance: + record["balance"] = balance + if not has_data and not balance and not record["limits"]: + record["ready"] = False + record["usageStatusText"] = "No OpenRouter usage found" + record["authHelpText"] = AUTH_HELP + return record + + +def main() -> int: + ap = argparse.ArgumentParser(description="Print the OpenRouter usage record as JSON") + # --force rescans everything and rewrites the cache; --limits-only + # reuses a recent local scan and only the balance/limits probe is fresh. + # The flags exist so every collector accepts the same invocation. + ap.add_argument("--force", action="store_true", help="rescan transcripts and re-probe limits, ignoring caches") + ap.add_argument("--limits-only", action="store_true", help="reuse any recent transcript scan; only the limits probe must be fresh") + args = ap.parse_args() + try: + record = scan(args) + except OpenRouterError as error: + record = base_record(usageStatusText="OpenRouter unavailable", authHelpText=str(error)) + except Exception as error: + record = base_record(usageStatusText="OpenRouter unavailable", authHelpText="OpenRouter scan failed") + print(f"omarchy-agent-usage-openrouter: {type(error).__name__}", file=sys.stderr) + print(json.dumps(record, separators=(",", ":"))) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/shell/plugins/agents/README.md b/shell/plugins/agents/README.md index 9daa59da97c..b36d5c443ef 100644 --- a/shell/plugins/agents/README.md +++ b/shell/plugins/agents/README.md @@ -79,6 +79,7 @@ light surfaces — and the bar glyph stands in when there is none. | `zai` | Z.ai's coding-plan monitor endpoint (5-hour session + 7-day weekly) | pi/omp sessions on the Z.ai provider (`glm-*` models) | | `ollama` | Ollama Cloud's usage endpoint (5-hour session + 7-day weekly) | pi and omp sessions on the `ollama-cloud` provider | | `minimax` | MiniMax Token Plan API (5-hour session + 7-day weekly) | Not available from the quota API | +| `openrouter` | Prepaid balance from `/credits`, key-limit meter from `/key` | opencode sessions on the OpenRouter provider, or the `/activity` billing API with a management key | When `~/.local/state/omarchy/agents/accounts/.json` registers more than one account, the `claude`, `codex`, and `grok` records @@ -173,6 +174,38 @@ accounts. Without a configured `fundedAmount` the tab still shows token usage, just no balance. With a live ledger, `fundedAmount` is optional and only adds the meter and the spent-of-funded line under the real figure. +### OpenRouter balance and history + +OpenRouter reads `OPENROUTER_API_KEY` first, then the key +opencode stores in `~/.local/share/opencode/auth.json` when OpenRouter is +signed in there, then `apiKey` in `~/.config/omarchy/agents/openrouter.json`. +A management key in `OPENROUTER_MANAGEMENT_KEY` (or `managementKey` in that +file, from `openrouter.ai/settings/management-keys`) switches the token +history from the local opencode scan to the account-global `/activity` API. + +The collector reads the prepaid ledger from `GET /api/v1/credits` +(`remaining = total_credits - total_usage`) and reports it as a live, +non-estimated balance. A key with a credit limit additionally gets a +draining key-limit meter from `GET /api/v1/key`. + +Token history comes in two tiers, selected automatically by key type: + +```json +{ + "apiKey": "", + "managementKey": "" +} +``` + +Without a management key, tokens come from the local opencode database +(`providerID == "openrouter"`, read-only), like the claude/codex collectors. +With `OPENROUTER_MANAGEMENT_KEY` (or `managementKey` in +`~/.config/omarchy/agents/openrouter.json`), the last 30 days come from +`GET /api/v1/activity` instead: account-global per-day/per-model tokens that +cover every machine and client. That record carries `"scope": "account"` and +`hasPromptStats: false`, like Fireworks, so synced aggregation merges it by +widest value instead of summing it. + ## Interactions - Bar icon: left = panel, right = launch agent, middle = refresh. It turns diff --git a/shell/plugins/agents/manifest.json b/shell/plugins/agents/manifest.json index bc42f51feef..2969b557a30 100644 --- a/shell/plugins/agents/manifest.json +++ b/shell/plugins/agents/manifest.json @@ -27,7 +27,8 @@ "agy": { "enabled": true }, "zai": { "enabled": true }, "ollama": { "enabled": true }, - "minimax": { "enabled": true } + "minimax": { "enabled": true }, + "openrouter": { "enabled": true } }, "refreshIntervalSec": 900, "syncMode": "Off", diff --git a/test/shell.d/agent-usage-openrouter-scanner-test.sh b/test/shell.d/agent-usage-openrouter-scanner-test.sh new file mode 100755 index 00000000000..ea3b2680310 --- /dev/null +++ b/test/shell.d/agent-usage-openrouter-scanner-test.sh @@ -0,0 +1,136 @@ +#!/bin/bash + +source "$(dirname "$0")/base-test.sh" + +require_command jq +require_command python3 +require_command sqlite3 + +COLLECTOR="$ROOT/bin/omarchy-agent-usage-openrouter" + +[[ -x $COLLECTOR ]] || fail "collector missing/executable" "$COLLECTOR" + +TEST_HOME=$(mktemp -d) +trap 'rm -rf "$TEST_HOME"' EXIT + +export HOME="$TEST_HOME" +export XDG_CONFIG_HOME="$TEST_HOME/.config" +export XDG_DATA_HOME="$TEST_HOME/.local/share" +export XDG_CACHE_HOME="$TEST_HOME/.cache" +export XDG_STATE_HOME="$TEST_HOME/.local/state" +export OPENROUTER_API_KEY="" +export OPENROUTER_MANAGEMENT_KEY="" +# Bucket dates resolve in local time, so pin the zone or the fixtures below +# would shift by a day depending on where the test runs. +export TZ="UTC" + +# 1. Without credentials: valid record, ready=false +no_key=$("$COLLECTOR") +[[ $(jq -r '.id + ":" + (.ready | tostring)' <<<"$no_key") == "openrouter:false" ]] || + fail "prints a valid record without credentials" "$no_key" +pass "prints a valid record without credentials" + +# 2. Fixture opencode DB: 2 assistant rows (openrouter) + 1 other provider (ignored) +mkdir -p "$XDG_DATA_HOME/opencode" +sqlite3 "$XDG_DATA_HOME/opencode/opencode.db" <<'SQL' +CREATE TABLE message (id TEXT PRIMARY KEY, session_id TEXT NOT NULL, time_created INTEGER NOT NULL, time_updated INTEGER NOT NULL, data TEXT NOT NULL); +INSERT INTO message VALUES + ('m1','ses_aaa',1787000000000,1787000000000,'{"role":"assistant","providerID":"openrouter","modelID":"openai/gpt-4o-mini","tokens":{"input":100,"output":20,"reasoning":5,"cache":{"read":10,"write":0}},"time":{"created":1787000000000}}'), + ('m2','ses_aaa',1787000001000,1787000001000,'{"role":"user","providerID":"openrouter","modelID":"openai/gpt-4o-mini","time":{"created":1787000001000}}'), + ('m3','ses_bbb',1787000002000,1787000002000,'{"role":"assistant","providerID":"openai","modelID":"gpt-5","tokens":{"input":999,"output":999},"time":{"created":1787000002000}}'); +SQL + +result=$(python3 - "$COLLECTOR" <<'PY' +import importlib.machinery, importlib.util, json, sys +from datetime import datetime, timezone +from pathlib import Path +loader = importlib.machinery.SourceFileLoader("or_collector", sys.argv[1]) +spec = importlib.util.spec_from_loader(loader.name, loader) +mod = importlib.util.module_from_spec(spec) +loader.exec_module(mod) + +today = datetime.now().astimezone() +scan = mod.LocalScan(today) +complete = mod.scan_opencode(scan) +stats = scan.snapshot() + +# /activity fold: prompt+reasoning -> input, completion -> output +activity = mod.summarize_activity([ + {"date": today.strftime("%Y-%m-%d"), "model": "openai/gpt-4o-mini", + "prompt_tokens": 50, "completion_tokens": 25, "reasoning_tokens": 10}, + {"date": "2020-01-01", "model": "x/y-model", + "prompt_tokens": 7, "completion_tokens": 3, "reasoning_tokens": 0}, + "garbage-row", +], today) + +out = { + "complete": complete, + "totalPrompts": stats["totalPrompts"], + "totalTokens": stats["todayTotalTokens"] + sum(d["messageCount"] for d in stats["recentDays"] if d["date"] != today.strftime("%Y-%m-%d")), + "model": sorted(stats["modelUsage"].keys()), + "hasPromptStats_tier1": True, + "activityToday": activity["todayTotalTokens"], + "activityModels": sorted(activity["modelUsage"].keys()), + "activityOldDay": activity["activeDates"], + "keyMeterEmpty": mod.key_limit_meter({}) == [], + "keyMeter": mod.key_limit_meter({"limit": 10, "limit_remaining": 4, "limit_reset": "weekly"}), + "scopeNote": "scope set only in tier2 path", +} +print(json.dumps(out)) +PY +) + +check() { + local actual="$1" expected="$2" label="$3" + [[ $actual == "$expected" ]] || fail "$label (got: $actual, want: $expected)" "$result" + pass "$label" +} + +check "$(jq -r .complete <<<"$result")" "true" "opencode scan completes" +check "$(jq -r .totalPrompts <<<"$result")" "1" "counts only openrouter assistant rows" +check "$(jq -r '.model | join(",")' <<<"$result")" "gpt-4o-mini" "strips provider prefix from model" +check "$(jq -r .activityToday <<<"$result")" "85" "activity folds prompt+completion+reasoning" +check "$(jq -r '.activityModels | join(",")' <<<"$result")" "gpt-4o-mini,y-model" "activity keeps per-model buckets" +check "$(jq -r '.activityOldDay | length' <<<"$result")" "2" "activity tracks active dates" +check "$(jq -r .keyMeterEmpty <<<"$result")" "true" "no meter without key limit" +check "$(jq -r '.keyMeter[0].percent' <<<"$result")" "0.6" "key meter drains toward empty" +check "$(jq -r '.keyMeter[0].resetsAt' <<<"$result")" "weekly" "key meter maps limit_reset" + +# 3. Tier-2 path sets scope=account + hasPromptStats=false (stub api_get) +scope_check=$(python3 - "$COLLECTOR" <<'PY' +import importlib.machinery, importlib.util, json, sys +from datetime import datetime +loader = importlib.machinery.SourceFileLoader("or_collector2", sys.argv[1]) +spec = importlib.util.spec_from_loader(loader.name, loader) +mod = importlib.util.module_from_spec(spec) +loader.exec_module(mod) + +captured = {} +def fake_api_get(path, key): + captured[path] = key + if path == "/activity": + return {"data": [{"date": "2026-09-08", "model": "a/b", "prompt_tokens": 10, + "completion_tokens": 5, "reasoning_tokens": 0}]} + if path == "/key": + return {"data": {"limit": None, "limit_remaining": None}} + if path == "/credits": + return {"data": {"total_credits": 10, "total_usage": 1}} + raise AssertionError(path) +mod.api_get = fake_api_get + +class Args: force = True; limits_only = False +import os +os.environ["OPENROUTER_API_KEY"] = "sk-test" +os.environ["OPENROUTER_MANAGEMENT_KEY"] = "sk-mgmt" +rec = mod.scan(Args()) +print(json.dumps({"scope": rec.get("scope"), "hasPromptStats": rec.get("hasPromptStats"), + "balance": rec.get("balance", {}).get("remaining"), + "called": sorted(captured.keys())})) +PY +) +check "$(jq -r .scope <<<"$scope_check")" "account" "tier-2 record is scope=account" +check "$(jq -r .hasPromptStats <<<"$scope_check")" "false" "tier-2 hides prompt counts" +check "$(jq -r .balance <<<"$scope_check")" "9.0" "tier-2 keeps balance from /credits" +check "$(jq -r '.called | join(",")' <<<"$scope_check")" "/activity,/credits,/key" "tier-2 probes activity + key + credits" + +pass "ALL OPENROUTER SCANNER TESTS PASSED" From 8e45c3ce8741d16c43dd7e613f144e08e66c0278 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Fri, 2 Oct 2026 06:14:32 -0400 Subject: [PATCH 12/19] Keep a stale OpenRouter key-limit meter and indent with two spaces A failed /key lookup no longer throws away the token stats and balance: the last good key-limit meter stays, marked with limitsStale and limitsFetchedAt. limit_reset names a period (daily, weekly, monthly), not a time, so it labels the meter instead of filling resetsAt. Without a key the record has no plan, so the panel skips it. The manual mentions OpenRouter. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-openrouter | 863 +++++++++--------- manual/17-ai.md | 2 +- shell/plugins/agents/README.md | 3 +- .../agent-usage-openrouter-scanner-test.sh | 38 +- 4 files changed, 483 insertions(+), 423 deletions(-) diff --git a/bin/omarchy-agent-usage-openrouter b/bin/omarchy-agent-usage-openrouter index bd6b908e330..0bbcd05651b 100755 --- a/bin/omarchy-agent-usage-openrouter +++ b/bin/omarchy-agent-usage-openrouter @@ -28,6 +28,7 @@ import os import sqlite3 import sys import tempfile +import time import urllib.request from datetime import datetime, timedelta, timezone from pathlib import Path @@ -50,514 +51,538 @@ class OpenRouterError(Exception): def number(value: Any) -> int: - try: - return max(0, int(value or 0)) - except (TypeError, ValueError): - return 0 + try: + return max(0, int(value or 0)) + except (TypeError, ValueError): + return 0 def local_day(value) -> str: - now = datetime.now() - if value is None: - return now.strftime("%Y-%m-%d") - if isinstance(value, (int, float)): - ts = value / 1000 if value > 10_000_000_000 else value - return datetime.fromtimestamp(ts).astimezone().strftime("%Y-%m-%d") - try: - text = str(value) - dt = datetime.fromisoformat(text.replace("Z", "+00:00")) - if dt.tzinfo is not None: - dt = dt.astimezone() - return dt.strftime("%Y-%m-%d") - except Exception: - return now.strftime("%Y-%m-%d") + now = datetime.now() + if value is None: + return now.strftime("%Y-%m-%d") + if isinstance(value, (int, float)): + ts = value / 1000 if value > 10_000_000_000 else value + return datetime.fromtimestamp(ts).astimezone().strftime("%Y-%m-%d") + try: + text = str(value) + dt = datetime.fromisoformat(text.replace("Z", "+00:00")) + if dt.tzinfo is not None: + dt = dt.astimezone() + return dt.strftime("%Y-%m-%d") + except Exception: + return now.strftime("%Y-%m-%d") def model_name(raw: Any) -> str: - value = str(raw or "openrouter").rstrip("/").split("/")[-1] - return value or "openrouter" + value = str(raw or "openrouter").rstrip("/").split("/")[-1] + return value or "openrouter" def empty_bucket() -> dict[str, int]: - return { - "inputTokens": 0, - "outputTokens": 0, - "cacheReadInputTokens": 0, - "cacheCreationInputTokens": 0, - } + return { + "inputTokens": 0, + "outputTokens": 0, + "cacheReadInputTokens": 0, + "cacheCreationInputTokens": 0, + } def empty_stats() -> dict[str, Any]: - return { - "todayPrompts": 0, - "todaySessions": 0, - "todayTotalTokens": 0, - "todayTokensByModel": {}, - "recentDays": [], - "totalPrompts": 0, - "totalSessions": 0, - "activeDays": 0, - "activeDates": [], - "modelUsage": {}, - } + return { + "todayPrompts": 0, + "todaySessions": 0, + "todayTotalTokens": 0, + "todayTokensByModel": {}, + "recentDays": [], + "totalPrompts": 0, + "totalSessions": 0, + "activeDays": 0, + "activeDates": [], + "modelUsage": {}, + } # ---------------------------------------------------------------- auth def config_path() -> Path: - config_home = Path(os.environ.get("XDG_CONFIG_HOME") or (Path.home() / ".config")) - return config_home / "omarchy" / "agents" / "openrouter.json" + config_home = Path(os.environ.get("XDG_CONFIG_HOME") or (Path.home() / ".config")) + return config_home / "omarchy" / "agents" / "openrouter.json" def read_config() -> dict[str, Any]: - try: - parsed = json.loads(config_path().read_text()) - return parsed if isinstance(parsed, dict) else {} - except (OSError, json.JSONDecodeError): - return {} + try: + parsed = json.loads(config_path().read_text()) + return parsed if isinstance(parsed, dict) else {} + except (OSError, json.JSONDecodeError): + return {} def opencode_auth_path() -> Path: - data_home = Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) - return data_home / "opencode" / "auth.json" + data_home = Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) + return data_home / "opencode" / "auth.json" def read_opencode_key(path: Path) -> str: - try: - parsed = json.loads(path.read_text()) - except (OSError, json.JSONDecodeError): - return "" - entry = parsed.get("openrouter") if isinstance(parsed, dict) else None - if not isinstance(entry, dict): - return "" - return str(entry.get("key") or "").strip() + try: + parsed = json.loads(path.read_text()) + except (OSError, json.JSONDecodeError): + return "" + entry = parsed.get("openrouter") if isinstance(parsed, dict) else None + if not isinstance(entry, dict): + return "" + return str(entry.get("key") or "").strip() def api_key(config: dict[str, Any]) -> str: - # Explicit key or opencode login wins; the config file is the last - # resort for headless/refresh-timer runs without either. - return ( - str(os.environ.get("OPENROUTER_API_KEY", "")).strip() - or read_opencode_key(opencode_auth_path()) - or str(config.get("apiKey") or "").strip() - ) + # Explicit key or opencode login wins; the config file is the last + # resort for headless/refresh-timer runs without either. + return ( + str(os.environ.get("OPENROUTER_API_KEY", "")).strip() + or read_opencode_key(opencode_auth_path()) + or str(config.get("apiKey") or "").strip() + ) def management_key(config: dict[str, Any]) -> str: - # A management key unlocks /activity (30-day account history). A normal - # inference key answers /key + /credits only and fails /activity with - # 403 -- never send it there speculatively, just skip tier 2. - return ( - str(os.environ.get("OPENROUTER_MANAGEMENT_KEY", "")).strip() - or str(config.get("managementKey") or "").strip() - ) + # A management key unlocks /activity (30-day account history). A normal + # inference key answers /key + /credits only and fails /activity with + # 403 -- never send it there speculatively, just skip tier 2. + return ( + str(os.environ.get("OPENROUTER_MANAGEMENT_KEY", "")).strip() + or str(config.get("managementKey") or "").strip() + ) # ---------------------------------------------------------------- API class OpenRouterTransportError(OpenRouterError): - pass + pass def api_get(path: str, key: str) -> dict: - req = urllib.request.Request( - API_BASE + path, headers={"Authorization": "Bearer " + key, "Accept": "application/json"} - ) - try: - with urllib.request.urlopen(req, timeout=15) as resp: - payload = json.load(resp) - return payload if isinstance(payload, dict) else {} - except urllib.request.HTTPError as error: - if error.code in (401, 403): - # 403 from /activity with an inference key is the expected - # tier-1 signal ("Only management keys can perform this - # operation"), not a broken setup -- callers decide. - raise OpenRouterError(f"HTTP_{error.code}") - raise OpenRouterError(f"OpenRouter API returned HTTP {error.code}") - except urllib.request.URLError as error: - raise OpenRouterTransportError("Could not reach the OpenRouter API") from error - except (json.JSONDecodeError, ValueError) as error: - raise OpenRouterError("OpenRouter returned an invalid response") from error + req = urllib.request.Request( + API_BASE + path, headers={"Authorization": "Bearer " + key, "Accept": "application/json"} + ) + try: + with urllib.request.urlopen(req, timeout=15) as resp: + payload = json.load(resp) + return payload if isinstance(payload, dict) else {} + except urllib.request.HTTPError as error: + if error.code in (401, 403): + # 403 from /activity with an inference key is the expected + # tier-1 signal ("Only management keys can perform this + # operation"), not a broken setup -- callers decide. + raise OpenRouterError(f"HTTP_{error.code}") + raise OpenRouterError(f"OpenRouter API returned HTTP {error.code}") + except urllib.request.URLError as error: + raise OpenRouterTransportError("Could not reach the OpenRouter API") from error + except (json.JSONDecodeError, ValueError) as error: + raise OpenRouterError("OpenRouter returned an invalid response") from error def summarize_activity(items: list[Any], today: datetime) -> dict[str, Any]: - """Fold /activity rows into the panel's stats dict. - - One row per (date, endpoint): prompt + reasoning tokens count as input - (they are generated), completion as output. The billing API reports - tokens, never prompt or session counts, so hasPromptStats stays False - and the panel keeps those numbers out of today's tooltip. - """ - today_date = today.strftime("%Y-%m-%d") - recent_dates = [(today - timedelta(days=offset)).strftime("%Y-%m-%d") for offset in range(6, -1, -1)] - recent = {day: 0 for day in recent_dates} - today_by_model: dict[str, int] = {} - model_usage: dict[str, dict[str, int]] = {} - active_dates: set[str] = set() - - for raw in items: - if not isinstance(raw, dict): - continue - day = str(raw.get("date") or "") - model = model_name(str(raw.get("model_permaslug") or raw.get("model") or "")) - prompt = number(raw.get("prompt_tokens")) - reasoning = number(raw.get("reasoning_tokens")) - output = number(raw.get("completion_tokens")) - total = prompt + reasoning + output - if total <= 0: - continue - bucket = model_usage.setdefault(model, empty_bucket()) - bucket["inputTokens"] += prompt - bucket["outputTokens"] += output + reasoning - if day: - active_dates.add(day) - if day in recent: - recent[day] += total - if day == today_date: - today_by_model[model] = today_by_model.get(model, 0) + total - - return { - "todayTokensByModel": today_by_model, - "todayTotalTokens": sum(today_by_model.values()), - "recentDays": [{"date": day, "messageCount": recent[day]} for day in recent_dates], - "activeDays": len(active_dates), - "activeDates": sorted(active_dates), - "modelUsage": model_usage, - } + """Fold /activity rows into the panel's stats dict. + + One row per (date, endpoint): prompt + reasoning tokens count as input + (they are generated), completion as output. The billing API reports + tokens, never prompt or session counts, so hasPromptStats stays False + and the panel keeps those numbers out of today's tooltip. + """ + today_date = today.strftime("%Y-%m-%d") + recent_dates = [(today - timedelta(days=offset)).strftime("%Y-%m-%d") for offset in range(6, -1, -1)] + recent = {day: 0 for day in recent_dates} + today_by_model: dict[str, int] = {} + model_usage: dict[str, dict[str, int]] = {} + active_dates: set[str] = set() + + for raw in items: + if not isinstance(raw, dict): + continue + day = str(raw.get("date") or "") + model = model_name(str(raw.get("model_permaslug") or raw.get("model") or "")) + prompt = number(raw.get("prompt_tokens")) + reasoning = number(raw.get("reasoning_tokens")) + output = number(raw.get("completion_tokens")) + total = prompt + reasoning + output + if total <= 0: + continue + bucket = model_usage.setdefault(model, empty_bucket()) + bucket["inputTokens"] += prompt + bucket["outputTokens"] += output + reasoning + if day: + active_dates.add(day) + if day in recent: + recent[day] += total + if day == today_date: + today_by_model[model] = today_by_model.get(model, 0) + total + + return { + "todayTokensByModel": today_by_model, + "todayTotalTokens": sum(today_by_model.values()), + "recentDays": [{"date": day, "messageCount": recent[day]} for day in recent_dates], + "activeDays": len(active_dates), + "activeDates": sorted(active_dates), + "modelUsage": model_usage, + } def activity_stats(key: str, today: datetime) -> dict[str, Any] | None: - """Tier 2: 30-day account history via management key, or None. + """Tier 2: 30-day account history via management key, or None. - Returns None when there is no management key (plain tier-1 run) or the - key is rejected -- the caller then falls back to the local DB scan. - """ - if not key: - return None - try: - payload = api_get("/activity", key) - except OpenRouterError: - return None - items = payload.get("data") - if not isinstance(items, list): - return None - stats = empty_stats() - stats.update(summarize_activity(items, today)) - return stats + Returns None when there is no management key (plain tier-1 run) or the + key is rejected -- the caller then falls back to the local DB scan. + """ + if not key: + return None + try: + payload = api_get("/activity", key) + except OpenRouterError: + return None + items = payload.get("data") + if not isinstance(items, list): + return None + stats = empty_stats() + stats.update(summarize_activity(items, today)) + return stats # ---------------------------------------------------------------- local stats (opencode DB) class LocalScan: - """Aggregated opencode usage. Module-level in the original prototype; - an instance keeps the scanner testable without global state.""" - - def __init__(self, today: datetime): - self.today = today.strftime("%Y-%m-%d") - self.recent_dates = [(today - timedelta(days=o)).strftime("%Y-%m-%d") for o in range(6, -1, -1)] - self.recent = {day: {"date": day, "messageCount": 0} for day in self.recent_dates} - self.today_tokens_by_model: dict[str, int] = {} - self.model_usage: dict[str, dict[str, int]] = {} - self.today_sessions: set[str] = set() - self.active_days: set[str] = set() - self.today_prompts = 0 - self.today_total_tokens = 0 - self.total_prompts = 0 - self.total_sessions: set[str] = set() - - def add_usage(self, day, session_key, model, in_tok, out_tok, cache_read, cache_write): - total = in_tok + out_tok + cache_read + cache_write - self.total_prompts += 1 - self.total_sessions.add(session_key) - self.active_days.add(day) - bucket = self.model_usage.setdefault(model, empty_bucket()) - bucket["inputTokens"] += in_tok - bucket["outputTokens"] += out_tok - bucket["cacheReadInputTokens"] += cache_read - bucket["cacheCreationInputTokens"] += cache_write - if day in self.recent: - self.recent[day]["messageCount"] += total - if day == self.today: - self.today_prompts += 1 - self.today_sessions.add(session_key) - self.today_total_tokens += total - self.today_tokens_by_model[model] = self.today_tokens_by_model.get(model, 0) + total - - def snapshot(self) -> dict: - return { - "todayPrompts": self.today_prompts, - "todaySessions": len(self.today_sessions), - "todayTotalTokens": self.today_total_tokens, - "todayTokensByModel": self.today_tokens_by_model, - "recentDays": [self.recent[d] for d in self.recent_dates], - "totalPrompts": self.total_prompts, - "totalSessions": len(self.total_sessions), - "activeDays": len(self.active_days), - "activeDates": sorted(self.active_days), - "modelUsage": self.model_usage, - } + """Aggregated opencode usage. Module-level in the original prototype; + an instance keeps the scanner testable without global state.""" + + def __init__(self, today: datetime): + self.today = today.strftime("%Y-%m-%d") + self.recent_dates = [(today - timedelta(days=o)).strftime("%Y-%m-%d") for o in range(6, -1, -1)] + self.recent = {day: {"date": day, "messageCount": 0} for day in self.recent_dates} + self.today_tokens_by_model: dict[str, int] = {} + self.model_usage: dict[str, dict[str, int]] = {} + self.today_sessions: set[str] = set() + self.active_days: set[str] = set() + self.today_prompts = 0 + self.today_total_tokens = 0 + self.total_prompts = 0 + self.total_sessions: set[str] = set() + + def add_usage(self, day, session_key, model, in_tok, out_tok, cache_read, cache_write): + total = in_tok + out_tok + cache_read + cache_write + self.total_prompts += 1 + self.total_sessions.add(session_key) + self.active_days.add(day) + bucket = self.model_usage.setdefault(model, empty_bucket()) + bucket["inputTokens"] += in_tok + bucket["outputTokens"] += out_tok + bucket["cacheReadInputTokens"] += cache_read + bucket["cacheCreationInputTokens"] += cache_write + if day in self.recent: + self.recent[day]["messageCount"] += total + if day == self.today: + self.today_prompts += 1 + self.today_sessions.add(session_key) + self.today_total_tokens += total + self.today_tokens_by_model[model] = self.today_tokens_by_model.get(model, 0) + total + + def snapshot(self) -> dict: + return { + "todayPrompts": self.today_prompts, + "todaySessions": len(self.today_sessions), + "todayTotalTokens": self.today_total_tokens, + "todayTokensByModel": self.today_tokens_by_model, + "recentDays": [self.recent[d] for d in self.recent_dates], + "totalPrompts": self.total_prompts, + "totalSessions": len(self.total_sessions), + "activeDays": len(self.active_days), + "activeDates": sorted(self.active_days), + "modelUsage": self.model_usage, + } def opencode_db_path() -> Path: - return Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) / "opencode" / "opencode.db" + return Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) / "opencode" / "opencode.db" def scan_opencode(scan: LocalScan, db: Path | None = None) -> bool: - # Returns whether the scan ran to completion: a scan cut short by a - # database error still contributes what it read, but must not be cached - # as if it were the whole story. - db = db or opencode_db_path() - if not db.is_file(): - return True - try: - conn = sqlite3.connect(db.resolve().as_uri() + "?mode=ro", uri=True, timeout=2) - except sqlite3.Error: - return False - try: - conn.execute("PRAGMA query_only = ON") - # The LIKE gates skip rows whose JSON cannot contain the two - # key/value pairs, avoiding the JSON parse of huge message blobs; - # json_extract is the authority for the rows that reach Python. - rows = conn.execute( - "SELECT session_id, data FROM message" - " WHERE data LIKE '%\"role\"%:%\"assistant\"%'" - " AND data LIKE '%\"providerID\"%:%\"openrouter\"%'" - " AND CASE WHEN json_valid(data) THEN json_extract(data, '$.role') END = 'assistant'" - " AND CASE WHEN json_valid(data) THEN json_extract(data, '$.providerID') END = 'openrouter'" - ) - for session_id, raw in rows: - try: - entry = json.loads(raw) - if not isinstance(entry, dict) or entry.get("role") != "assistant": - continue - if str(entry.get("providerID") or "") != "openrouter": - continue - tokens = entry.get("tokens") or {} - cache = tokens.get("cache") or {} - in_tok = number(tokens.get("input")) - # opencode keeps thinking tokens out of output; both are generated. - out_tok = number(tokens.get("output")) + number(tokens.get("reasoning")) - cache_read = number(cache.get("read")) - cache_write = number(cache.get("write")) - if not (in_tok or out_tok or cache_read or cache_write): - continue - day = local_day((entry.get("time") or {}).get("created")) - model = model_name(str(entry.get("modelID") or "")) - except Exception: - continue - scan.add_usage(day, "opencode:" + str(session_id), model, in_tok, out_tok, cache_read, cache_write) - except sqlite3.Error: - return False - finally: - conn.close() + # Returns whether the scan ran to completion: a scan cut short by a + # database error still contributes what it read, but must not be cached + # as if it were the whole story. + db = db or opencode_db_path() + if not db.is_file(): return True + try: + conn = sqlite3.connect(db.resolve().as_uri() + "?mode=ro", uri=True, timeout=2) + except sqlite3.Error: + return False + try: + conn.execute("PRAGMA query_only = ON") + # The LIKE gates skip rows whose JSON cannot contain the two + # key/value pairs, avoiding the JSON parse of huge message blobs; + # json_extract is the authority for the rows that reach Python. + rows = conn.execute( + "SELECT session_id, data FROM message" + " WHERE data LIKE '%\"role\"%:%\"assistant\"%'" + " AND data LIKE '%\"providerID\"%:%\"openrouter\"%'" + " AND CASE WHEN json_valid(data) THEN json_extract(data, '$.role') END = 'assistant'" + " AND CASE WHEN json_valid(data) THEN json_extract(data, '$.providerID') END = 'openrouter'" + ) + for session_id, raw in rows: + try: + entry = json.loads(raw) + if not isinstance(entry, dict) or entry.get("role") != "assistant": + continue + if str(entry.get("providerID") or "") != "openrouter": + continue + tokens = entry.get("tokens") or {} + cache = tokens.get("cache") or {} + in_tok = number(tokens.get("input")) + # opencode keeps thinking tokens out of output; both are generated. + out_tok = number(tokens.get("output")) + number(tokens.get("reasoning")) + cache_read = number(cache.get("read")) + cache_write = number(cache.get("write")) + if not (in_tok or out_tok or cache_read or cache_write): + continue + day = local_day((entry.get("time") or {}).get("created")) + model = model_name(str(entry.get("modelID") or "")) + except Exception: + continue + scan.add_usage(day, "opencode:" + str(session_id), model, in_tok, out_tok, cache_read, cache_write) + except sqlite3.Error: + return False + finally: + conn.close() + return True def cache_root() -> Path: - root = Path(os.environ.get("XDG_CACHE_HOME") or (Path.home() / ".cache")) / "omarchy" / "agent-usage" - root.mkdir(parents=True, exist_ok=True) - return root + root = Path(os.environ.get("XDG_CACHE_HOME") or (Path.home() / ".cache")) / "omarchy" / "agent-usage" + root.mkdir(parents=True, exist_ok=True) + return root def scan_cache_paths() -> tuple[Path, Path]: - db = opencode_db_path() - digest = hashlib.sha1(str(db).encode("utf-8")).hexdigest()[:16] - root = cache_root() - return root / f"openrouter-scan-{digest}.json", root / f"openrouter-scan-{digest}.lock" + db = opencode_db_path() + digest = hashlib.sha1(str(db).encode("utf-8")).hexdigest()[:16] + root = cache_root() + return root / f"openrouter-scan-{digest}.json", root / f"openrouter-scan-{digest}.lock" def read_fresh_json(path: Path, max_age_seconds: float) -> dict | None: - if max_age_seconds <= 0 or not path.exists(): - return None - try: - age = datetime.now().timestamp() - path.stat().st_mtime - if 0 <= age <= max_age_seconds: - parsed = json.loads(path.read_text(encoding="utf-8")) - return parsed if isinstance(parsed, dict) else None - except Exception: - return None + if max_age_seconds <= 0 or not path.exists(): + return None + try: + age = datetime.now().timestamp() - path.stat().st_mtime + if 0 <= age <= max_age_seconds: + parsed = json.loads(path.read_text(encoding="utf-8")) + return parsed if isinstance(parsed, dict) else None + except Exception: return None + return None def write_json(path: Path, payload: dict) -> None: - 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=(",", ":")) + "\n") - tmp.chmod(0o644) - tmp.replace(path) - except BaseException: - tmp.unlink(missing_ok=True) - raise + 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=(",", ":")) + "\n") + tmp.chmod(0o644) + tmp.replace(path) + except BaseException: + tmp.unlink(missing_ok=True) + raise def cached_local_stats(today: datetime, max_age: float) -> tuple[dict, bool]: - """Local stats with the cache as a pure optimization. - - Any cache-layer failure degrades to a direct scan and a warning on - stderr. The JSON record is the contract; the cache is not. - """ - scan = LocalScan(today) - try: - cache_file, _ = scan_cache_paths() - cached = read_fresh_json(cache_file, max_age) - if ( - isinstance(cached, dict) - and cached.get("schemaVersion") == 1 - and cached.get("scanDate") == scan.today - and isinstance(cached.get("stats"), dict) - ): - return cached["stats"], True - complete = scan_opencode(scan) - stats = scan.snapshot() - if complete: - try: - write_json(cache_file, {"schemaVersion": 1, "scanDate": scan.today, "stats": stats}) - except Exception as exc: - print(f"omarchy-agent-usage-openrouter: could not write usage cache ({exc})", file=sys.stderr) - return stats, complete - except Exception as exc: - print(f"omarchy-agent-usage-openrouter: cache unavailable ({exc}); scanning directly", file=sys.stderr) - complete = scan_opencode(scan) - return scan.snapshot(), complete + """Local stats with the cache as a pure optimization. + + Any cache-layer failure degrades to a direct scan and a warning on + stderr. The JSON record is the contract; the cache is not. + """ + scan = LocalScan(today) + try: + cache_file, _ = scan_cache_paths() + cached = read_fresh_json(cache_file, max_age) + if ( + isinstance(cached, dict) + and cached.get("schemaVersion") == 1 + and cached.get("scanDate") == scan.today + and isinstance(cached.get("stats"), dict) + ): + return cached["stats"], True + complete = scan_opencode(scan) + stats = scan.snapshot() + if complete: + try: + write_json(cache_file, {"schemaVersion": 1, "scanDate": scan.today, "stats": stats}) + except Exception as exc: + print(f"omarchy-agent-usage-openrouter: could not write usage cache ({exc})", file=sys.stderr) + return stats, complete + except Exception as exc: + print(f"omarchy-agent-usage-openrouter: cache unavailable ({exc}); scanning directly", file=sys.stderr) + complete = scan_opencode(scan) + return scan.snapshot(), complete # ---------------------------------------------------------------- record def base_record(**overrides: Any) -> dict: - record: dict[str, Any] = { - "schemaVersion": 1, - "id": AGENT_ID, - "name": AGENT_NAME, - "updatedAt": datetime.now(timezone.utc).isoformat(), - "ready": False, - "hasLocalStats": False, - # Tier-2 numbers are account-global (billing API), like Fireworks: - # every synced device reports the same truth, so aggregation must - # not sum them. Tier 1 (local DB scan) is machine-local instead. - "hasPromptStats": True, - "tierLabel": "Prepaid", - "usageStatusText": "", - "authHelpText": "", - "limits": [], - } - record.update(empty_stats()) - record.update(overrides) - return record + record: dict[str, Any] = { + "schemaVersion": 1, + "id": AGENT_ID, + "name": AGENT_NAME, + "updatedAt": datetime.now(timezone.utc).isoformat(), + "ready": False, + "hasLocalStats": False, + # Tier-2 numbers are account-global (billing API), like Fireworks: + # every synced device reports the same truth, so aggregation must + # not sum them. Tier 1 (local DB scan) is machine-local instead. + "hasPromptStats": True, + "tierLabel": "", + "usageStatusText": "", + "authHelpText": "", + "limits": [], + "limitsStale": False, + "limitsFetchedAt": 0, + } + record.update(empty_stats()) + record.update(overrides) + return record + + +# A key with a credit limit drains toward empty. limit_reset names how often +# the limit starts over (daily, weekly, monthly), not when, so it labels the +# meter rather than giving it a reset time. +KEY_LIMIT_PERIODS = {"daily": "Daily", "weekly": "Weekly", "monthly": "Monthly"} def key_limit_meter(key_info: dict) -> list[dict]: - # A key with a credit limit drains toward empty; map limit_reset through - # so the panel does not read a bare "" as "now". - try: - limit = key_info.get("limit") - remaining = key_info.get("limit_remaining") - if limit is None or remaining is None or float(limit) <= 0: - return [] - lim = float(limit) - rem = max(0.0, float(remaining)) - entry: dict[str, Any] = {"label": "Key limit", "percent": (lim - rem) / lim} - reset = str(key_info.get("limit_reset") or "").strip() - if reset: - entry["resetsAt"] = reset - return [entry] - except (TypeError, ValueError): - return [] + try: + limit = key_info.get("limit") + remaining = key_info.get("limit_remaining") + if limit is None or remaining is None or float(limit) <= 0: + return [] + lim = float(limit) + rem = max(0.0, float(remaining)) + except (TypeError, ValueError): + return [] + period = KEY_LIMIT_PERIODS.get(str(key_info.get("limit_reset") or "").strip().lower()) + label = f"{period} key limit" if period else "Key limit" + return [{"label": label, "title": label, "percent": (lim - rem) / lim, "resetsAt": ""}] -def scan(args: Any) -> dict[str, Any]: - config = read_config() - key = api_key(config) - if not key: - return base_record(usageStatusText="OpenRouter unavailable", authHelpText=AUTH_HELP) - - today = datetime.now().astimezone() - - # Tier 2 first: with a management key the 30-day account history covers - # every machine and client. Without one (or a rejected one) fall back to - # the local opencode DB scan below. - managed = activity_stats(management_key(config), today) - if managed is not None: - stats = managed - complete = True - has_prompt_stats = False - # Account-global numbers must merge by widest value, not by sum. - stats["scope"] = "account" - else: - max_age = 0 if args.force else (LIMITS_ONLY_REUSE_SECONDS if args.limits_only else SCAN_REUSE_SECONDS) - stats, complete = cached_local_stats(today, max_age) - has_prompt_stats = True - has_data = number(stats.get("totalPrompts")) > 0 or bool(stats.get("modelUsage")) +def limits_cache() -> Path: + return cache_root() / "openrouter-limits.json" + +def scan(args: Any) -> dict[str, Any]: + config = read_config() + key = api_key(config) + if not key: + return base_record(usageStatusText="OpenRouter unavailable", authHelpText=AUTH_HELP) + + today = datetime.now().astimezone() + + # Tier 2 first: with a management key the 30-day account history covers + # every machine and client. Without one (or a rejected one) fall back to + # the local opencode DB scan below. + managed = activity_stats(management_key(config), today) + if managed is not None: + stats = managed + complete = True + has_prompt_stats = False + # Account-global numbers must merge by widest value, not by sum. + stats["scope"] = "account" + else: + max_age = 0 if args.force else (LIMITS_ONLY_REUSE_SECONDS if args.limits_only else SCAN_REUSE_SECONDS) + stats, complete = cached_local_stats(today, max_age) + has_prompt_stats = True + has_data = number(stats.get("totalPrompts")) > 0 or bool(stats.get("modelUsage")) + + # The key-limit meter. A failed lookup keeps the last good meter, marked + # stale, and leaves the token stats standing. + limits, stale, fetched_at = [], False, 0 + status = help_text = "" + try: + key_info = api_get("/key", key).get("data", {}) + if not isinstance(key_info, dict): + key_info = {} + limits = key_limit_meter(key_info) + fetched_at = round(time.time() * 1000) try: - key_info = api_get("/key", key).get("data", {}) - if not isinstance(key_info, dict): - key_info = {} - except OpenRouterError as error: - if str(error) == "HTTP_401": - return base_record(usageStatusText="OpenRouter unavailable", - authHelpText="OpenRouter rejected the API key. " + AUTH_HELP) - return base_record(usageStatusText="OpenRouter unavailable", authHelpText=str(error)) - - # Balance: total purchased minus total used. A transport failure here - # must not take the token stats down -- retryAdvised asks the shell for - # one sooner retry instead of waiting out the full refresh interval. - balance = None - retry_advised = False - try: - credits = api_get("/credits", key).get("data", {}) - total_credits = float(credits.get("total_credits") or 0) - total_usage = float(credits.get("total_usage") or 0) - if total_credits > 0: - balance = { - "remaining": round(max(0.0, total_credits - total_usage), 6), - "funded": round(total_credits, 6), - "spent": round(max(0.0, total_usage), 6), - "currency": "USD", - "estimated": False, - } - except OpenRouterError as error: - if str(error) == "Could not reach the OpenRouter API": - retry_advised = True - else: - print(f"omarchy-agent-usage-openrouter: credits lookup failed ({error})", file=sys.stderr) - - record = base_record( - ready=True, - hasLocalStats=complete, - hasPromptStats=has_prompt_stats, - limits=key_limit_meter(key_info), - ) - if retry_advised: - record["retryAdvised"] = True - record.update(stats) - if balance: - record["balance"] = balance - if not has_data and not balance and not record["limits"]: - record["ready"] = False - record["usageStatusText"] = "No OpenRouter usage found" - record["authHelpText"] = AUTH_HELP - return record + write_json(limits_cache(), {"fetchedAtMs": fetched_at, "limits": limits}) + except Exception as exc: + print(f"omarchy-agent-usage-openrouter: could not write limits cache ({exc})", file=sys.stderr) + except OpenRouterError as error: + cached = read_fresh_json(limits_cache(), float("inf")) or {} + limits = cached.get("limits") if isinstance(cached.get("limits"), list) else [] + stale, fetched_at = True, number(cached.get("fetchedAtMs")) + status = "OpenRouter unavailable" + help_text = "OpenRouter rejected the API key. " + AUTH_HELP if str(error) == "HTTP_401" else str(error) + + # Balance: total purchased minus total used. A transport failure here + # must not take the token stats down -- retryAdvised asks the shell for + # one sooner retry instead of waiting out the full refresh interval. + balance = None + retry_advised = False + try: + credits = api_get("/credits", key).get("data", {}) + total_credits = float(credits.get("total_credits") or 0) + total_usage = float(credits.get("total_usage") or 0) + if total_credits > 0: + balance = { + "remaining": round(max(0.0, total_credits - total_usage), 6), + "funded": round(total_credits, 6), + "spent": round(max(0.0, total_usage), 6), + "currency": "USD", + "estimated": False, + } + except OpenRouterError as error: + if str(error) == "Could not reach the OpenRouter API": + retry_advised = True + else: + print(f"omarchy-agent-usage-openrouter: credits lookup failed ({error})", file=sys.stderr) + + record = base_record( + ready=True, + tierLabel="Prepaid", + hasLocalStats=complete, + hasPromptStats=has_prompt_stats, + limits=limits, + limitsStale=stale, + limitsFetchedAt=fetched_at, + usageStatusText=status, + authHelpText=help_text, + ) + if retry_advised: + record["retryAdvised"] = True + record.update(stats) + if balance: + record["balance"] = balance + if not has_data and not balance and not record["limits"] and not status: + record["ready"] = False + record["usageStatusText"] = "No OpenRouter usage found" + record["authHelpText"] = AUTH_HELP + return record def main() -> int: - ap = argparse.ArgumentParser(description="Print the OpenRouter usage record as JSON") - # --force rescans everything and rewrites the cache; --limits-only - # reuses a recent local scan and only the balance/limits probe is fresh. - # The flags exist so every collector accepts the same invocation. - ap.add_argument("--force", action="store_true", help="rescan transcripts and re-probe limits, ignoring caches") - ap.add_argument("--limits-only", action="store_true", help="reuse any recent transcript scan; only the limits probe must be fresh") - args = ap.parse_args() - try: - record = scan(args) - except OpenRouterError as error: - record = base_record(usageStatusText="OpenRouter unavailable", authHelpText=str(error)) - except Exception as error: - record = base_record(usageStatusText="OpenRouter unavailable", authHelpText="OpenRouter scan failed") - print(f"omarchy-agent-usage-openrouter: {type(error).__name__}", file=sys.stderr) - print(json.dumps(record, separators=(",", ":"))) - return 0 + ap = argparse.ArgumentParser(description="Print the OpenRouter usage record as JSON") + # --force rescans everything and rewrites the cache; --limits-only + # reuses a recent local scan and only the balance/limits probe is fresh. + # The flags exist so every collector accepts the same invocation. + ap.add_argument("--force", action="store_true", help="rescan transcripts and re-probe limits, ignoring caches") + ap.add_argument("--limits-only", action="store_true", help="reuse any recent transcript scan; only the limits probe must be fresh") + args = ap.parse_args() + try: + record = scan(args) + except OpenRouterError as error: + record = base_record(usageStatusText="OpenRouter unavailable", authHelpText=str(error)) + except Exception as error: + record = base_record(usageStatusText="OpenRouter unavailable", authHelpText="OpenRouter scan failed") + print(f"omarchy-agent-usage-openrouter: {type(error).__name__}", file=sys.stderr) + print(json.dumps(record, separators=(",", ":"))) + return 0 if __name__ == "__main__": - raise SystemExit(main()) + raise SystemExit(main()) diff --git a/manual/17-ai.md b/manual/17-ai.md index aa45804c9de..6d43a8f2fca 100644 --- a/manual/17-ai.md +++ b/manual/17-ai.md @@ -38,7 +38,7 @@ There are terminal shortcuts too: `a` runs the default agent inline in the curre ### The agents panel -The top bar has an agents icon. Before you've set up an agent, its panel walks you through signing in to Claude, Codex, or Grok, and the first one becomes your default agent. After that, the panel puts every subscription on one page: each agent and account with its plan and the percentage used of its 5-hour session and weekly limits (or the remaining prepaid balance), under a line that rotates through what you've used across all of them, like tokens this week and your most used model. Under _Make something_ it offers starter prompts for a new theme, plugin, or app that hand the task to your default agent. Apps are built the way Omarchy's own Hype, Monologue, and Omacut are, in C++ and Qt Quick, following the omarchy-app agent skill, and the compiler and Qt pieces they need come installed. Claude Code, Codex, Grok, and Fireworks are covered out of the box, as is Antigravity with its 5-hour and weekly limits, and Hermes shows its sessions and tokens, from every profile, once you've used it. Provider plans and keys you use through pi, omp, or opencode show up as well: the Z.ai GLM Coding Plan, Ollama Cloud, and the MiniMax Token Plan with their 5-hour and weekly limits. +The top bar has an agents icon. Before you've set up an agent, its panel walks you through signing in to Claude, Codex, or Grok, and the first one becomes your default agent. After that, the panel puts every subscription on one page: each agent and account with its plan and the percentage used of its 5-hour session and weekly limits (or the remaining prepaid balance), under a line that rotates through what you've used across all of them, like tokens this week and your most used model. Under _Make something_ it offers starter prompts for a new theme, plugin, or app that hand the task to your default agent. Apps are built the way Omarchy's own Hype, Monologue, and Omacut are, in C++ and Qt Quick, following the omarchy-app agent skill, and the compiler and Qt pieces they need come installed. Claude Code, Codex, Grok, and Fireworks are covered out of the box, as is Antigravity with its 5-hour and weekly limits, and Hermes shows its sessions and tokens, from every profile, once you've used it. Provider plans and keys you use through pi, omp, or opencode show up as well: the Z.ai GLM Coding Plan, Ollama Cloud, and the MiniMax Token Plan with their 5-hour and weekly limits, and OpenRouter with its prepaid balance and any credit limit on the key. Left-click the bar icon for the panel, right-click to launch your default agent. The usage records behind it are regenerated every 15 minutes by `omarchy agent usage-update`, and the panel can even merge usage from your other machines via a synced folder. See the README under `$OMARCHY_PATH/shell/plugins/agents/` for the full settings. diff --git a/shell/plugins/agents/README.md b/shell/plugins/agents/README.md index b36d5c443ef..465eb850351 100644 --- a/shell/plugins/agents/README.md +++ b/shell/plugins/agents/README.md @@ -186,7 +186,8 @@ history from the local opencode scan to the account-global `/activity` API. The collector reads the prepaid ledger from `GET /api/v1/credits` (`remaining = total_credits - total_usage`) and reports it as a live, non-estimated balance. A key with a credit limit additionally gets a -draining key-limit meter from `GET /api/v1/key`. +draining key-limit meter from `GET /api/v1/key`, labeled with how often the +limit starts over; a failed lookup keeps the last meter, dimmed as stale. Token history comes in two tiers, selected automatically by key type: diff --git a/test/shell.d/agent-usage-openrouter-scanner-test.sh b/test/shell.d/agent-usage-openrouter-scanner-test.sh index ea3b2680310..d9a54b15c66 100755 --- a/test/shell.d/agent-usage-openrouter-scanner-test.sh +++ b/test/shell.d/agent-usage-openrouter-scanner-test.sh @@ -26,7 +26,7 @@ export TZ="UTC" # 1. Without credentials: valid record, ready=false no_key=$("$COLLECTOR") -[[ $(jq -r '.id + ":" + (.ready | tostring)' <<<"$no_key") == "openrouter:false" ]] || +[[ $(jq -r '.id + ":" + (.ready | tostring) + ":" + .tierLabel' <<<"$no_key") == "openrouter:false:" ]] || fail "prints a valid record without credentials" "$no_key" pass "prints a valid record without credentials" @@ -94,7 +94,7 @@ check "$(jq -r '.activityModels | join(",")' <<<"$result")" "gpt-4o-mini,y-model check "$(jq -r '.activityOldDay | length' <<<"$result")" "2" "activity tracks active dates" check "$(jq -r .keyMeterEmpty <<<"$result")" "true" "no meter without key limit" check "$(jq -r '.keyMeter[0].percent' <<<"$result")" "0.6" "key meter drains toward empty" -check "$(jq -r '.keyMeter[0].resetsAt' <<<"$result")" "weekly" "key meter maps limit_reset" +check "$(jq -r '.keyMeter[0].label + ":" + .keyMeter[0].resetsAt' <<<"$result")" "Weekly key limit:" "key meter names its period, not a reset time" # 3. Tier-2 path sets scope=account + hasPromptStats=false (stub api_get) scope_check=$(python3 - "$COLLECTOR" <<'PY' @@ -133,4 +133,38 @@ check "$(jq -r .hasPromptStats <<<"$scope_check")" "false" "tier-2 hides prompt check "$(jq -r .balance <<<"$scope_check")" "9.0" "tier-2 keeps balance from /credits" check "$(jq -r '.called | join(",")' <<<"$scope_check")" "/activity,/credits,/key" "tier-2 probes activity + key + credits" +# 4. A failed /key lookup keeps the last good meter, marked stale. +stale_check=$(python3 - "$COLLECTOR" <<'PY' +import importlib.machinery, importlib.util, json, os, sys +loader = importlib.machinery.SourceFileLoader("or_collector3", sys.argv[1]) +spec = importlib.util.spec_from_loader(loader.name, loader) +mod = importlib.util.module_from_spec(spec) +loader.exec_module(mod) + +state = {"up": True} +def fake_api_get(path, key): + if not state["up"]: + raise mod.OpenRouterTransportError("Could not reach the OpenRouter API") + if path == "/key": + return {"data": {"limit": 10, "limit_remaining": 4, "limit_reset": "monthly"}} + if path == "/credits": + return {"data": {"total_credits": 10, "total_usage": 1}} + raise AssertionError(path) +mod.api_get = fake_api_get + +class Args: force = True; limits_only = False +os.environ["OPENROUTER_API_KEY"] = "sk-test" +os.environ["OPENROUTER_MANAGEMENT_KEY"] = "" +live = mod.scan(Args()) +state["up"] = False +stale = mod.scan(Args()) +pick = lambda r: {"stale": r["limitsStale"], "fetchedAt": r["limitsFetchedAt"], "percent": r["limits"][0]["percent"] if r["limits"] else None, "ready": r["ready"], "tier": r["tierLabel"]} +print(json.dumps({"live": pick(live), "stale": pick(stale), "status": stale["usageStatusText"]})) +PY +) +check "$(jq -c '.live | {stale, percent, tier}' <<<"$stale_check")" '{"stale":false,"percent":0.6,"tier":"Prepaid"}' "a live key meter is fresh" +check "$(jq -c '.stale | {stale, percent, ready}' <<<"$stale_check")" '{"stale":true,"percent":0.6,"ready":true}' "a failed lookup keeps the last meter, marked stale" +check "$(jq -r '.stale.fetchedAt == .live.fetchedAt and .live.fetchedAt > 0' <<<"$stale_check")" "true" "stale meter keeps when it was fetched" +check "$(jq -r .status <<<"$stale_check")" "OpenRouter unavailable" "a failed lookup says so" + pass "ALL OPENROUTER SCANNER TESTS PASSED" From 7a1b2630ac2ca5e6e000104003076534fe435598 Mon Sep 17 00:00:00 2001 From: Alexander Holbreich Date: Fri, 2 Oct 2026 06:14:53 -0400 Subject: [PATCH 13/19] Add a DeepSeek usage collector for the agents panel Counts shared DeepSeek days once and keeps refreshes quiet. --- bin/omarchy-agent-usage-deepseek | 476 ++++++++++++++++++ shell/plugins/agents/README.md | 8 + shell/plugins/agents/assets/deepseek.svg | 4 + shell/plugins/agents/manifest.json | 3 +- .../agent-usage-deepseek-scanner-test.sh | 214 ++++++++ 5 files changed, 704 insertions(+), 1 deletion(-) create mode 100755 bin/omarchy-agent-usage-deepseek create mode 100644 shell/plugins/agents/assets/deepseek.svg create mode 100755 test/shell.d/agent-usage-deepseek-scanner-test.sh diff --git a/bin/omarchy-agent-usage-deepseek b/bin/omarchy-agent-usage-deepseek new file mode 100755 index 00000000000..71efa7c2667 --- /dev/null +++ b/bin/omarchy-agent-usage-deepseek @@ -0,0 +1,476 @@ +#!/usr/bin/python3 +# omarchy:summary=Print the DeepSeek usage record as JSON +# omarchy:args=[--force] [--limits-only] +# omarchy:hidden=true +"""Collect DeepSeek usage into one display-ready JSON record. + +Local stats come from pi/omp session files whose assistant messages ran on the +deepseek provider, and from opencode sessions that ran on a deepseek provider. +DeepSeek billing is prepaid, so the record carries the live credit balance from +GET https://api.deepseek.com/user/balance when a key is available, plus the +optional funded/spent detail derived from ~/.config/omarchy/agents/deepseek.json +(the same shape the fireworks collector uses). The agents panel only ever reads +the JSON this prints. + +Token accounting deliberately matches pi's own totals: input + output + +cacheRead + cacheWrite, with the separate "reasoning" field left out, exactly +like the codex collector treats pi sessions. +""" + +from __future__ import annotations + +import argparse +import json +import os +import shutil +import subprocess +import sys +import urllib.error +import urllib.request +from datetime import datetime, timedelta, timezone +from decimal import Decimal, InvalidOperation +from pathlib import Path +from typing import Any + +AGENT_ID = "deepseek" +AGENT_NAME = "DeepSeek" +API_BASE_URL = "https://api.deepseek.com" + + +class DeepSeekError(Exception): + pass + + +def number(value: Any) -> int: + try: + return int(value or 0) + except Exception: + return 0 + + +def model_name(raw: Any) -> str: + value = str(raw or AGENT_ID) + return value if value else AGENT_ID + + +def local_day(value: Any) -> str: + if value is None: + return datetime.now().strftime("%Y-%m-%d") + if isinstance(value, (int, float)): + # pi message timestamps are milliseconds; other clients use 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 runtime_env() -> dict[str, str]: + 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: str) -> str | None: + return shutil.which(name, path=ENV.get("PATH")) + + +def empty_bucket() -> dict[str, int]: + return { + "inputTokens": 0, + "outputTokens": 0, + "cacheReadInputTokens": 0, + "cacheCreationInputTokens": 0, + } + + +def empty_stats() -> dict[str, Any]: + return { + "todayPrompts": 0, + "todaySessions": 0, + "todayTotalTokens": 0, + "todayTokensByModel": {}, + "recentDays": [], + "totalPrompts": 0, + "totalSessions": 0, + "activeDays": 0, + "activeDates": [], + "modelUsage": {}, + } + + +def base_record(**overrides: Any) -> dict[str, Any]: + record: dict[str, Any] = { + "schemaVersion": 1, + "id": AGENT_ID, + "name": AGENT_NAME, + "updatedAt": datetime.now(timezone.utc).isoformat(), + "ready": True, + "hasLocalStats": True, + # Prompt and session counts are real: every pi session message carries + # them, so today's tooltip may show "N prompts · M sessions". + "hasPromptStats": True, + "tierLabel": "Prepaid", + "usageStatusText": "", + "authHelpText": "", + "limits": [], + } + record.update(empty_stats()) + record.update(overrides) + return record + + +class Stats: + """Token and activity totals shared by every local source. + + Sessions and days are sets, so a day active in both pi and opencode counts + once toward activeDays. + """ + + def __init__(self) -> None: + now = datetime.now() + self.today = now.strftime("%Y-%m-%d") + self.recent_dates = [ + (now - timedelta(days=offset)).strftime("%Y-%m-%d") + for offset in range(6, -1, -1) + ] + self.recent = {day: 0 for day in self.recent_dates} + self.today_prompts = 0 + self.today_total_tokens = 0 + self.today_sessions: set[str] = set() + self.today_tokens_by_model: dict[str, int] = {} + self.total_prompts = 0 + self.total_sessions: set[str] = set() + self.active_days: set[str] = set() + self.model_usage: dict[str, dict[str, int]] = {} + + def add(self, session: str, day: str, model: str, input_tokens: int, output_tokens: int, cache_read: int, cache_write: int) -> None: + bucket = self.model_usage.setdefault(model, empty_bucket()) + bucket["inputTokens"] += input_tokens + bucket["outputTokens"] += output_tokens + bucket["cacheReadInputTokens"] += cache_read + bucket["cacheCreationInputTokens"] += cache_write + + total_tokens = input_tokens + output_tokens + cache_read + cache_write + self.total_prompts += 1 + self.total_sessions.add(session) + self.active_days.add(day) + if day in self.recent: + self.recent[day] += total_tokens + if day == self.today: + self.today_prompts += 1 + self.today_sessions.add(session) + self.today_total_tokens += total_tokens + self.today_tokens_by_model[model] = self.today_tokens_by_model.get(model, 0) + total_tokens + + def record(self) -> dict[str, Any]: + return { + "todayPrompts": self.today_prompts, + "todaySessions": len(self.today_sessions), + "todayTotalTokens": self.today_total_tokens, + "todayTokensByModel": self.today_tokens_by_model, + "recentDays": [{"date": day, "messageCount": self.recent[day]} for day in self.recent_dates], + "totalPrompts": self.total_prompts, + "totalSessions": len(self.total_sessions), + "activeDays": len(self.active_days), + "activeDates": sorted(self.active_days), + "modelUsage": self.model_usage, + } + + +def pi_session_roots() -> list[Path]: + return [ + Path.home() / ".pi" / "agent" / "sessions", + Path.home() / ".omp" / "agent" / "sessions", + ] + + +def scan_pi_sessions(stats: Stats) -> None: + """Add pi/omp session messages that ran on the deepseek provider.""" + seen_messages: set[str] = set() + + rg = find_command("rg") or "rg" + for root in pi_session_roots(): + if not root.exists(): + continue + try: + proc = subprocess.Popen( + [rg, "--json", "-e", r'"provider"\s*:\s*"deepseek"', str(root)], + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + text=True, + errors="replace", + env=ENV, + ) + except FileNotFoundError: + continue + + 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_messages: + continue + seen_messages.add(message_key) + message = entry.get("message") or {} + if message.get("role") != "assistant": + continue + if str(message.get("provider") or "") != "deepseek": + 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")) + stats.add(path, day, model_name(message.get("model")), input_tokens, output_tokens, cache_read, cache_write) + + try: + proc.wait(timeout=1) + except Exception: + proc.kill() + + +def opencode_db_path() -> Path: + data_home = Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) + return data_home / "opencode" / "opencode.db" + + +def scan_opencode_sessions(stats: Stats) -> None: + """Add opencode messages on a deepseek provider, mirroring the codex scan. + + Read-only: opencode may be writing right now. A scan cut short by a + database error keeps what it already read. + """ + db = opencode_db_path() + if not db.is_file(): + return + + import sqlite3 + + 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") + # The json_extract conditions are the authority; the LIKE gates are pure + # acceleration that skip rows whose JSON cannot contain the pair, the same + # trade the codex collector makes on this table. + for session_id, raw in conn.execute( + "SELECT session_id, data FROM message" + " WHERE data LIKE '%\"role\"%:%\"assistant\"%'" + " AND data LIKE '%\"providerID\"%:%\"deepseek\"%'" + " AND CASE WHEN json_valid(data) THEN json_extract(data, '$.role') END = 'assistant'" + " AND CASE WHEN json_valid(data) THEN json_extract(data, '$.providerID') END = 'deepseek'" + ): + try: + entry = json.loads(raw) + if not isinstance(entry, dict) or entry.get("role") != "assistant": + continue + if str(entry.get("providerID") or "") != "deepseek": + 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]) + stats.add("opencode:" + str(session_id), day, model, input_tokens, output_tokens, cache_read, cache_write) + except Exception: + continue + except sqlite3.Error: + pass + finally: + conn.close() + + +# -------------------------------------------------------------- credentials + +def read_key_from_dict(parsed: Any, provider: str) -> str: + if not isinstance(parsed, dict): + return "" + entry = parsed.get(provider) + if not isinstance(entry, dict): + return "" + return str(entry.get("key") or entry.get("apiKey") or entry.get("api_key") or "").strip() + + +def read_json_key(path: Path, provider: str) -> str: + try: + return read_key_from_dict(json.loads(path.read_text()), provider) + except (OSError, json.JSONDecodeError): + return "" + + +def api_key() -> str: + env_key = str(os.environ.get("DEEPSEEK_API_KEY", "")).strip() + if env_key: + return env_key + key = read_json_key(Path.home() / ".pi" / "agent" / "auth.json", "deepseek") + if key: + return key + data_home = Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) + return read_json_key(data_home / "opencode" / "auth.json", "deepseek") + + +def config_path() -> Path: + config_home = Path(os.environ.get("XDG_CONFIG_HOME") or (Path.home() / ".config")) + return config_home / "omarchy" / "agents" / "deepseek.json" + + +def read_config() -> dict[str, Any]: + try: + parsed = json.loads(config_path().read_text()) + return parsed if isinstance(parsed, dict) else {} + except (OSError, json.JSONDecodeError): + return {} + + +# ----------------------------------------------------------------- balance + +def fetch_balance(api_base_url: str, key: str) -> dict[str, Any]: + """GET /user/balance: total_balance, granted_balance, topped_up_balance.""" + url = api_base_url.rstrip("/") + "/user/balance" + request = urllib.request.Request( + url, + headers={"Authorization": "Bearer " + key, "Accept": "application/json"}, + ) + try: + with urllib.request.urlopen(request, timeout=15) as response: + payload = json.load(response) + except urllib.error.HTTPError as error: + raise DeepSeekError(f"DeepSeek balance endpoint returned HTTP {error.code}") + except urllib.error.URLError as error: + raise DeepSeekError("Could not reach the DeepSeek API") from error + except (json.JSONDecodeError, TimeoutError) as error: + raise DeepSeekError("DeepSeek returned an invalid balance response") from error + + if not isinstance(payload, dict): + raise DeepSeekError("DeepSeek returned an invalid balance response") + infos = payload.get("balance_infos") + if not isinstance(infos, list) or not infos or not isinstance(infos[0], dict): + raise DeepSeekError("DeepSeek balance response carried no balance_infos") + first = infos[0] + + def money(value: Any) -> Decimal: + try: + return Decimal(str(value or "0")) + except InvalidOperation: + return Decimal("0") + + return { + "remaining": money(first.get("total_balance")), + "granted": money(first.get("granted_balance")), + "toppedUp": money(first.get("topped_up_balance")), + "currency": str(first.get("currency") or "CNY"), + } + + +def balance_record(key: str, api_base_url: str) -> dict[str, Any]: + config = read_config() + try: + funded = Decimal(str(config.get("fundedAmount") or "0")) + except InvalidOperation: + funded = Decimal("0") + if not funded.is_finite() or funded < 0: + funded = Decimal("0") + + live = fetch_balance(api_base_url, key) + remaining = live["remaining"] + spent = max(Decimal("0"), funded - remaining) if funded > 0 else Decimal("0") + return { + "remaining": float(max(Decimal("0"), remaining)), + "funded": float(funded), + "spent": float(spent), + "currency": str(live["currency"] or "CNY"), + # funded/spent derive from the configured fundedAmount, not DeepSeek's + # own ledger, so they are labelled estimates; remaining is live. + "estimated": funded > 0, + } + + +def scan(api_base_url: str) -> dict[str, Any]: + stats = Stats() + scan_pi_sessions(stats) + scan_opencode_sessions(stats) + record = base_record() + record.update(stats.record()) + + # Without a key the record carries local stats only; the panel leaves the + # balance section out rather than reporting an error. + key = api_key() + if key: + try: + record["balance"] = balance_record(key, api_base_url) + except DeepSeekError as error: + record["usageStatusText"] = "Balance unavailable" + record["authHelpText"] = str(error) + return record + + +def main() -> int: + parser = argparse.ArgumentParser(description="Print the DeepSeek usage record as JSON") + # Stats are local file scans and the balance is one API call; 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.add_argument("--api-base-url", default=os.environ.get("DEEPSEEK_API_BASE_URL", API_BASE_URL)) + args = parser.parse_args() + + try: + record = scan(args.api_base_url) + except Exception as error: + record = base_record(usageStatusText="DeepSeek unavailable", authHelpText="DeepSeek usage scan failed") + print(f"omarchy-agent-usage-deepseek: {type(error).__name__}", file=sys.stderr) + print(json.dumps(record, separators=(",", ":"))) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/shell/plugins/agents/README.md b/shell/plugins/agents/README.md index 465eb850351..6a86368baef 100644 --- a/shell/plugins/agents/README.md +++ b/shell/plugins/agents/README.md @@ -80,6 +80,7 @@ light surfaces — and the bar glyph stands in when there is none. | `ollama` | Ollama Cloud's usage endpoint (5-hour session + 7-day weekly) | pi and omp sessions on the `ollama-cloud` provider | | `minimax` | MiniMax Token Plan API (5-hour session + 7-day weekly) | Not available from the quota API | | `openrouter` | Prepaid balance from `/credits`, key-limit meter from `/key` | opencode sessions on the OpenRouter provider, or the `/activity` billing API with a management key | +| `deepseek` | Live prepaid balance from `GET https://api.deepseek.com/user/balance`, with the optional `fundedAmount` estimate for the funded-versus-spent line | pi/omp session files whose assistant messages ran on the deepseek provider, plus opencode sessions on a deepseek provider | When `~/.local/state/omarchy/agents/accounts/.json` registers more than one account, the `claude`, `codex`, and `grok` records @@ -146,6 +147,13 @@ when MiniMax is signed in there. It uses the region and resource URL saved by `mmx`, or the global `https://api.minimax.io` endpoint by default; `MINIMAX_BASE_URL` overrides both. +DeepSeek reads `DEEPSEEK_API_KEY` first, then the key pi +stores in `~/.pi/agent/auth.json`, then the key opencode stores in its own +`auth.json` when DeepSeek is signed in there. Its `fundedAmount` estimate +lives in `~/.config/omarchy/agents/deepseek.json`, the same shape the +fireworks collector uses; without it the panel shows the live remaining +credit and skips the funded-versus-spent line. + ### Fireworks balance The collector first asks the account's `:getBalance` endpoint for the real diff --git a/shell/plugins/agents/assets/deepseek.svg b/shell/plugins/agents/assets/deepseek.svg new file mode 100644 index 00000000000..57f9b482977 --- /dev/null +++ b/shell/plugins/agents/assets/deepseek.svg @@ -0,0 +1,4 @@ + + DeepSeek + + diff --git a/shell/plugins/agents/manifest.json b/shell/plugins/agents/manifest.json index 2969b557a30..44bfdb98dcb 100644 --- a/shell/plugins/agents/manifest.json +++ b/shell/plugins/agents/manifest.json @@ -28,7 +28,8 @@ "zai": { "enabled": true }, "ollama": { "enabled": true }, "minimax": { "enabled": true }, - "openrouter": { "enabled": true } + "openrouter": { "enabled": true }, + "deepseek": { "enabled": true } }, "refreshIntervalSec": 900, "syncMode": "Off", diff --git a/test/shell.d/agent-usage-deepseek-scanner-test.sh b/test/shell.d/agent-usage-deepseek-scanner-test.sh new file mode 100755 index 00000000000..b003f82a60b --- /dev/null +++ b/test/shell.d/agent-usage-deepseek-scanner-test.sh @@ -0,0 +1,214 @@ +#!/bin/bash + +source "$(dirname "$0")/base-test.sh" + +require_command jq +require_command python3 + +TEST_HOME=$(mktemp -d) +trap 'rm -rf "$TEST_HOME"' EXIT + +timestamp="$(date +%Y-%m-%d)T12:00:00Z" + +# ---------------------------------------------------------------- pi and omp + +PI_HOME=$(mktemp -d) +trap 'rm -rf "$TEST_HOME" "$PI_HOME"' EXIT +mkdir -p "$PI_HOME/.pi/agent/sessions/project" "$PI_HOME/.omp/agent/sessions/project" + +cat >"$PI_HOME/.pi/agent/sessions/project/pi.jsonl" <"$PI_HOME/.omp/agent/sessions/project/omp.jsonl" </dev/null) + +[[ $(jq -r '.todayTotalTokens' <<<"$result") == "49" ]] || + fail "DeepSeek collector counts usage from pi and omp sessions" "$result" +pass "DeepSeek collector counts usage from pi and omp sessions" + +[[ $(jq -c '.modelUsage' <<<"$result") == '{"deepseek-pi":{"inputTokens":10,"outputTokens":4,"cacheReadInputTokens":3,"cacheCreationInputTokens":2},"deepseek-omp":{"inputTokens":20,"outputTokens":5,"cacheReadInputTokens":4,"cacheCreationInputTokens":1}}' ]] || + fail "DeepSeek collector filters pi and omp sessions to DeepSeek providers" "$result" +pass "DeepSeek collector filters pi and omp sessions to DeepSeek providers" + +[[ $(jq -r '"\(.todayPrompts):\(.todaySessions)"' <<<"$result") == "2:2" ]] || + fail "DeepSeek collector counts each message once per session" "$result" +pass "DeepSeek collector counts each message once per session" + +# Without a key the record must still be complete and hidden-by-default: +# the panel shows local stats and simply omits the balance section. +[[ $(jq -r '.id + ":" + (.ready|tostring) + ":" + (.hasPromptStats|tostring) + ":" + (has("balance")|tostring)' <<<"$result") == "deepseek:true:true:false" ]] || + fail "DeepSeek collector prints a valid record without a key" "$result" +pass "DeepSeek collector prints a valid record without a key" + +# ---------------------------------------------------------------- opencode + +OPENCODE_HOME=$(mktemp -d) +trap 'rm -rf "$TEST_HOME" "$PI_HOME" "$OPENCODE_HOME"' EXIT + +python3 - "$OPENCODE_HOME/.local/share/opencode/opencode.db" <<'PY' +import json +import sqlite3 +import sys +import time +from pathlib import Path + +db = Path(sys.argv[1]) +db.parent.mkdir(parents=True, exist_ok=True) +conn = sqlite3.connect(db) +conn.execute("CREATE TABLE message (id text PRIMARY KEY, session_id text NOT NULL, time_created integer NOT NULL, time_updated integer NOT NULL, data text NOT NULL)") +now_ms = int(time.time() * 1000) + +def message(id, provider, model, input=0, output=0, reasoning=0, read=0, write=0): + return (id, "ses_1", now_ms, now_ms, json.dumps({ + "role": "assistant", + "providerID": provider, + "modelID": model, + "tokens": {"input": input, "output": output, "reasoning": reasoning, "cache": {"read": read, "write": write}}, + "time": {"created": now_ms}, + })) + +conn.executemany( + "INSERT INTO message VALUES (?, ?, ?, ?, ?)", + [ + message("m-1", "deepseek", "deepseek-model", input=100, output=20, reasoning=5, read=30, write=2), + message("m-2", "openai", "gpt-model", input=999, output=999), + message("m-3", "anthropic", "claude-model", input=999, output=999), + ], +) +conn.commit() +conn.close() +PY + +result=$(HOME="$OPENCODE_HOME" XDG_DATA_HOME="$OPENCODE_HOME/.local/share" \ + "$ROOT/bin/omarchy-agent-usage-deepseek" 2>/dev/null) + +[[ $(jq -r '.todayTotalTokens' <<<"$result") == "157" ]] || + fail "DeepSeek collector counts opencode messages on a DeepSeek provider" "$result" +pass "DeepSeek collector counts opencode messages on a DeepSeek provider" + +[[ $(jq -c '.modelUsage["deepseek-model"]' <<<"$result") == '{"inputTokens":100,"outputTokens":25,"cacheReadInputTokens":30,"cacheCreationInputTokens":2}' ]] || + fail "DeepSeek collector keeps reasoning with output and cache apart in opencode" "$result" +pass "DeepSeek collector keeps reasoning with output and cache apart in opencode" + +# A day active in both pi and opencode is still one active day. +mkdir -p "$OPENCODE_HOME/.pi/agent/sessions/project" +cp "$PI_HOME/.pi/agent/sessions/project/pi.jsonl" "$OPENCODE_HOME/.pi/agent/sessions/project/" +result=$(HOME="$OPENCODE_HOME" XDG_DATA_HOME="$OPENCODE_HOME/.local/share" \ + "$ROOT/bin/omarchy-agent-usage-deepseek" 2>/dev/null) + +[[ $(jq -r '"\(.todayTotalTokens):\(.totalSessions):\(.activeDays)"' <<<"$result") == "176:2:1" ]] || + fail "DeepSeek collector counts a day shared by pi and opencode once" "$result" +pass "DeepSeek collector counts a day shared by pi and opencode once" + +# ---------------------------------------------------------------- balance + +result=$(HOME="$TEST_HOME" XDG_DATA_HOME="$TEST_HOME/.local/share" python3 - "$ROOT/bin/omarchy-agent-usage-deepseek" "$TEST_HOME/.config" <<'PY' +import importlib.machinery +import importlib.util +import json +import os +import sys +from decimal import Decimal +from pathlib import Path + +collector_path = str(Path(sys.argv[1])) +os.environ["XDG_CONFIG_HOME"] = sys.argv[2] + +loader = importlib.machinery.SourceFileLoader("deepseek_collector", collector_path) +spec = importlib.util.spec_from_loader(loader.name, loader) +scanner = importlib.util.module_from_spec(spec) +loader.exec_module(scanner) +real_fetch_balance = scanner.fetch_balance + +config_dir = Path(sys.argv[2]) / "omarchy" / "agents" +config_dir.mkdir(parents=True, exist_ok=True) +(config_dir / "deepseek.json").write_text(json.dumps({"fundedAmount": 20})) + +scanner.fetch_balance = lambda base_url, key: { + "remaining": Decimal("8.15"), + "granted": Decimal("0"), + "toppedUp": Decimal("8.15"), + "currency": "USD", +} +balance = scanner.balance_record("test-key", "https://example.invalid") +summary = {"withConfig": balance} + +scanner.fetch_balance = lambda base_url, key: (_ for _ in ()).throw( + scanner.DeepSeekError("balance endpoint unavailable") +) +scanner.api_key = lambda: "test-key" +record = scanner.scan("https://example.invalid") +summary["onFailure"] = { + "hasBalance": "balance" in record, + "status": record["usageStatusText"], + "help": record["authHelpText"], +} +scanner.fetch_balance = real_fetch_balance + +# The real request must carry the bearer key and hit /user/balance. +captured = {} + +class FakeResponse: + def __init__(self, payload): + self.payload = payload + + def __enter__(self): + return self + + def __exit__(self, *args): + return False + + def read(self): + return json.dumps({ + "is_available": True, + "balance_infos": [{ + "currency": "CNY", + "total_balance": "42.50", + "granted_balance": "2.50", + "topped_up_balance": "40.00", + }], + }).encode("utf-8") + +import urllib.request +real_urlopen = urllib.request.urlopen +urllib.request.urlopen = lambda request, timeout=15: ( + captured.update({"url": request.full_url, "auth": request.get_header("Authorization")}) + or FakeResponse(None) +) +try: + live = scanner.fetch_balance("https://api.deepseek.com", "key-123") +finally: + urllib.request.urlopen = real_urlopen + +summary["liveShape"] = { + "remaining": float(live["remaining"]), + "granted": float(live["granted"]), + "toppedUp": float(live["toppedUp"]), + "currency": live["currency"], +} +summary["requestShape"] = {"url": captured.get("url"), "auth": captured.get("auth")} +print(json.dumps(summary, separators=(",", ":"))) +PY +) + +[[ $(jq -c '.withConfig | {remaining, funded, spent, currency, estimated}' <<<"$result") == '{"remaining":8.15,"funded":20.0,"spent":11.85,"currency":"USD","estimated":true}' ]] || + fail "DeepSeek collector derives the estimated ledger from fundedAmount" "$result" +pass "DeepSeek collector derives the estimated ledger from fundedAmount" + +[[ $(jq -c '.onFailure' <<<"$result") == '{"hasBalance":false,"status":"Balance unavailable","help":"balance endpoint unavailable"}' ]] || + fail "DeepSeek collector reports a failed balance lookup in the record" "$result" +pass "DeepSeek collector reports a failed balance lookup in the record" + +[[ $(jq -c '.liveShape' <<<"$result") == '{"remaining":42.5,"granted":2.5,"toppedUp":40.0,"currency":"CNY"}' ]] || + fail "DeepSeek collector parses the live balance response" "$result" +pass "DeepSeek collector parses the live balance response" + +[[ $(jq -c '.requestShape' <<<"$result") == '{"url":"https://api.deepseek.com/user/balance","auth":"Bearer key-123"}' ]] || + fail "DeepSeek collector requests /user/balance with the bearer key" "$result" +pass "DeepSeek collector requests /user/balance with the bearer key" From 52e2412e60226fd117cfc6e2fce881ea622fedc6 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Fri, 2 Oct 2026 06:15:37 -0400 Subject: [PATCH 14/19] Skip the DeepSeek record on machines without DeepSeek The record claimed to be ready with a Prepaid plan for everyone, which put a DeepSeek section in every panel. Ready and the plan now follow the key. The test no longer reaches DeepSeek when DEEPSEEK_API_KEY is set, and the manual mentions DeepSeek. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-deepseek | 10 ++++++---- manual/17-ai.md | 2 +- test/shell.d/agent-usage-deepseek-scanner-test.sh | 13 ++++++++++++- 3 files changed, 19 insertions(+), 6 deletions(-) diff --git a/bin/omarchy-agent-usage-deepseek b/bin/omarchy-agent-usage-deepseek index 71efa7c2667..fedc514e346 100755 --- a/bin/omarchy-agent-usage-deepseek +++ b/bin/omarchy-agent-usage-deepseek @@ -124,12 +124,12 @@ def base_record(**overrides: Any) -> dict[str, Any]: "id": AGENT_ID, "name": AGENT_NAME, "updatedAt": datetime.now(timezone.utc).isoformat(), - "ready": True, + "ready": False, "hasLocalStats": True, # Prompt and session counts are real: every pi session message carries # them, so today's tooltip may show "N prompts · M sessions". "hasPromptStats": True, - "tierLabel": "Prepaid", + "tierLabel": "", "usageStatusText": "", "authHelpText": "", "limits": [], @@ -441,10 +441,12 @@ def scan(api_base_url: str) -> dict[str, Any]: record = base_record() record.update(stats.record()) - # Without a key the record carries local stats only; the panel leaves the - # balance section out rather than reporting an error. + # Without a key the record carries local stats only, with no plan, so the + # panel leaves out the balance section, and skips the record altogether on + # a machine that has never used DeepSeek. key = api_key() if key: + record.update(ready=True, tierLabel="Prepaid") try: record["balance"] = balance_record(key, api_base_url) except DeepSeekError as error: diff --git a/manual/17-ai.md b/manual/17-ai.md index 6d43a8f2fca..4dc6a94aa71 100644 --- a/manual/17-ai.md +++ b/manual/17-ai.md @@ -38,7 +38,7 @@ There are terminal shortcuts too: `a` runs the default agent inline in the curre ### The agents panel -The top bar has an agents icon. Before you've set up an agent, its panel walks you through signing in to Claude, Codex, or Grok, and the first one becomes your default agent. After that, the panel puts every subscription on one page: each agent and account with its plan and the percentage used of its 5-hour session and weekly limits (or the remaining prepaid balance), under a line that rotates through what you've used across all of them, like tokens this week and your most used model. Under _Make something_ it offers starter prompts for a new theme, plugin, or app that hand the task to your default agent. Apps are built the way Omarchy's own Hype, Monologue, and Omacut are, in C++ and Qt Quick, following the omarchy-app agent skill, and the compiler and Qt pieces they need come installed. Claude Code, Codex, Grok, and Fireworks are covered out of the box, as is Antigravity with its 5-hour and weekly limits, and Hermes shows its sessions and tokens, from every profile, once you've used it. Provider plans and keys you use through pi, omp, or opencode show up as well: the Z.ai GLM Coding Plan, Ollama Cloud, and the MiniMax Token Plan with their 5-hour and weekly limits, and OpenRouter with its prepaid balance and any credit limit on the key. +The top bar has an agents icon. Before you've set up an agent, its panel walks you through signing in to Claude, Codex, or Grok, and the first one becomes your default agent. After that, the panel puts every subscription on one page: each agent and account with its plan and the percentage used of its 5-hour session and weekly limits (or the remaining prepaid balance), under a line that rotates through what you've used across all of them, like tokens this week and your most used model. Under _Make something_ it offers starter prompts for a new theme, plugin, or app that hand the task to your default agent. Apps are built the way Omarchy's own Hype, Monologue, and Omacut are, in C++ and Qt Quick, following the omarchy-app agent skill, and the compiler and Qt pieces they need come installed. Claude Code, Codex, Grok, and Fireworks are covered out of the box, as is Antigravity with its 5-hour and weekly limits, and Hermes shows its sessions and tokens, from every profile, once you've used it. Provider plans and keys you use through pi, omp, or opencode show up as well: the Z.ai GLM Coding Plan, Ollama Cloud, and the MiniMax Token Plan with their 5-hour and weekly limits, OpenRouter with its prepaid balance and any credit limit on the key, and DeepSeek with its prepaid balance. Left-click the bar icon for the panel, right-click to launch your default agent. The usage records behind it are regenerated every 15 minutes by `omarchy agent usage-update`, and the panel can even merge usage from your other machines via a synced folder. See the README under `$OMARCHY_PATH/shell/plugins/agents/` for the full settings. diff --git a/test/shell.d/agent-usage-deepseek-scanner-test.sh b/test/shell.d/agent-usage-deepseek-scanner-test.sh index b003f82a60b..91226a78aa1 100755 --- a/test/shell.d/agent-usage-deepseek-scanner-test.sh +++ b/test/shell.d/agent-usage-deepseek-scanner-test.sh @@ -8,6 +8,9 @@ require_command python3 TEST_HOME=$(mktemp -d) trap 'rm -rf "$TEST_HOME"' EXIT +# A key in the environment would send the collector to DeepSeek. +unset DEEPSEEK_API_KEY DEEPSEEK_API_BASE_URL + timestamp="$(date +%Y-%m-%d)T12:00:00Z" # ---------------------------------------------------------------- pi and omp @@ -42,10 +45,18 @@ pass "DeepSeek collector counts each message once per session" # Without a key the record must still be complete and hidden-by-default: # the panel shows local stats and simply omits the balance section. -[[ $(jq -r '.id + ":" + (.ready|tostring) + ":" + (.hasPromptStats|tostring) + ":" + (has("balance")|tostring)' <<<"$result") == "deepseek:true:true:false" ]] || +[[ $(jq -r '.id + ":" + (.ready|tostring) + ":" + .tierLabel + ":" + (.hasPromptStats|tostring) + ":" + (has("balance")|tostring)' <<<"$result") == "deepseek:false::true:false" ]] || fail "DeepSeek collector prints a valid record without a key" "$result" pass "DeepSeek collector prints a valid record without a key" +# Nothing used and no key: a record the panel skips. +EMPTY_HOME=$(mktemp -d) +result=$(HOME="$EMPTY_HOME" XDG_DATA_HOME="$EMPTY_HOME/.local/share" "$ROOT/bin/omarchy-agent-usage-deepseek" 2>/dev/null) +rm -rf "$EMPTY_HOME" +[[ $(jq -c '{ready, tierLabel, totalPrompts, activeDays, balance: has("balance")}' <<<"$result") == '{"ready":false,"tierLabel":"","totalPrompts":0,"activeDays":0,"balance":false}' ]] || + fail "DeepSeek collector prints a record the panel skips without DeepSeek" "$result" +pass "DeepSeek collector prints a record the panel skips without DeepSeek" + # ---------------------------------------------------------------- opencode OPENCODE_HOME=$(mktemp -d) From 3fda50cbde9724debf428869556c2f8bd8cee06f Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Fri, 2 Oct 2026 06:15:49 -0400 Subject: [PATCH 15/19] Describe the agents panel without listing every agent Co-Authored-By: Claude Opus 5.5 --- shell/plugins/agents/manifest.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/shell/plugins/agents/manifest.json b/shell/plugins/agents/manifest.json index 44bfdb98dcb..156c79ed970 100644 --- a/shell/plugins/agents/manifest.json +++ b/shell/plugins/agents/manifest.json @@ -5,7 +5,7 @@ "version": "1.0.0", "author": "Omarchy", "license": "MIT", - "description": "Claude Code, Codex, Grok, Fireworks, Hermes, and Antigravity usage, limits, and pace in a native Omarchy bar panel.", + "description": "Usage, limits, pace, and prepaid balances for every AI agent and provider in a native Omarchy bar panel.", "kinds": ["bar-widget"], "activation": "on-demand", "entryPoints": { From a6f33974819f8acb9b3d098df42fa26842fd21fd Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Fri, 2 Oct 2026 06:29:41 -0400 Subject: [PATCH 16/19] Treat OpenRouter read timeouts as transport failures A server stalling on headers or body raised TimeoutError past the OpenRouterError handlers, so the record went blank instead of keeping its stats and last key-limit meter and asking for a retry. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-openrouter | 4 ++ .../agent-usage-openrouter-scanner-test.sh | 39 +++++++++++++++++++ 2 files changed, 43 insertions(+) diff --git a/bin/omarchy-agent-usage-openrouter b/bin/omarchy-agent-usage-openrouter index 0bbcd05651b..f63467e499d 100755 --- a/bin/omarchy-agent-usage-openrouter +++ b/bin/omarchy-agent-usage-openrouter @@ -176,6 +176,10 @@ def api_get(path: str, key: str) -> dict: raise OpenRouterError(f"OpenRouter API returned HTTP {error.code}") except urllib.request.URLError as error: raise OpenRouterTransportError("Could not reach the OpenRouter API") from error + # A server that stalls on its headers or body times out the read itself, + # outside URLError; that is a network failure like any other. + except (TimeoutError, ConnectionError) as error: + raise OpenRouterTransportError("Could not reach the OpenRouter API") from error except (json.JSONDecodeError, ValueError) as error: raise OpenRouterError("OpenRouter returned an invalid response") from error diff --git a/test/shell.d/agent-usage-openrouter-scanner-test.sh b/test/shell.d/agent-usage-openrouter-scanner-test.sh index d9a54b15c66..b8f19435ff7 100755 --- a/test/shell.d/agent-usage-openrouter-scanner-test.sh +++ b/test/shell.d/agent-usage-openrouter-scanner-test.sh @@ -167,4 +167,43 @@ check "$(jq -c '.stale | {stale, percent, ready}' <<<"$stale_check")" '{"stale": check "$(jq -r '.stale.fetchedAt == .live.fetchedAt and .live.fetchedAt > 0' <<<"$stale_check")" "true" "stale meter keeps when it was fetched" check "$(jq -r .status <<<"$stale_check")" "OpenRouter unavailable" "a failed lookup says so" +# 5. A read that times out (headers or body) is a transport failure: the +# record keeps its stats and last meter instead of going blank. +timeout_check=$(python3 - "$COLLECTOR" <<'PY' +import importlib.machinery, importlib.util, json, os, sys +loader = importlib.machinery.SourceFileLoader("or_collector4", sys.argv[1]) +spec = importlib.util.spec_from_loader(loader.name, loader) +mod = importlib.util.module_from_spec(spec) +loader.exec_module(mod) + +class StalledBody: + def __enter__(self): return self + def __exit__(self, *args): return False + def read(self, *args): raise TimeoutError("The read operation timed out") + +results = {} +for name, opener in ( + ("headers", lambda request, timeout=None: (_ for _ in ()).throw(TimeoutError("timed out"))), + ("body", lambda request, timeout=None: StalledBody()), +): + mod.urllib.request.urlopen = opener + try: + mod.api_get("/key", "sk-test") + results[name] = "no error" + except mod.OpenRouterTransportError: + results[name] = "transport" + except Exception as error: + results[name] = type(error).__name__ + +class Args: force = True; limits_only = False +os.environ["OPENROUTER_API_KEY"] = "sk-test" +os.environ["OPENROUTER_MANAGEMENT_KEY"] = "" +record = mod.scan(Args()) +results["record"] = {"ready": record["ready"], "stale": record["limitsStale"], "meter": bool(record["limits"]), "retry": record.get("retryAdvised", False), "prompts": record["totalPrompts"]} +print(json.dumps(results)) +PY +) +check "$(jq -r '.headers + ":" + .body' <<<"$timeout_check")" "transport:transport" "read timeouts are transport failures" +check "$(jq -c .record <<<"$timeout_check")" '{"ready":true,"stale":true,"meter":true,"retry":true,"prompts":1}' "a timed-out probe keeps stats and the last meter and asks for a retry" + pass "ALL OPENROUTER SCANNER TESTS PASSED" From 7008b8539e5051f47e2d0092d0105634a848432d Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Fri, 2 Oct 2026 06:30:49 -0400 Subject: [PATCH 17/19] Collect OpenRouter usage with only a management key A management key without an inference key returned an empty record before /activity was asked. Now either key collects: the management key alone brings the 30-day history and the balance, with no key-limit meter, since that belongs to an inference key. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-openrouter | 54 +++++++++++-------- shell/plugins/agents/README.md | 3 +- .../agent-usage-openrouter-scanner-test.sh | 29 ++++++++++ 3 files changed, 62 insertions(+), 24 deletions(-) diff --git a/bin/omarchy-agent-usage-openrouter b/bin/omarchy-agent-usage-openrouter index f63467e499d..836dc6cdea1 100755 --- a/bin/omarchy-agent-usage-openrouter +++ b/bin/omarchy-agent-usage-openrouter @@ -478,10 +478,32 @@ def limits_cache() -> Path: return cache_root() / "openrouter-limits.json" +# The key-limit meter as (limits, stale, fetchedAtMs, status, help). A failed +# lookup keeps the last good meter, marked stale, and leaves the token stats +# standing. +def key_meter(key: str) -> tuple[list[dict], bool, int, str, str]: + try: + key_info = api_get("/key", key).get("data", {}) + except OpenRouterError as error: + cached = read_fresh_json(limits_cache(), float("inf")) or {} + limits = cached.get("limits") if isinstance(cached.get("limits"), list) else [] + help_text = "OpenRouter rejected the API key. " + AUTH_HELP if str(error) == "HTTP_401" else str(error) + return limits, True, number(cached.get("fetchedAtMs")), "OpenRouter unavailable", help_text + + limits = key_limit_meter(key_info if isinstance(key_info, dict) else {}) + fetched_at = round(time.time() * 1000) + try: + write_json(limits_cache(), {"fetchedAtMs": fetched_at, "limits": limits}) + except Exception as exc: + print(f"omarchy-agent-usage-openrouter: could not write limits cache ({exc})", file=sys.stderr) + return limits, False, fetched_at, "", "" + + def scan(args: Any) -> dict[str, Any]: config = read_config() key = api_key(config) - if not key: + manager = management_key(config) + if not key and not manager: return base_record(usageStatusText="OpenRouter unavailable", authHelpText=AUTH_HELP) today = datetime.now().astimezone() @@ -489,7 +511,7 @@ def scan(args: Any) -> dict[str, Any]: # Tier 2 first: with a management key the 30-day account history covers # every machine and client. Without one (or a rejected one) fall back to # the local opencode DB scan below. - managed = activity_stats(management_key(config), today) + managed = activity_stats(manager, today) if managed is not None: stats = managed complete = True @@ -502,26 +524,12 @@ def scan(args: Any) -> dict[str, Any]: has_prompt_stats = True has_data = number(stats.get("totalPrompts")) > 0 or bool(stats.get("modelUsage")) - # The key-limit meter. A failed lookup keeps the last good meter, marked - # stale, and leaves the token stats standing. - limits, stale, fetched_at = [], False, 0 - status = help_text = "" - try: - key_info = api_get("/key", key).get("data", {}) - if not isinstance(key_info, dict): - key_info = {} - limits = key_limit_meter(key_info) - fetched_at = round(time.time() * 1000) - try: - write_json(limits_cache(), {"fetchedAtMs": fetched_at, "limits": limits}) - except Exception as exc: - print(f"omarchy-agent-usage-openrouter: could not write limits cache ({exc})", file=sys.stderr) - except OpenRouterError as error: - cached = read_fresh_json(limits_cache(), float("inf")) or {} - limits = cached.get("limits") if isinstance(cached.get("limits"), list) else [] - stale, fetched_at = True, number(cached.get("fetchedAtMs")) - status = "OpenRouter unavailable" - help_text = "OpenRouter rejected the API key. " + AUTH_HELP if str(error) == "HTTP_401" else str(error) + # The meter describes an inference key's own credit limit, so with only a + # management key there is none to show. + if key: + limits, stale, fetched_at, status, help_text = key_meter(key) + else: + limits, stale, fetched_at, status, help_text = [], False, 0, "", "" # Balance: total purchased minus total used. A transport failure here # must not take the token stats down -- retryAdvised asks the shell for @@ -529,7 +537,7 @@ def scan(args: Any) -> dict[str, Any]: balance = None retry_advised = False try: - credits = api_get("/credits", key).get("data", {}) + credits = api_get("/credits", key or manager).get("data", {}) total_credits = float(credits.get("total_credits") or 0) total_usage = float(credits.get("total_usage") or 0) if total_credits > 0: diff --git a/shell/plugins/agents/README.md b/shell/plugins/agents/README.md index 6a86368baef..d6bf6679bd4 100644 --- a/shell/plugins/agents/README.md +++ b/shell/plugins/agents/README.md @@ -189,7 +189,8 @@ opencode stores in `~/.local/share/opencode/auth.json` when OpenRouter is signed in there, then `apiKey` in `~/.config/omarchy/agents/openrouter.json`. A management key in `OPENROUTER_MANAGEMENT_KEY` (or `managementKey` in that file, from `openrouter.ai/settings/management-keys`) switches the token -history from the local opencode scan to the account-global `/activity` API. +history from the local opencode scan to the account-global `/activity` API; +on its own it also reads the balance, with no key-limit meter. The collector reads the prepaid ledger from `GET /api/v1/credits` (`remaining = total_credits - total_usage`) and reports it as a live, diff --git a/test/shell.d/agent-usage-openrouter-scanner-test.sh b/test/shell.d/agent-usage-openrouter-scanner-test.sh index b8f19435ff7..a6b93b71f3d 100755 --- a/test/shell.d/agent-usage-openrouter-scanner-test.sh +++ b/test/shell.d/agent-usage-openrouter-scanner-test.sh @@ -206,4 +206,33 @@ PY check "$(jq -r '.headers + ":" + .body' <<<"$timeout_check")" "transport:transport" "read timeouts are transport failures" check "$(jq -c .record <<<"$timeout_check")" '{"ready":true,"stale":true,"meter":true,"retry":true,"prompts":1}' "a timed-out probe keeps stats and the last meter and asks for a retry" +# 6. A management key alone still collects: the 30-day history and the +# balance, with no inference key's meter to show. +manager_check=$(python3 - "$COLLECTOR" <<'PY' +import importlib.machinery, importlib.util, json, os, sys +loader = importlib.machinery.SourceFileLoader("or_collector5", sys.argv[1]) +spec = importlib.util.spec_from_loader(loader.name, loader) +mod = importlib.util.module_from_spec(spec) +loader.exec_module(mod) + +calls = [] +def fake_api_get(path, key): + calls.append(f"{path}={key}") + if path == "/activity": + return {"data": [{"date": "2026-09-08", "model": "a/b", "prompt_tokens": 10, "completion_tokens": 5}]} + if path == "/credits": + return {"data": {"total_credits": 10, "total_usage": 4}} + raise AssertionError(path) +mod.api_get = fake_api_get + +class Args: force = True; limits_only = False +os.environ["OPENROUTER_API_KEY"] = "" +os.environ["OPENROUTER_MANAGEMENT_KEY"] = "sk-mgmt" +record = mod.scan(Args()) +print(json.dumps({"ready": record["ready"], "scope": record.get("scope"), "balance": record.get("balance", {}).get("remaining"), "limits": record["limits"], "status": record["usageStatusText"], "calls": sorted(calls)})) +PY +) +check "$(jq -c '{ready, scope, balance, limits, status}' <<<"$manager_check")" '{"ready":true,"scope":"account","balance":6.0,"limits":[],"status":""}' "a management key alone collects history and balance" +check "$(jq -r '.calls | join(",")' <<<"$manager_check")" "/activity=sk-mgmt,/credits=sk-mgmt" "a management key alone skips the inference key's meter" + pass "ALL OPENROUTER SCANNER TESTS PASSED" From 5421e2e08124513203f83fe53c2503ebeabf73d5 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Fri, 2 Oct 2026 06:31:18 -0400 Subject: [PATCH 18/19] Reuse Ollama's session scan only for its day and model filter The cached scan was reused by age alone, so just after midnight yesterday's tokens showed as today's, and a changed modelPrefix kept showing excluded models for up to 15 minutes. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-ollama | 8 +++++--- .../agent-usage-ollama-scanner-test.sh | 19 +++++++++++++++++++ 2 files changed, 24 insertions(+), 3 deletions(-) diff --git a/bin/omarchy-agent-usage-ollama b/bin/omarchy-agent-usage-ollama index 9ed1503da0c..f4db9d81939 100755 --- a/bin/omarchy-agent-usage-ollama +++ b/bin/omarchy-agent-usage-ollama @@ -169,11 +169,13 @@ def scan_sessions(max_age_seconds: float, model_prefix: str) -> dict[str, Any] | Path.home() / ".omp" / "agent" / "sessions", ] cache_file = cache_root() / "ollama-pi-sessions.json" + # A scan only answers for the day it ran on (its today* counts would be + # yesterday's after midnight) and for the model filter it applied. + today = local_date_string() cached = read_fresh_json(cache_file, max_age_seconds) - if cached is not None: + if cached is not None and cached.get("scanDate") == today and cached.get("modelPrefix") == model_prefix: return cached.get("stats") - today = local_date_string() recent_dates = recent_date_strings() recent = {day: {"date": day, "messageCount": 0} for day in recent_dates} @@ -263,7 +265,7 @@ def scan_sessions(max_age_seconds: float, model_prefix: str) -> dict[str, Any] | "activeDays": len(active_days), "activeDates": sorted(active_days), } - write_json(cache_file, {"stats": stats}) + write_json(cache_file, {"scanDate": today, "modelPrefix": model_prefix, "stats": stats}) return stats diff --git a/test/shell.d/agent-usage-ollama-scanner-test.sh b/test/shell.d/agent-usage-ollama-scanner-test.sh index feaa7ce0d22..7437010c7bf 100755 --- a/test/shell.d/agent-usage-ollama-scanner-test.sh +++ b/test/shell.d/agent-usage-ollama-scanner-test.sh @@ -59,6 +59,25 @@ result=$(HOME="$PI_HOME" XDG_CONFIG_HOME="$PI_HOME/.config" XDG_CACHE_HOME="$PI_ fail "Ollama collector honors the configured model prefix" "$result" pass "Ollama collector honors the configured model prefix" +# A cached scan is reused only on the day it ran and with the filter it +# applied: a changed prefix, or yesterday's scan, means scanning again. +result=$(HOME="$PI_HOME" XDG_CONFIG_HOME="$PI_HOME/.config" XDG_CACHE_HOME="$PI_HOME/.cache" \ + OLLAMA_API_KEY="" "$ROOT/bin/omarchy-agent-usage-ollama" --limits-only) +[[ $(jq -c '.modelUsage | keys' <<<"$result") == '["kimi-k3"]' ]] || + fail "Ollama collector reuses a scan made with the same filter today" "$result" +rm "$PI_HOME/.config/omarchy/agents/ollama.json" +result=$(HOME="$PI_HOME" XDG_CONFIG_HOME="$PI_HOME/.config" XDG_CACHE_HOME="$PI_HOME/.cache" \ + OLLAMA_API_KEY="" "$ROOT/bin/omarchy-agent-usage-ollama" --limits-only) +[[ $(jq -c '.modelUsage | keys' <<<"$result") == '["deepseek-v4-pro","kimi-k3"]' ]] || + fail "Ollama collector rescans when the model filter changes" "$result" +cache_file="$PI_HOME/.cache/omarchy/agent-usage/ollama-pi-sessions.json" +jq -c '.scanDate = "2000-01-01" | .stats.todayTotalTokens = 999' "$cache_file" >"$cache_file.tmp" && mv "$cache_file.tmp" "$cache_file" +result=$(HOME="$PI_HOME" XDG_CONFIG_HOME="$PI_HOME/.config" XDG_CACHE_HOME="$PI_HOME/.cache" \ + OLLAMA_API_KEY="" "$ROOT/bin/omarchy-agent-usage-ollama" --limits-only) +[[ $(jq -r '.todayTotalTokens' <<<"$result") == "55" ]] || + fail "Ollama collector rescans past a scan from another day" "$result" +pass "Ollama collector reuses a session scan only for the same day and filter" + # probe_limits reaches Ollama, so the reader that interprets its answer is # exercised on its own: the collector loads as a module, and a recorded # payload stands in for the response. From 2ffdf9a4dd5fb2c6a3ce977b24ee3988f154adb9 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Fri, 2 Oct 2026 06:31:55 -0400 Subject: [PATCH 19/19] Drop cached MiniMax windows that have reset During an outage the fallback republished cached windows even after they had reset, so an exhausted session could stay on screen. Windows past their reset are now dropped; open ones and those without a reset time stay. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-minimax | 21 +++++++++++++++++++- test/shell.d/agent-usage-minimax-test.sh | 25 ++++++++++++++++++++++-- 2 files changed, 43 insertions(+), 3 deletions(-) diff --git a/bin/omarchy-agent-usage-minimax b/bin/omarchy-agent-usage-minimax index 71117f311c5..942041cea35 100755 --- a/bin/omarchy-agent-usage-minimax +++ b/bin/omarchy-agent-usage-minimax @@ -252,6 +252,25 @@ def fetch_quota(api_key: str, base_url: str) -> dict[str, Any]: return payload +# A window kept from an earlier check is dropped once it has reset, since +# what it said no longer holds; one with no reset time is kept. +def open_windows(limits: Any) -> list[dict[str, Any]]: + now = datetime.now(timezone.utc) + kept = [] + for limit in limits if isinstance(limits, list) else []: + if not isinstance(limit, dict): + continue + try: + resets = datetime.fromisoformat(str(limit.get("resetsAt") or "").replace("Z", "+00:00")) + except ValueError: + resets = None + if resets and resets.tzinfo is None: + resets = resets.replace(tzinfo=timezone.utc) + if not resets or resets > now: + kept.append(limit) + return kept + + # Nobody with a MiniMax key here: an empty record, which the panel skips. # With one, a failed check keeps the last good limits, marked stale. def collect(api_key: str, base_url: str) -> dict[str, Any]: @@ -264,7 +283,7 @@ def collect(api_key: str, base_url: str) -> dict[str, Any]: except MiniMaxError as error: cached = read_json(limits_cache()) or {} record.update( - limits=cached.get("limits") or [], + limits=open_windows(cached.get("limits")), limitsStale=True, limitsFetchedAt=int(cached.get("fetchedAtMs") or 0), usageStatusText="MiniMax limits unavailable", diff --git a/test/shell.d/agent-usage-minimax-test.sh b/test/shell.d/agent-usage-minimax-test.sh index a30acb560d8..37ff0f769a5 100755 --- a/test/shell.d/agent-usage-minimax-test.sh +++ b/test/shell.d/agent-usage-minimax-test.sh @@ -82,8 +82,13 @@ opencode_key, opencode_base_url = scanner.credentials() class WorkingClient: pass -scanner.fetch_quota = lambda api_key, base_url: payload -record = scanner.collect("mmx_test", "https://example.invalid") +# Windows still open, so the cached copy outlives the failed check below. +import copy, time +current = copy.deepcopy(payload) +current["model_remains"][0]["end_time"] = round((time.time() + 3600) * 1000) +current["model_remains"][0]["weekly_end_time"] = round((time.time() + 86400) * 1000) +scanner.fetch_quota = lambda api_key, base_url: current +record =scanner.collect("mmx_test", "https://example.invalid") def unreachable(api_key, base_url): raise scanner.MiniMaxError("Could not reach the MiniMax API") @@ -91,6 +96,17 @@ def unreachable(api_key, base_url): scanner.fetch_quota = unreachable stale = scanner.collect("mmx_test", "https://example.invalid") +# A cached window that has since reset is dropped; one still open, or with +# no reset time, stays. +from datetime import datetime, timedelta, timezone +now = datetime.now(timezone.utc) +scanner.write_json(scanner.limits_cache(), {"fetchedAtMs": 1, "limits": [ + {"label": "Session (5-hour)", "percent": 1.0, "resetsAt": (now - timedelta(hours=1)).isoformat()}, + {"label": "Weekly (7-day)", "percent": 0.5, "resetsAt": (now + timedelta(days=2)).isoformat()}, + {"label": "Other", "percent": 0.2, "resetsAt": ""}, +]}) +reset = scanner.collect("mmx_test", "https://example.invalid") + print(json.dumps({ "limits": limits, "savedKey": saved_key, @@ -100,6 +116,7 @@ print(json.dumps({ "opencodeKey": opencode_key, "record": record, "stale": stale, + "reset": reset, })) PY ) @@ -129,3 +146,7 @@ pass "MiniMax collector prints the display-ready record contract" [[ $(jq -r '.stale.limitsFetchedAt == .record.limitsFetchedAt and .record.limitsFetchedAt > 0' <<<"$result") == "true" ]] || fail "stale MiniMax limits keep when they were fetched" "$result" pass "MiniMax collector keeps the last limits, marked stale, after a failed check" + +[[ $(jq -c '[.reset.limits[].label]' <<<"$result") == '["Weekly (7-day)","Other"]' ]] || + fail "MiniMax collector drops kept windows that have reset" "$result" +pass "MiniMax collector drops kept windows that have reset"