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 001/546] 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 817bb4b775511a6341891592e6d08894ddab24c2 Mon Sep 17 00:00:00 2001 From: Mark Xiong Date: Sat, 19 Sep 2026 17:09:37 -0700 Subject: [PATCH 002/546] Add a Cursor collector to the agents panel Cursor keeps no token counts on disk, so the collector reads the same dashboard RPCs the Cursor CLI itself calls: GetCurrentPeriodUsage for the meters and the billing cycle, GetPlanInfo for the plan name, and one GetAggregatedUsageEvents call per day for tokens by day and by model. The meters are Cursor's own percentages rather than spend divided by the allowance: provider bonus usage is spend with no allowance behind it, so dollars over the plan limit read far past 100% while Cursor's own meters do not. The on-demand row stays a real ratio of dollars against the limit the account sets. Co-Authored-By: Claude Opus 5 --- bin/omarchy-agent-usage-cursor | 337 +++++++++++++++++++ shell/plugins/agents/README.md | 20 +- shell/plugins/agents/assets/cursor-light.svg | 6 + shell/plugins/agents/assets/cursor.svg | 6 + shell/plugins/agents/manifest.json | 3 +- test/shell.d/agent-usage-cursor-test.sh | 154 +++++++++ 6 files changed, 524 insertions(+), 2 deletions(-) create mode 100755 bin/omarchy-agent-usage-cursor create mode 100644 shell/plugins/agents/assets/cursor-light.svg create mode 100644 shell/plugins/agents/assets/cursor.svg create mode 100755 test/shell.d/agent-usage-cursor-test.sh diff --git a/bin/omarchy-agent-usage-cursor b/bin/omarchy-agent-usage-cursor new file mode 100755 index 00000000000..ee632871b77 --- /dev/null +++ b/bin/omarchy-agent-usage-cursor @@ -0,0 +1,337 @@ +#!/usr/bin/python3 +# omarchy:summary=Print the Cursor usage record as JSON +# omarchy:args=[--force] [--limits-only] +# omarchy:hidden=true +"""Collect Cursor usage into one display-ready JSON record. + +Cursor keeps no token counts on disk, so everything here comes from the same +dashboard RPCs the Cursor CLI itself calls: the plan and the period meters +from GetCurrentPeriodUsage and GetPlanInfo, tokens per model and per day from +GetAggregatedUsageEvents. The agents panel only ever reads the JSON this +prints. +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +import urllib.error +import urllib.request +from concurrent.futures import ThreadPoolExecutor +from datetime import date, datetime, time, timedelta, timezone +from pathlib import Path +from typing import Any + +AGENT_ID = "cursor" +AGENT_NAME = "Cursor" +AUTH_HELP = "Run `cursor-agent login` to authenticate." +API_BASE_URL = "https://api2.cursor.sh" +SERVICE = "aiserver.v1.DashboardService" +REQUEST_TIMEOUT = 20 +RECENT_DAYS = 7 + + +class CursorError(Exception): + pass + + +def number(value: Any) -> int: + try: + return max(0, round(float(value or 0))) + except (TypeError, ValueError): + return 0 + + +def fraction(value: Any) -> float: + # Cursor reports whole percentages; the panel's meters take 0..1. + try: + return max(0.0, float(value or 0) / 100) + except (TypeError, ValueError): + return 0.0 + + +def state_home() -> Path: + return Path(os.environ.get("XDG_STATE_HOME") or (Path.home() / ".local" / "state")) + + +def record_path() -> Path: + return state_home() / "omarchy" / "agents" / "usage" / f"{AGENT_ID}.json" + + +def default_auth_path() -> Path: + config_home = Path(os.environ.get("XDG_CONFIG_HOME") or (Path.home() / ".config")) + return config_home / "cursor" / "auth.json" + + +def read_token(path: Path) -> str: + """The credential the Cursor CLI stores, or the key the user exported.""" + key = str(os.environ.get("CURSOR_API_KEY") or "").strip() + if key: + return key + if not path.is_file(): + return "" + try: + return str(json.loads(path.read_text()).get("accessToken") or "").strip() + except (OSError, ValueError): + return "" + + +def iso(epoch_millis: Any) -> str | None: + millis = number(epoch_millis) + if millis <= 0: + return None + return datetime.fromtimestamp(millis / 1000, timezone.utc).isoformat() + + +def empty_bucket() -> dict[str, int]: + return { + "inputTokens": 0, + "outputTokens": 0, + "cacheReadInputTokens": 0, + "cacheCreationInputTokens": 0, + } + + +def empty_stats() -> dict[str, Any]: + return { + "hasLocalStats": False, + "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, + # Dashboard numbers are account-global, not machine-local: every synced + # device reports the same truth, so aggregation must not sum them. + "scope": "account", + # The dashboard reports usage events, never prompt or session counts; the + # panel keeps those numbers out of today's tooltip when this is false. + "hasPromptStats": False, + "tierLabel": "", + "usageStatusText": "", + "authHelpText": "", + "limits": [], + } + record.update(empty_stats()) + record.update(overrides) + return record + + +class CursorClient: + def __init__(self, token: str, api_base_url: str) -> None: + self.token = token + self.api_base_url = api_base_url.rstrip("/") + + def call(self, method: str, payload: dict[str, Any]) -> dict[str, Any]: + request = urllib.request.Request( + f"{self.api_base_url}/{SERVICE}/{method}", + data=json.dumps(payload).encode("utf-8"), + headers={ + "Content-Type": "application/json", + "Connect-Protocol-Version": "1", + "Authorization": f"Bearer {self.token}", + "x-cursor-client-type": "cli", + }, + method="POST", + ) + try: + with urllib.request.urlopen(request, timeout=REQUEST_TIMEOUT) as response: + body = response.read().decode("utf-8") + except urllib.error.HTTPError as error: + if error.code in (401, 403): + raise CursorError(AUTH_HELP) from error + raise CursorError(f"Cursor API returned {error.code}") from error + except (urllib.error.URLError, TimeoutError) as error: + raise CursorError("Cursor API unreachable") from error + try: + return json.loads(body or "{}") + except ValueError as error: + raise CursorError("Cursor API returned malformed JSON") from error + + def period(self) -> dict[str, Any]: + return self.call("GetCurrentPeriodUsage", {}) + + def plan_name(self) -> str: + try: + plan = self.call("GetPlanInfo", {}).get("planInfo") + except CursorError: + return "" + return str((plan or {}).get("planName") or "") + + def aggregations(self, start_millis: int, end_millis: int) -> list[Any]: + payload = {"startDate": str(start_millis), "endDate": str(end_millis)} + rows = self.call("GetAggregatedUsageEvents", payload).get("aggregations") + return rows if isinstance(rows, list) else [] + + +def day_windows(today: date, count: int) -> list[tuple[str, int, int]]: + """The last `count` local days, oldest first, as (date, startMs, endMs).""" + windows = [] + for offset in range(count - 1, -1, -1): + day = today - timedelta(days=offset) + start = datetime.combine(day, time.min).astimezone() + end = start + timedelta(days=1) + windows.append((day.isoformat(), int(start.timestamp() * 1000), int(end.timestamp() * 1000))) + return windows + + +def summarize_day(rows: list[Any]) -> dict[str, dict[str, int]]: + """One day's aggregations under the token names every collector prints.""" + usage: dict[str, dict[str, int]] = {} + for row in rows: + if not isinstance(row, dict): + continue + model = str(row.get("modelIntent") or "cursor") + bucket = usage.setdefault(model, empty_bucket()) + bucket["inputTokens"] += number(row.get("inputTokens")) + bucket["outputTokens"] += number(row.get("outputTokens")) + bucket["cacheReadInputTokens"] += number(row.get("cacheReadTokens")) + bucket["cacheCreationInputTokens"] += number(row.get("cacheWriteTokens")) + return usage + + +def summarize_usage(days: list[tuple[str, dict[str, dict[str, int]]]]) -> dict[str, Any]: + recent_days = [] + active_dates = [] + model_usage: dict[str, dict[str, int]] = {} + today_by_model: dict[str, int] = {} + + for index, (day, usage) in enumerate(days): + tokens = 0 + for model, bucket in usage.items(): + carried = model_usage.setdefault(model, empty_bucket()) + for name, value in bucket.items(): + carried[name] += value + tokens += sum(bucket.values()) + recent_days.append({"date": day, "messageCount": tokens}) + if tokens > 0: + active_dates.append(day) + if index == len(days) - 1: + today_by_model = {model: sum(bucket.values()) for model, bucket in usage.items() if sum(bucket.values()) > 0} + + return { + "hasLocalStats": bool(model_usage), + "todayTotalTokens": sum(today_by_model.values()), + "todayTokensByModel": today_by_model, + "recentDays": recent_days, + "activeDays": len(active_dates), + "activeDates": active_dates, + "modelUsage": model_usage, + } + + +def carried_stats() -> dict[str, Any]: + """--limits-only promises fresh limits only, so the token sections stand.""" + stats = empty_stats() + try: + previous = json.loads(record_path().read_text()) + except (OSError, ValueError): + return stats + if not isinstance(previous, dict): + return stats + for key in stats: + if key in previous: + stats[key] = previous[key] + return stats + + +def collect_stats(client: CursorClient, today: date) -> dict[str, Any]: + windows = day_windows(today, RECENT_DAYS) + with ThreadPoolExecutor(max_workers=4) as pool: + rows = list(pool.map(lambda window: client.aggregations(window[1], window[2]), windows)) + return summarize_usage([(window[0], summarize_day(day_rows)) for window, day_rows in zip(windows, rows)]) + + +def summarize_limits(period: dict[str, Any]) -> list[dict[str, Any]]: + """Cursor's own percentages, never a percentage derived from spend. + + The dollar figures do not divide into the allowance: bonus usage the + providers hand out is spend without an allowance behind it, so + totalSpend/limit reads well past 100% while Cursor's own meters do not. + """ + plan = period.get("planUsage") + plan = plan if isinstance(plan, dict) else {} + spend = period.get("spendLimitUsage") + spend = spend if isinstance(spend, dict) else {} + resets_at = iso(period.get("billingCycleEnd")) + limits = [] + + for label, key in ( + ("Included total", "totalPercentUsed"), + ("Auto models", "autoPercentUsed"), + ("Named models", "apiPercentUsed"), + ): + if key in plan: + limits.append({"label": label, "percent": fraction(plan.get(key)), "resetsAt": resets_at}) + + limit_cents = number(spend.get("individualLimit")) + if limit_cents > 0: + used_cents = number(spend.get("individualUsed")) + limits.append({ + "label": f"On-demand ${used_cents / 100:.2f} / ${limit_cents / 100:.2f}", + "percent": min(1.0, used_cents / limit_cents), + "resetsAt": resets_at, + }) + return limits + + +def scan(api_base_url: str, auth_path: Path, limits_only: bool) -> dict[str, Any]: + token = read_token(auth_path) + if not token: + return base_record(usageStatusText="Cursor unavailable", authHelpText=AUTH_HELP) + + client = CursorClient(token, api_base_url) + period = client.period() + record = base_record( + ready=True, + tierLabel=client.plan_name(), + authHelpText=AUTH_HELP, + limits=summarize_limits(period), + ) + if limits_only: + record.update(carried_stats()) + else: + record.update(collect_stats(client, datetime.now().astimezone().date())) + return record + + +def main() -> int: + parser = argparse.ArgumentParser(description="Print the Cursor usage record as JSON") + # Every call is a live API call, so there is no cache to force past; the + # flag exists 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("--auth-path", default=os.environ.get("CURSOR_AUTH_PATH", "")) + parser.add_argument("--api-base-url", default=os.environ.get("CURSOR_API_ENDPOINT", API_BASE_URL)) + args = parser.parse_args() + + auth_path = Path(args.auth_path).expanduser() if args.auth_path else default_auth_path() + try: + record = scan(args.api_base_url, auth_path, args.limits_only) + except CursorError as error: + record = base_record(usageStatusText="Cursor unavailable", authHelpText=str(error)) + except Exception as error: + record = base_record(usageStatusText="Cursor unavailable", authHelpText="Cursor usage scan failed") + print(f"omarchy-agent-usage-cursor: {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 c9daebba45e..2964f88c5bb 100644 --- a/shell/plugins/agents/README.md +++ b/shell/plugins/agents/README.md @@ -73,6 +73,7 @@ light surfaces — and the bar glyph stands in when there is none. | `claude` | Anthropic's OAuth usage endpoint (5-hour session + 7-day weekly) | `~/.claude/projects` transcripts, opencode sessions on an Anthropic provider, plus `stats-cache.json` and `history.jsonl` as fallback | | `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 | +| `cursor` | Cursor's dashboard RPCs: included, auto-model, named-model, and on-demand meters | the same dashboard RPCs, one call per day for the last week | | `fireworks` | Estimated prepaid balance: configured funding minus rated account costs | Fireworks billing API, grouped by day and model for the last 30 days | When `~/.local/state/omarchy/agents/accounts/.json` @@ -99,7 +100,24 @@ again. Fireworks reads `FIREWORKS_API_KEY` and `FIREWORKS_ACCOUNT_ID` first, then `~/.fireworks/auth.ini` (which `firectl set-api-key` creates), then the key opencode stores in `~/.local/share/opencode/auth.json` when Fireworks is -signed in there. +signed in there. Cursor reads `CURSOR_API_KEY` first, then the token +`cursor-agent login` stores in `~/.config/cursor/auth.json`, and honors +`CURSOR_API_ENDPOINT`. + +### Cursor + +Cursor's session files on disk carry no token counts, so this is the one +collector with nothing to fall back on: without a credential the tab says +"Cursor unavailable" rather than showing a week of zeros. Tokens by day and by +model cover the last seven days, one aggregation call per day. + +The meters are Cursor's own percentages — included total, auto-bucket models, +and named models — so they always agree with the Cursor dashboard. Spend is +never divided into an allowance here: bonus usage the model providers hand out +is spend with no allowance behind it, so dollars over the plan limit read far +past 100% while Cursor's own meters do not. The on-demand row is the exception +and is a real ratio: dollars spent against the on-demand limit the account +sets. ### Fireworks balance diff --git a/shell/plugins/agents/assets/cursor-light.svg b/shell/plugins/agents/assets/cursor-light.svg new file mode 100644 index 00000000000..ecf8c981852 --- /dev/null +++ b/shell/plugins/agents/assets/cursor-light.svg @@ -0,0 +1,6 @@ + + + + + + diff --git a/shell/plugins/agents/assets/cursor.svg b/shell/plugins/agents/assets/cursor.svg new file mode 100644 index 00000000000..8652a86f9c0 --- /dev/null +++ b/shell/plugins/agents/assets/cursor.svg @@ -0,0 +1,6 @@ + + + + + + diff --git a/shell/plugins/agents/manifest.json b/shell/plugins/agents/manifest.json index 7ca52484b12..60af1e5a2fd 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, Cursor, and Fireworks usage, limits, and pace in a native Omarchy bar panel.", "kinds": ["bar-widget"], "activation": "on-demand", "entryPoints": { @@ -22,6 +22,7 @@ "claude": { "enabled": true }, "codex": { "enabled": true }, "grok": { "enabled": true }, + "cursor": { "enabled": true }, "fireworks": { "enabled": true } }, "refreshIntervalSec": 900, diff --git a/test/shell.d/agent-usage-cursor-test.sh b/test/shell.d/agent-usage-cursor-test.sh new file mode 100755 index 00000000000..026a3387907 --- /dev/null +++ b/test/shell.d/agent-usage-cursor-test.sh @@ -0,0 +1,154 @@ +#!/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/.config/cursor" +cat >"$TEST_HOME/.config/cursor/auth.json" <<'EOF' +{ "accessToken": "cursor_test", "refreshToken": "cursor_refresh" } +EOF + +# Without a credential the collector must still print a full, hidden-by-default +# record: the update runner writes whatever valid JSON appears on stdout. +no_token=$(HOME="$TEST_HOME" XDG_CONFIG_HOME="$TEST_HOME/missing" XDG_STATE_HOME="$TEST_HOME/.local/state" \ + CURSOR_API_KEY="" "$ROOT/bin/omarchy-agent-usage-cursor") + +[[ $(jq -r '.id + ":" + (.ready | tostring) + ":" + (.hasPromptStats | tostring)' <<<"$no_token") == "cursor:false:false" ]] || + fail "Cursor collector prints a valid record without credentials" "$no_token" +pass "Cursor collector prints a valid record without credentials" + +result=$(python3 - "$ROOT/bin/omarchy-agent-usage-cursor" "$TEST_HOME/.config" "$TEST_HOME/.local/state" <<'PY' +import importlib.machinery +import importlib.util +import json +import os +import sys +import time +from datetime import datetime, timedelta +from pathlib import Path + +collector_path = str(Path(sys.argv[1])) +os.environ["XDG_CONFIG_HOME"] = sys.argv[2] +os.environ["XDG_STATE_HOME"] = sys.argv[3] + +# Day windows resolve in local time, so pin the zone or the fixtures below +# would shift by a day depending on where the test runs. +os.environ["TZ"] = "UTC" +time.tzset() +os.environ.pop("CURSOR_API_KEY", None) + +loader = importlib.machinery.SourceFileLoader("cursor_collector", collector_path) +spec = importlib.util.spec_from_loader(loader.name, loader) +collector = importlib.util.module_from_spec(spec) +loader.exec_module(collector) + +today = datetime.now().astimezone().date() +yesterday = today - timedelta(days=1) + +# Spend well past the included limit, the way a plan with provider bonus usage +# reports it: 609.86 spent against a 400.00 allowance, while Cursor's own +# meters say 18.76% / 11.81% / 100%. +PERIOD = { + "billingCycleEnd": "1792099926000", + "planUsage": { + "totalSpend": 60986, + "includedSpend": 40000, + "limit": 40000, + "autoPercentUsed": 11.81, + "apiPercentUsed": 100, + "totalPercentUsed": 18.76, + }, + "spendLimitUsage": {"individualLimit": "1000", "individualUsed": "951", "limitType": "user"}, +} + +AGGREGATIONS = { + today.isoformat(): [ + {"modelIntent": "default", "inputTokens": "10", "outputTokens": "5", + "cacheReadTokens": "100", "cacheWriteTokens": "1", "totalCents": None}, + ], + yesterday.isoformat(): [ + {"modelIntent": "default", "inputTokens": "2", "outputTokens": "3", + "cacheReadTokens": "4", "cacheWriteTokens": "0"}, + {"modelIntent": "kimi-k3-max", "inputTokens": "1", "outputTokens": "1", + "cacheReadTokens": "0", "cacheWriteTokens": "0"}, + ], +} + +calls = [] + + +class StubClient(collector.CursorClient): + def call(self, method, payload): + calls.append((method, payload)) + if method == "GetCurrentPeriodUsage": + return PERIOD + if method == "GetPlanInfo": + return {"planInfo": {"planName": "Ultra"}} + if method == "GetAggregatedUsageEvents": + start = datetime.fromtimestamp(int(payload["startDate"]) / 1000).astimezone().date() + return {"aggregations": AGGREGATIONS.get(start.isoformat(), [])} + raise AssertionError(f"unexpected method: {method}") + + +collector.CursorClient = StubClient +record = collector.scan("https://api.example", Path(os.environ["XDG_CONFIG_HOME"]) / "cursor" / "auth.json", False) + +windows = collector.day_windows(today, collector.RECENT_DAYS) +first_start = datetime.fromtimestamp(windows[0][1] / 1000).astimezone() + +# --limits-only keeps the token sections of the record already on disk. +record_path = collector.record_path() +record_path.parent.mkdir(parents=True, exist_ok=True) +record_path.write_text(json.dumps(record)) +carried = collector.scan("https://api.example", Path(os.environ["XDG_CONFIG_HOME"]) / "cursor" / "auth.json", True) + +print(json.dumps({ + "record": record, + "aggregationCalls": sum(1 for method, _ in calls if method == "GetAggregatedUsageEvents"), + "windowIsLocalMidnight": (first_start.hour, first_start.minute, first_start.second) == (0, 0, 0), + "carriedRecentDays": carried["recentDays"] == record["recentDays"], + "carriedSkipsAggregation": sum(1 for method, _ in calls if method == "GetAggregatedUsageEvents") == collector.RECENT_DAYS, +})) +PY +) + +[[ $(jq -c '.record | {schemaVersion, id, name, ready, hasPromptStats, scope, tierLabel, usageStatusText}' <<<"$result") == '{"schemaVersion":1,"id":"cursor","name":"Cursor","ready":true,"hasPromptStats":false,"scope":"account","tierLabel":"Ultra","usageStatusText":""}' ]] || + fail "Cursor collector prints the display-ready record contract" "$result" +pass "Cursor collector prints the display-ready record contract" + +[[ $(jq -c '[.record.limits[] | {label, percent: (.percent * 10000 | round)}]' <<<"$result") == '[{"label":"Included total","percent":1876},{"label":"Auto models","percent":1181},{"label":"Named models","percent":10000},{"label":"On-demand $9.51 / $10.00","percent":9510}]' ]] || + fail "Cursor collector reports Cursor's own meters, not spend over the allowance" "$result" +pass "Cursor collector reports Cursor's own meters, not spend over the allowance" + +[[ $(jq -r '.record.limits[0].resetsAt' <<<"$result") == "2026-10-15T21:32:06+00:00" ]] || + fail "Cursor collector dates the reset from the billing cycle end" "$result" +pass "Cursor collector dates the reset from the billing cycle end" + +[[ $(jq -r '.aggregationCalls' <<<"$result") == "7" ]] || + fail "Cursor collector asks for seven daily windows" "$result" +pass "Cursor collector asks for seven daily windows" + +[[ $(jq -r '.windowIsLocalMidnight' <<<"$result") == "true" ]] || + fail "Cursor collector requests windows that start at local midnight" "$result" +pass "Cursor collector requests windows that start at local midnight" + +[[ $(jq -c '[.record.recentDays[] | .messageCount]' <<<"$result") == '[0,0,0,0,0,11,116]' ]] || + fail "Cursor collector builds the seven-day token series" "$result" +pass "Cursor collector builds the seven-day token series" + +[[ $(jq -c '.record.modelUsage' <<<"$result") == '{"default":{"inputTokens":12,"outputTokens":8,"cacheReadInputTokens":104,"cacheCreationInputTokens":1},"kimi-k3-max":{"inputTokens":1,"outputTokens":1,"cacheReadInputTokens":0,"cacheCreationInputTokens":0}}' ]] || + fail "Cursor collector maps Cursor token names onto the record contract" "$result" +pass "Cursor collector maps Cursor token names onto the record contract" + +[[ $(jq -c '{today: .record.todayTotalTokens, byModel: .record.todayTokensByModel, activeDays: .record.activeDays}' <<<"$result") == '{"today":116,"byModel":{"default":116},"activeDays":2}' ]] || + fail "Cursor collector separates today from the rest of the week" "$result" +pass "Cursor collector separates today from the rest of the week" + +[[ $(jq -r '.carriedRecentDays' <<<"$result") == "true" && $(jq -r '.carriedSkipsAggregation' <<<"$result") == "true" ]] || + fail "Cursor collector keeps the token sections on --limits-only" "$result" +pass "Cursor collector keeps the token sections on --limits-only" From af8c27e7355680c78e9e7301584452fdcb503dd6 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Thu, 1 Oct 2026 21:53:20 -0400 Subject: [PATCH 003/546] 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 77c50345e57b8bd4114b9fba9a0c0f5c32de83dc Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Thu, 1 Oct 2026 21:53:25 -0400 Subject: [PATCH 004/546] Read the Cursor editor's sign-in when the CLI has none Someone who only uses the Cursor editor never runs `cursor-agent login`, so fall back to the access token the editor keeps in its state.vscdb, opened read-only. The lookup comes from @AndrijaSkontra's Cursor collector in #7087. Co-Authored-By: AndrijaSkontra Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-cursor | 60 +++++++++++++++++++------ shell/plugins/agents/README.md | 5 ++- test/shell.d/agent-usage-cursor-test.sh | 40 +++++++++++++++++ 3 files changed, 90 insertions(+), 15 deletions(-) diff --git a/bin/omarchy-agent-usage-cursor b/bin/omarchy-agent-usage-cursor index ee632871b77..83f7fd43c31 100755 --- a/bin/omarchy-agent-usage-cursor +++ b/bin/omarchy-agent-usage-cursor @@ -7,8 +7,9 @@ Cursor keeps no token counts on disk, so everything here comes from the same dashboard RPCs the Cursor CLI itself calls: the plan and the period meters from GetCurrentPeriodUsage and GetPlanInfo, tokens per model and per day from -GetAggregatedUsageEvents. The agents panel only ever reads the JSON this -prints. +GetAggregatedUsageEvents. They are asked with the sign-in the Cursor CLI +stores, or failing that the one the Cursor editor keeps in its state.vscdb. +The agents panel only ever reads the JSON this prints. """ from __future__ import annotations @@ -16,6 +17,7 @@ from __future__ import annotations import argparse import json import os +import sqlite3 import sys import urllib.error import urllib.request @@ -26,7 +28,7 @@ from typing import Any AGENT_ID = "cursor" AGENT_NAME = "Cursor" -AUTH_HELP = "Run `cursor-agent login` to authenticate." +AUTH_HELP = "Sign in to the Cursor editor, or run `cursor-agent login`." API_BASE_URL = "https://api2.cursor.sh" SERVICE = "aiserver.v1.DashboardService" REQUEST_TIMEOUT = 20 @@ -65,17 +67,49 @@ def default_auth_path() -> Path: return config_home / "cursor" / "auth.json" -def read_token(path: Path) -> str: - """The credential the Cursor CLI stores, or the key the user exported.""" - key = str(os.environ.get("CURSOR_API_KEY") or "").strip() - if key: - return key - if not path.is_file(): +def default_state_db() -> Path: + config_home = Path(os.environ.get("XDG_CONFIG_HOME") or (Path.home() / ".config")) + return config_home / "Cursor" / "User" / "globalStorage" / "state.vscdb" + + +def read_ide_token(state_db: Path) -> str: + """The sign-in the Cursor editor keeps, for people who never ran the CLI. + + Read-only: the editor may be writing its database right now. No + immutable=1, so a live database with an active WAL still reads. + """ + if not state_db.is_file(): return "" try: - return str(json.loads(path.read_text()).get("accessToken") or "").strip() - except (OSError, ValueError): + conn = sqlite3.connect(state_db.resolve().as_uri() + "?mode=ro", uri=True, timeout=2) + except sqlite3.Error: return "" + try: + conn.execute("PRAGMA query_only = ON") + row = conn.execute("SELECT value FROM ItemTable WHERE key = 'cursorAuth/accessToken' LIMIT 1").fetchone() + except sqlite3.Error: + return "" + finally: + conn.close() + value = row[0] if row else None + if isinstance(value, bytes): + value = value.decode("utf-8", errors="replace") + return str(value or "").strip() + + +def read_token(path: Path, state_db: Path | None = None) -> str: + """The key the user exported, the credential the Cursor CLI stores, or the + editor's sign-in.""" + key = str(os.environ.get("CURSOR_API_KEY") or "").strip() + if key: + return key + token = "" + if path.is_file(): + try: + token = str(json.loads(path.read_text()).get("accessToken") or "").strip() + except (OSError, ValueError, AttributeError): + token = "" + return token or read_ide_token(state_db or default_state_db()) def iso(epoch_millis: Any) -> str | None: @@ -291,8 +325,8 @@ def summarize_limits(period: dict[str, Any]) -> list[dict[str, Any]]: return limits -def scan(api_base_url: str, auth_path: Path, limits_only: bool) -> dict[str, Any]: - token = read_token(auth_path) +def scan(api_base_url: str, auth_path: Path, limits_only: bool, state_db: Path | None = None) -> dict[str, Any]: + token = read_token(auth_path, state_db) if not token: return base_record(usageStatusText="Cursor unavailable", authHelpText=AUTH_HELP) diff --git a/shell/plugins/agents/README.md b/shell/plugins/agents/README.md index 2964f88c5bb..e546ff46557 100644 --- a/shell/plugins/agents/README.md +++ b/shell/plugins/agents/README.md @@ -101,8 +101,9 @@ again. Fireworks reads `~/.fireworks/auth.ini` (which `firectl set-api-key` creates), then the key opencode stores in `~/.local/share/opencode/auth.json` when Fireworks is signed in there. Cursor reads `CURSOR_API_KEY` first, then the token -`cursor-agent login` stores in `~/.config/cursor/auth.json`, and honors -`CURSOR_API_ENDPOINT`. +`cursor-agent login` stores in `~/.config/cursor/auth.json`, then the +editor's own sign-in in `~/.config/Cursor/User/globalStorage/state.vscdb` +(opened read-only), and honors `CURSOR_API_ENDPOINT`. ### Cursor diff --git a/test/shell.d/agent-usage-cursor-test.sh b/test/shell.d/agent-usage-cursor-test.sh index 026a3387907..840437d5284 100755 --- a/test/shell.d/agent-usage-cursor-test.sh +++ b/test/shell.d/agent-usage-cursor-test.sh @@ -152,3 +152,43 @@ pass "Cursor collector separates today from the rest of the week" [[ $(jq -r '.carriedRecentDays' <<<"$result") == "true" && $(jq -r '.carriedSkipsAggregation' <<<"$result") == "true" ]] || fail "Cursor collector keeps the token sections on --limits-only" "$result" pass "Cursor collector keeps the token sections on --limits-only" + +# Someone who only uses the Cursor editor has no CLI auth file; the editor's +# own sign-in in state.vscdb stands in, read without touching the database. +ide=$(python3 - "$ROOT/bin/omarchy-agent-usage-cursor" "$TEST_HOME/ide" "$TEST_HOME/.config/cursor/auth.json" <<'PY' +import hashlib +import importlib.machinery +import importlib.util +import json +import os +import sqlite3 +import sys +from pathlib import Path + +os.environ.pop("CURSOR_API_KEY", None) +loader = importlib.machinery.SourceFileLoader("cursor_collector", sys.argv[1]) +spec = importlib.util.spec_from_loader(loader.name, loader) +collector = importlib.util.module_from_spec(spec) +loader.exec_module(collector) + +root = Path(sys.argv[2]) +db = root / "Cursor" / "User" / "globalStorage" / "state.vscdb" +db.parent.mkdir(parents=True) +conn = sqlite3.connect(db) +conn.execute("CREATE TABLE ItemTable (key TEXT UNIQUE ON CONFLICT REPLACE, value BLOB)") +conn.execute("INSERT INTO ItemTable VALUES ('cursorAuth/accessToken', 'ide_token')") +conn.commit() +conn.close() +before = hashlib.sha256(db.read_bytes()).hexdigest() + +os.environ["XDG_CONFIG_HOME"] = str(root) +print(json.dumps({ + "ide": collector.read_token(root / "cursor" / "auth.json"), + "cliWins": collector.read_token(Path(sys.argv[3])), + "untouched": hashlib.sha256(db.read_bytes()).hexdigest() == before and sorted(p.name for p in db.parent.iterdir()) == ["state.vscdb"], +})) +PY +) +[[ $(jq -c . <<<"$ide") == '{"ide":"ide_token","cliWins":"cursor_test","untouched":true}' ]] || + fail "Cursor collector falls back to the editor's sign-in, read-only" "$ide" +pass "Cursor collector falls back to the editor's sign-in, read-only" 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 005/546] 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 c340da9d8f4af2f8bb619a462626387f02467f4b Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Thu, 1 Oct 2026 21:55:19 -0400 Subject: [PATCH 006/546] Keep the last good Cursor limits through a failed check A failed dashboard call kept nothing: the record lost its meters, plan, and week of tokens. Cache each sign-in's last good limits and plan, and when a check fails, show them marked limitsStale with their limitsFetchedAt so the panel dims them, keep the token sections already on disk, and say when the sign-in was refused. Mention Cursor in the manual. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-cursor | 86 +++++++++++++++++++++++-- manual/17-ai.md | 2 +- shell/plugins/agents/README.md | 6 +- test/shell.d/agent-usage-cursor-test.sh | 32 +++++++++ 4 files changed, 119 insertions(+), 7 deletions(-) diff --git a/bin/omarchy-agent-usage-cursor b/bin/omarchy-agent-usage-cursor index 83f7fd43c31..3611619bd7b 100755 --- a/bin/omarchy-agent-usage-cursor +++ b/bin/omarchy-agent-usage-cursor @@ -15,10 +15,13 @@ The agents panel only ever reads the JSON this prints. from __future__ import annotations import argparse +import base64 +import hashlib import json import os import sqlite3 import sys +import tempfile import urllib.error import urllib.request from concurrent.futures import ThreadPoolExecutor @@ -36,7 +39,9 @@ RECENT_DAYS = 7 class CursorError(Exception): - pass + def __init__(self, message: str, auth: bool = False) -> None: + super().__init__(message) + self.auth = auth def number(value: Any) -> int: @@ -161,6 +166,8 @@ def base_record(**overrides: Any) -> dict[str, Any]: "usageStatusText": "", "authHelpText": "", "limits": [], + "limitsStale": False, + "limitsFetchedAt": 0, } record.update(empty_stats()) record.update(overrides) @@ -189,7 +196,7 @@ class CursorClient: body = response.read().decode("utf-8") except urllib.error.HTTPError as error: if error.code in (401, 403): - raise CursorError(AUTH_HELP) from error + raise CursorError(AUTH_HELP, auth=True) from error raise CursorError(f"Cursor API returned {error.code}") from error except (urllib.error.URLError, TimeoutError) as error: raise CursorError("Cursor API unreachable") from error @@ -325,23 +332,94 @@ def summarize_limits(period: dict[str, Any]) -> list[dict[str, Any]]: return limits +def cache_root() -> Path: + return Path(os.environ.get("XDG_CACHE_HOME") or (Path.home() / ".cache")) / "omarchy" / "agent-usage" + + +def read_json(path: Path) -> dict[str, Any]: + try: + data = json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError): + return {} + return data if isinstance(data, dict) else {} + + +# A temporary file of its own per write, so overlapping runs can't rename +# each other's away. +def write_json(path: Path, payload: dict[str, 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 + + +# Keyed by who the token signs in as (the JWT's subject), so a machine signed +# in to someone else never shows the last one's limits, and a refreshed token +# keeps the same cache. An opaque API key keys by its own hash. +def limits_cache(token: str) -> Path: + subject = "" + parts = token.split(".") + if len(parts) == 3: + try: + claims = json.loads(base64.urlsafe_b64decode(parts[1] + "=" * (-len(parts[1]) % 4))) + subject = str(claims.get("sub") or "") if isinstance(claims, dict) else "" + except (ValueError, UnicodeDecodeError): + subject = "" + key = hashlib.sha256((subject or token).encode("utf-8")).hexdigest()[:16] + return cache_root() / f"cursor-limits-{key}.json" + + +def now_ms() -> int: + return round(datetime.now(timezone.utc).timestamp() * 1000) + + def scan(api_base_url: str, auth_path: Path, limits_only: bool, state_db: Path | None = None) -> dict[str, Any]: token = read_token(auth_path, state_db) if not token: return base_record(usageStatusText="Cursor unavailable", authHelpText=AUTH_HELP) client = CursorClient(token, api_base_url) - period = client.period() + cache = limits_cache(token) + try: + period = client.period() + except CursorError as error: + # A failed check keeps the last good limits, marked stale so the panel + # dims them, and the week of tokens already on disk. + cached = read_json(cache) + limits = cached.get("limits") or [] + status = "Waiting for auth" if error.auth else ("" if limits else "Cursor limits unavailable") + record = base_record( + ready=True, + tierLabel=str(cached.get("tierLabel") or ""), + usageStatusText=status, + authHelpText=str(error), + limits=limits, + limitsStale=True, + limitsFetchedAt=number(cached.get("fetchedAtMs")), + ) + record.update(carried_stats()) + return record + record = base_record( ready=True, tierLabel=client.plan_name(), authHelpText=AUTH_HELP, limits=summarize_limits(period), + limitsFetchedAt=now_ms(), ) + write_json(cache, {"fetchedAtMs": record["limitsFetchedAt"], "tierLabel": record["tierLabel"], "limits": record["limits"]}) if limits_only: record.update(carried_stats()) else: - record.update(collect_stats(client, datetime.now().astimezone().date())) + try: + record.update(collect_stats(client, datetime.now().astimezone().date())) + except CursorError: + record.update(carried_stats()) return record diff --git a/manual/17-ai.md b/manual/17-ai.md index ea8a8d65eaf..ee75897f5c8 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, Cursor, and Fireworks are covered out of the box. 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 e546ff46557..be115a31b9b 100644 --- a/shell/plugins/agents/README.md +++ b/shell/plugins/agents/README.md @@ -108,8 +108,10 @@ editor's own sign-in in `~/.config/Cursor/User/globalStorage/state.vscdb` ### Cursor Cursor's session files on disk carry no token counts, so this is the one -collector with nothing to fall back on: without a credential the tab says -"Cursor unavailable" rather than showing a week of zeros. Tokens by day and by +collector with nothing to fall back on: without a credential the record stays +empty and the panel skips it rather than showing a week of zeros. A failed +check keeps the last good meters, dimmed, and the week already on disk; a +refused sign-in also says so under the name. Tokens by day and by model cover the last seven days, one aggregation call per day. The meters are Cursor's own percentages — included total, auto-bucket models, diff --git a/test/shell.d/agent-usage-cursor-test.sh b/test/shell.d/agent-usage-cursor-test.sh index 840437d5284..ff82d4a1b4c 100755 --- a/test/shell.d/agent-usage-cursor-test.sh +++ b/test/shell.d/agent-usage-cursor-test.sh @@ -7,6 +7,7 @@ require_command python3 TEST_HOME=$(mktemp -d) trap 'rm -rf "$TEST_HOME"' EXIT +export XDG_CACHE_HOME="$TEST_HOME/.cache" mkdir -p "$TEST_HOME/.config/cursor" cat >"$TEST_HOME/.config/cursor/auth.json" <<'EOF' @@ -107,7 +108,26 @@ record_path.parent.mkdir(parents=True, exist_ok=True) record_path.write_text(json.dumps(record)) carried = collector.scan("https://api.example", Path(os.environ["XDG_CONFIG_HOME"]) / "cursor" / "auth.json", True) + +# A failed check keeps the last good limits and the week on disk, marked stale. +class DownClient(StubClient): + def call(self, method, payload): + raise collector.CursorError("Cursor API unreachable") + + +class RefusedClient(StubClient): + def call(self, method, payload): + raise collector.CursorError(collector.AUTH_HELP, auth=True) + + +collector.CursorClient = DownClient +down = collector.scan("https://api.example", Path(os.environ["XDG_CONFIG_HOME"]) / "cursor" / "auth.json", False) +collector.CursorClient = RefusedClient +refused = collector.scan("https://api.example", Path(os.environ["XDG_CONFIG_HOME"]) / "cursor" / "auth.json", False) + print(json.dumps({ + "down": down, + "refused": refused, "record": record, "aggregationCalls": sum(1 for method, _ in calls if method == "GetAggregatedUsageEvents"), "windowIsLocalMidnight": (first_start.hour, first_start.minute, first_start.second) == (0, 0, 0), @@ -153,6 +173,18 @@ pass "Cursor collector separates today from the rest of the week" fail "Cursor collector keeps the token sections on --limits-only" "$result" pass "Cursor collector keeps the token sections on --limits-only" +[[ $(jq -c '.record | {limitsStale, fetched: (.limitsFetchedAt > 0)}' <<<"$result") == '{"limitsStale":false,"fetched":true}' ]] || + fail "Cursor collector stamps fresh limits" "$result" +pass "Cursor collector stamps fresh limits" + +[[ $(jq -c '{same: (.down.limits == .record.limits), stamp: (.down.limitsFetchedAt >= .record.limitsFetchedAt), stale: .down.limitsStale, tier: .down.tierLabel, status: .down.usageStatusText, week: (.down.recentDays == .record.recentDays)}' <<<"$result") == '{"same":true,"stamp":true,"stale":true,"tier":"Ultra","status":"","week":true}' ]] || + fail "Cursor collector keeps the last good limits, stale, after a failed check" "$result" +pass "Cursor collector keeps the last good limits, stale, after a failed check" + +[[ $(jq -c '{stale: .refused.limitsStale, status: .refused.usageStatusText, kept: (.refused.limits | length)}' <<<"$result") == '{"stale":true,"status":"Waiting for auth","kept":4}' ]] || + fail "Cursor collector says a refused sign-in needs attention" "$result" +pass "Cursor collector says a refused sign-in needs attention" + # Someone who only uses the Cursor editor has no CLI auth file; the editor's # own sign-in in state.vscdb stands in, read without touching the database. ide=$(python3 - "$ROOT/bin/omarchy-agent-usage-cursor" "$TEST_HOME/ide" "$TEST_HOME/.config/cursor/auth.json" <<'PY' From 58c199045d43e1c0c1b3d8666130a8f257f48e17 Mon Sep 17 00:00:00 2001 From: Mauricio Wolff Date: Mon, 28 Sep 2026 13:58:16 +0200 Subject: [PATCH 007/546] Add an OpenCode agent usage collector Sums token usage on the opencode and opencode-go providers into the agents panel's record contract: pi and omp JSONL transcripts first, because they can spend the subscription without the opencode CLI ever running, then opencode's own message store for sessions the CLI recorded itself, reading `message` first and falling back to `session_message`. Read the rolling, weekly, and monthly meters from the Zen gateway's Go usage endpoint. The key comes from OPENCODE_API_KEY, then opencode's own credential store, then pi's. The probe is reused for 15 seconds, cached windows drop once they reset, the last good numbers survive a failed probe, and a transport failure tells the panel to retry shortly. A rejected key is told apart from a key with no Go subscription behind it. --- bin/omarchy-agent-usage-opencode | 654 ++++++++++++++++++ manual/17-ai.md | 2 +- shell/plugins/agents/README.md | 7 + .../plugins/agents/assets/opencode-light.svg | 1 + shell/plugins/agents/assets/opencode.svg | 1 + .../agent-usage-opencode-limits-test.sh | 181 +++++ .../agent-usage-opencode-scanner-test.sh | 158 +++++ 7 files changed, 1003 insertions(+), 1 deletion(-) create mode 100755 bin/omarchy-agent-usage-opencode create mode 100644 shell/plugins/agents/assets/opencode-light.svg create mode 100644 shell/plugins/agents/assets/opencode.svg create mode 100644 test/shell.d/agent-usage-opencode-limits-test.sh create mode 100644 test/shell.d/agent-usage-opencode-scanner-test.sh diff --git a/bin/omarchy-agent-usage-opencode b/bin/omarchy-agent-usage-opencode new file mode 100755 index 00000000000..10eb1970214 --- /dev/null +++ b/bin/omarchy-agent-usage-opencode @@ -0,0 +1,654 @@ +#!/usr/bin/python3 +# omarchy:summary=Print the OpenCode usage record as JSON +# omarchy:args=[--force] [--limits-only] +# omarchy:hidden=true +"""Collect OpenCode usage into one display-ready JSON record. + +Local stats come from pi/omp sessions that ran through an opencode provider +and from opencode's own message database. The agents panel only ever reads +the JSON this prints. +""" + +import argparse +import fcntl +import hashlib +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 + +AGENT_ID = "opencode" +AGENT_NAME = "OpenCode" + +# The subscription meters live on the Zen "go" gateway; the plain gateway has +# no usage endpoint. Limits are per-account, so they come from the same key +# regardless of which provider id a given run used. +DEFAULT_USAGE_URL = "https://opencode.ai/zen/go/v1/usage" +AUTH_HELP = "Run `opencode providers login`, set OPENCODE_API_KEY, or sign in with pi to read subscription limits." + +# A panel opened and closed repeatedly must not become one request per flick. +PROBE_MIN_INTERVAL_SECONDS = 15 + +# The subscription has two provider ids: "opencode" for the plain Zen +# gateway and "opencode-go" for the Go plan. Both spend the same account and +# both appear in pi sessions and in opencode's own database. +PROVIDERS = ("opencode", "opencode-go") + +# 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, however low +# refreshIntervalSec is set. --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 + + +def number(value): + try: + return int(value or 0) + except Exception: + return 0 + + +# Usage percentages arrive as 0-100; the record contract wants 0-1. An +# unreadable value returns -1 so the caller can drop the window. +def fraction(value): + try: + return float(value) / 100.0 + except (TypeError, ValueError): + return -1.0 + + +def model_name(raw): + # opencode stores the model as a {providerID, modelID} object in its own + # database and as a plain string in pi sessions. A variant suffix such as + # ":high" names the thinking level, not the model. + if isinstance(raw, dict): + raw = raw.get("modelID") or raw.get("id") or "" + value = str(raw or "").rstrip("/").split("/")[-1].split(":")[0] or "opencode" + # The panel title-cases words but cannot produce all-caps acronyms. + if value[:3].lower() == "glm": + value = "GLM" + value[3:] + return value + + +def runtime_env(): + home = str(Path.home()) + path_parts = [ + os.environ.get("PATH", ""), + f"{home}/.local/bin", + f"{home}/.npm-global/bin", + f"{home}/.local/share/mise/shims", + ] + env = os.environ.copy() + env["PATH"] = os.pathsep.join(part for part in path_parts if part) + return env + + +ENV = runtime_env() + + +def find_command(name): + return shutil.which(name, path=ENV.get("PATH")) + + +def local_day(value): + if value is None: + return datetime.now().strftime("%Y-%m-%d") + if isinstance(value, (int, float)): + # pi message timestamps are milliseconds; opencode writes milliseconds too. + 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") + + +now = datetime.now() +today = now.strftime("%Y-%m-%d") +recent_dates = [(now - timedelta(days=offset)).strftime("%Y-%m-%d") for offset in range(6, -1, -1)] +recent = {day: {"date": day, "messageCount": 0} for day in recent_dates} +today_tokens_by_model = {} +model_usage = {} +today_sessions = set() +active_days = set() + +today_prompts = 0 +today_total_tokens = 0 +total_prompts = 0 +total_sessions = set() +seen_pi_messages = set() + + +def add_usage(day, session_key, model, input_tokens, output_tokens, cache_read, cache_write): + global today_prompts, today_total_tokens, total_prompts + total = input_tokens + output_tokens + cache_read + cache_write + total_prompts += 1 + total_sessions.add(session_key) + active_days.add(day) + + bucket = model_usage.setdefault(model, { + "inputTokens": 0, + "outputTokens": 0, + "cacheReadInputTokens": 0, + "cacheCreationInputTokens": 0, + }) + bucket["inputTokens"] += input_tokens + bucket["outputTokens"] += output_tokens + bucket["cacheReadInputTokens"] += cache_read + bucket["cacheCreationInputTokens"] += cache_write + + if day in recent: + recent[day]["messageCount"] += total + + if day == today: + today_prompts += 1 + today_sessions.add(session_key) + today_total_tokens += total + today_tokens_by_model[model] = today_tokens_by_model.get(model, 0) + total + + +def scan_pi_sessions(): + # pi and omp can spend the subscription without opencode's CLI ever + # running, so its own database cannot be the only source. Their JSONL + # transcripts carry the provider, model, and token usage per message. + roots = [ + Path.home() / ".pi" / "agent" / "sessions", + Path.home() / ".omp" / "agent" / "sessions", + ] + rg = find_command("rg") or "rg" + for root in roots: + if not root.exists(): + continue + try: + proc = subprocess.Popen( + [rg, "--json", "-e", r'"provider"\s*:\s*"opencode(-go)?"', str(root)], + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + text=True, + errors="replace", + env=ENV, + ) + except FileNotFoundError: + return + + assert proc.stdout is not None + for raw in proc.stdout: + try: + event = json.loads(raw) + if event.get("type") != "match": + continue + line = event.get("data", {}).get("lines", {}).get("text", "") + path = event.get("data", {}).get("path", {}).get("text", "pi-session") + entry = json.loads(line) + except Exception: + continue + + if entry.get("type") != "message": + continue + message_key = path + ":" + str(entry.get("id") or "") + if message_key in seen_pi_messages: + continue + seen_pi_messages.add(message_key) + message = entry.get("message") or {} + if message.get("role") != "assistant": + continue + if str(message.get("provider") or "") not in PROVIDERS: + continue + + usage = message.get("usage") or {} + if not usage: + continue + input_tokens = number(usage.get("input")) + # pi folds thinking into its own field; opencode counts it as generated. + output_tokens = number(usage.get("output")) + number(usage.get("reasoning")) + cache_read = number(usage.get("cacheRead")) + cache_write = number(usage.get("cacheWrite")) + if not (input_tokens or output_tokens or cache_read or cache_write): + continue + + day = local_day(entry.get("timestamp") or message.get("timestamp")) + # A model id may carry a thinking variant suffix ("deepseek-v4.1-flash:high"). + model = model_name(message.get("model")) + add_usage(day, path, model, input_tokens, output_tokens, cache_read, cache_write) + + try: + proc.wait(timeout=1) + except Exception: + proc.kill() + + +def table_columns(conn, table): + try: + return [row[1] for row in conn.execute(f"PRAGMA table_info({table})")] + except sqlite3.Error: + return [] + + +def store_row(session_id, entry, created): + # Shared by both message tables: same JSON shape, different containers. + if not isinstance(entry, dict) or entry.get("role") != "assistant": + return False + # Exact match: opencode provider ids are free-form, and a custom + # "opencode-local" gateway is not this subscription. + if str(entry.get("providerID") or "") not in PROVIDERS: + return False + 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): + return False + day = local_day((entry.get("time") or {}).get("created") or created) + model = model_name(entry.get("model") or entry.get("modelID")) + add_usage(day, "opencode:" + str(session_id), model, input_tokens, output_tokens, cache_read, cache_write) + return True + + +def scan_message_rows(conn): + # Current opencode writes one row per chat message here. The LIKE gates + # skip rows whose JSON cannot hold the key/value pairs before any parse. + # json_extract is the authority after them: it RAISES on malformed JSON, so + # each call sits inside a json_valid CASE rather than a bare AND term. + found = 0 + for session_id, raw in conn.execute( + "SELECT session_id, data FROM message" + " WHERE data LIKE '%\"role\"%:%\"assistant\"%'" + " AND (data LIKE '%\"providerID\"%:%\"opencode\"%' OR data LIKE '%\"providerID\"%:%\"opencode-go\"%')" + " 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 IN ('opencode', 'opencode-go'))" + ): + try: + entry = json.loads(raw) + except ValueError: + continue + if store_row(session_id, entry, None): + found += 1 + return found + + +def scan_session_message_rows(conn): + # Some builds move the same message JSON into session_message with the role + # in a `type` column. It is empty on the current release, so this only runs + # when `message` held no assistant rows; that keeps the two stores from + # double counting a session if both ever carry it. + columns = table_columns(conn, "session_message") + if "type" not in columns or "data" not in columns: + return 0 + found = 0 + for session_id, row_type, created, raw in conn.execute("SELECT session_id, type, time_created, data FROM session_message"): + try: + entry = json.loads(raw) + except ValueError: + continue + if isinstance(entry, dict) and not entry.get("role"): + entry["role"] = str(row_type or "") + if store_row(session_id, entry, created): + found += 1 + return found + + +def scan_opencode_db(): + # A subscription burned entirely through opencode leaves no pi session + # files, but opencode records per-message provider, model, and token usage + # in its own database. Read-only: opencode may be writing right now. + # 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 = Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) / "opencode" / "opencode.db" + 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") + # Read the current store first, the migrated one only when `message` held + # nothing, so a store that has both cannot count a session twice. + found = scan_message_rows(conn) if table_columns(conn, "message") else 0 + if found == 0: + scan_session_message_rows(conn) + except sqlite3.Error: + # Transient lock, schema migration, corruption: the numbers stop here, + # incomplete. + return False + finally: + conn.close() + return True + + +def cache_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(): + db = Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) / "opencode" / "opencode.db" + # The digest covers every data path the scan reads: the opencode DB and, + # via Path.home(), the pi/omp session roots. + digest = hashlib.sha1((str(Path.home()) + "\n" + str(db)).encode("utf-8")).hexdigest()[:16] + root = cache_root() + return root / f"opencode-scan-{digest}.json", root / f"opencode-scan-{digest}.lock" + + +def read_fresh_json(path, max_age_seconds): + 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, payload): + # 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=(",", ":")) + "\n") + # mkstemp opens at 0600; nothing in the cache is sensitive, so open it + # up to the usual 0644. + tmp.chmod(0o644) + tmp.replace(path) + except BaseException: + tmp.unlink(missing_ok=True) + raise + + +# The cache payload is a versioned envelope around the local-stats dict, so a +# corrupted or foreign-shaped file is a cache miss (rescan + rewrite) instead +# of a crash or a garbage record. +def read_cached_stats(cache_file, max_age_seconds): + 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, + # not merely old, whatever its mtime says. + 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, stats): + try: + write_json(cache_file, {"schemaVersion": 1, "scanDate": today, "stats": stats}) + except Exception as exc: + print(f"omarchy-agent-usage-opencode: could not write usage cache ({exc})", file=sys.stderr) + + +def local_stats(): + """Snapshot the aggregated local usage into the record's stats dict.""" + return { + "todayPrompts": today_prompts, + "todaySessions": len(today_sessions), + "todayTotalTokens": today_total_tokens, + "todayTokensByModel": today_tokens_by_model, + "recentDays": [recent[day] for day in recent_dates], + "totalPrompts": total_prompts, + "totalSessions": len(total_sessions), + # Days with any recorded usage, for the all-time "N days" summary. The + # dates travel too: merging snapshots from several machines needs their + # union, which a count alone cannot give. + "activeDays": len(active_days), + "activeDates": sorted(active_days), + "modelUsage": model_usage, + } + + +def run_local_scans(): + scan_pi_sessions() + complete = scan_opencode_db() + return local_stats(), complete + + +def cached_local_stats(max_age): + """Local stats, with the cache as a pure optimization. + + The cache must never take the collector down: any cache-layer failure + (unwritable cache root, lock errors, disk full) degrades to a direct scan + and a warning on stderr. 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-opencode: cache unavailable ({exc}); scanning directly", file=sys.stderr) + stats, _ = run_local_scans() + return stats + + +def _cached_local_stats(max_age): + cache_file, lock_file = scan_cache_paths() + + cached = read_cached_stats(cache_file, max_age) + if cached is not None: + return cached + + with lock_file.open("w") as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + cached = read_cached_stats(cache_file, max_age) + if cached is not None: + return cached + stats, complete = run_local_scans() + # An interrupted scan still serves this run, but caching it would + # suppress the missing usage for every reader until the cache expires. + if complete: + write_cached_stats(cache_file, stats) + return stats + + +def read_auth_key(path): + try: + parsed = json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError): + return "" + if not isinstance(parsed, dict): + return "" + # Prefer the exact provider ids, then any opencode* entry: the login flow + # sometimes files the key under a different id than the runtime uses. + entries = [parsed[provider] for provider in PROVIDERS if provider in parsed] + entries += [entry for name, entry in parsed.items() if "opencode" in str(name).lower()] + for entry in entries: + if isinstance(entry, dict) and entry.get("key"): + return str(entry["key"]).strip() + return "" + + +def subscription_key(): + # An explicit key wins, then opencode's own credential store, then pi's. pi + # drives the same Zen subscription and records it under the same provider + # ids, so a machine that codes through pi never signs the opencode CLI in. + explicit = str(os.environ.get("OPENCODE_API_KEY", "")).strip() + if explicit: + return explicit + data_home = Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) + key = read_auth_key(data_home / "opencode" / "auth.json") + if key: + return key + return read_auth_key(Path.home() / ".pi" / "agent" / "auth.json") + + +def probe_limits(key): + url = os.environ.get("OPENCODE_USAGE_URL") or DEFAULT_USAGE_URL + request = urllib.request.Request(url, headers={ + "Authorization": "Bearer " + key, + "Accept": "application/json", + # Cloudflare answers the default urllib signature with a 1010 before the + # request reaches the API, so identify the collector instead. + "User-Agent": "omarchy-agent-usage-opencode/1.0", + }) + 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 == 401: + return {"ok": False, "helpText": "OpenCode rejected the saved key. Reconnect OpenCode Go with /connect in the TUI, or set OPENCODE_API_KEY."} + if error.code == 403: + return {"ok": False, "helpText": "The saved key has no OpenCode Go subscription behind it. Local OpenCode stats are still shown."} + return {"ok": False, "helpText": f"OpenCode's usage endpoint returned status {error.code}. Local OpenCode 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 OpenCode's usage endpoint. Retrying shortly. Local OpenCode stats are still shown."} + + usage = payload.get("usage") if isinstance(payload, dict) else None + if not isinstance(usage, dict): + return {"ok": False, "helpText": "OpenCode's usage endpoint returned no data. Local OpenCode stats are still shown."} + + limits = [] + for name, label in (("rolling", "Rolling (5h)"), ("weekly", "Weekly (7-day)"), ("monthly", "Monthly")): + window = usage.get(name) + if not isinstance(window, dict): + continue + percent = fraction(window.get("percent")) + if percent < 0: + continue + limits.append({ + "label": label, + "percent": max(0.0, min(1.0, percent)), + "resetsAt": str(window.get("resetsAt") or ""), + }) + if not limits: + return {"ok": False, "helpText": "OpenCode's usage endpoint returned no quota windows. Local OpenCode stats are still shown."} + return {"ok": True, "limits": limits} + + +def limit_window_open(entry, now): + # A cached percentage stops meaning anything once its window resets: a + # stale 78% would misreport an allowance that is now untouched. An + # unreadable timestamp is no reason to drop a real number. + 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 + + +def open_cached_limits(cached): + entries = cached.get("limits") if isinstance(cached, dict) else None + 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 collect_limits(force): + result = {"limits": [], "tierLabel": "", "usageStatusText": "Local usage only", "authHelpText": "", "retryAdvised": False} + probe_cache = cache_root() / "opencode-limits.json" + cached = read_fresh_json(probe_cache, float("inf")) or {} + fallback = open_cached_limits(cached) + + key = subscription_key() + if not key: + result["limits"] = fallback + result["authHelpText"] = AUTH_HELP + return result + + # A panel opened and closed repeatedly must not become one request per + # flick, so a recent answer is reused; --force skips the wait. + fetched_at = number(cached.get("fetchedAtMs")) / 1000.0 + if fallback and not force and 0 <= time.time() - fetched_at < PROBE_MIN_INTERVAL_SECONDS: + result["limits"] = fallback + result["tierLabel"] = "Go" + result["usageStatusText"] = "" + return result + + probe = probe_limits(key) + if probe.get("ok"): + result["limits"] = probe["limits"] + result["tierLabel"] = "Go" + result["usageStatusText"] = "" + try: + write_json(probe_cache, {"fetchedAtMs": round(time.time() * 1000), "limits": probe["limits"]}) + except Exception: + pass + return result + + if probe.get("transport"): + result["retryAdvised"] = True + # Last good numbers beat an empty limits section while the probe is down. + if fallback: + result["limits"] = fallback + result["tierLabel"] = "Go" + result["usageStatusText"] = "" + else: + result["usageStatusText"] = "OpenCode limits unavailable" + result["authHelpText"] = probe.get("helpText", "") + return result + + +def main(): + parser = argparse.ArgumentParser() + # --force rescans everything and rewrites the cache. --limits-only is kept + # for CLI compatibility with the panel's refreshLimits() call: only the + # limits probe must be fresh, so it may reuse a scan for far longer than a + # normal run, whose short window exists purely to dedup concurrent + # collector runs. + parser.add_argument("--force", action="store_true") + parser.add_argument("--limits-only", action="store_true") + args = parser.parse_args() + + max_age = 0 if args.force else (LIMITS_ONLY_REUSE_SECONDS if args.limits_only else SCAN_REUSE_SECONDS) + stats = cached_local_stats(max_age) + limits = collect_limits(args.force) + retry_advised = bool(limits.pop("retryAdvised", False)) + + record = { + "schemaVersion": 1, + "id": AGENT_ID, + "name": AGENT_NAME, + "updatedAt": datetime.now(timezone.utc).isoformat(), + "ready": True, + "hasLocalStats": True, + } + record.update(stats) + record.update(limits) + if retry_advised: + record["retryAdvised"] = True + print(json.dumps(record, separators=(",", ":"))) + + +if __name__ == "__main__": + main() diff --git a/manual/17-ai.md b/manual/17-ai.md index ee75897f5c8..f1fadbf93f7 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, Cursor, 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, Cursor, OpenCode, and Fireworks are covered out of the box. 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 be115a31b9b..039f998f910 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 | | `cursor` | Cursor's dashboard RPCs: included, auto-model, named-model, and on-demand meters | the same dashboard RPCs, one call per day for the last week | +| `opencode` | OpenCode Zen's Go usage endpoint (rolling + weekly + monthly) | pi and omp sessions on an opencode provider, plus opencode's own message store (`message`, falling back to `session_message`) | | `fireworks` | Estimated prepaid balance: configured funding minus rated account costs | Fireworks billing API, grouped by day and model for the last 30 days | When `~/.local/state/omarchy/agents/accounts/.json` @@ -122,6 +123,12 @@ past 100% while Cursor's own meters do not. The on-demand row is the exception and is a real ratio: dollars spent against the on-demand limit the account sets. +OpenCode limits read `OPENCODE_API_KEY` first, then opencode's own credential +store, then pi's. A recent probe answer is reused for 15 seconds, so opening +the panel repeatedly does not repeat the request, and cached windows are +dropped once they reset. Without a key the panel still shows local token +usage. + ### Fireworks balance The collector first asks the account's `:getBalance` endpoint for the real diff --git a/shell/plugins/agents/assets/opencode-light.svg b/shell/plugins/agents/assets/opencode-light.svg new file mode 100644 index 00000000000..b79140a5070 --- /dev/null +++ b/shell/plugins/agents/assets/opencode-light.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/shell/plugins/agents/assets/opencode.svg b/shell/plugins/agents/assets/opencode.svg new file mode 100644 index 00000000000..b79c7332e20 --- /dev/null +++ b/shell/plugins/agents/assets/opencode.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/test/shell.d/agent-usage-opencode-limits-test.sh b/test/shell.d/agent-usage-opencode-limits-test.sh new file mode 100644 index 00000000000..1f4b8ae1173 --- /dev/null +++ b/test/shell.d/agent-usage-opencode-limits-test.sh @@ -0,0 +1,181 @@ +#!/bin/bash + +source "$(dirname "$0")/base-test.sh" + +require_command python3 + +# collect_limits reaches the Zen gateway, so the reader that interprets its +# answer is exercised on its own: the collector loads as a module, and a fake +# urlopen stands in for the response. +TEST_HOME=$(mktemp -d) +trap 'rm -rf "$TEST_HOME"' EXIT + +COLLECTOR="$ROOT/bin/omarchy-agent-usage-opencode" TEST_HOME="$TEST_HOME" python3 - <<'PY' +import importlib.machinery +import importlib.util +import io +import json +import os +import pathlib +import sys +import 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) + +home = os.environ["TEST_HOME"] +os.environ["XDG_CACHE_HOME"] = os.path.join(home, "cache") +limits_cache = pathlib.Path(os.environ["XDG_CACHE_HOME"]) / "omarchy" / "agent-usage" / "opencode-limits.json" + + +def check(description, condition, detail=""): + if condition: + print("ok - " + description) + return + print("not ok - " + description, file=sys.stderr) + if detail: + print(detail, file=sys.stderr) + sys.exit(1) + + +def answer(payload=None, error=None): + def open_(request, timeout=None): + if error is not None: + raise error + return io.BytesIO(json.dumps(payload).encode()) + return open_ + + +def clear_cache(): + limits_cache.unlink(missing_ok=True) + + +def clear_key(): + os.environ.pop("OPENCODE_API_KEY", None) + os.environ.pop("XDG_DATA_HOME", None) + os.environ["HOME"] = os.path.join(home, "empty") + os.makedirs(os.environ["HOME"], exist_ok=True) + clear_cache() + + +os.environ["OPENCODE_API_KEY"] = "zen_test" +collector.urllib.request.urlopen = answer({"usage": { + "rolling": {"percent": 16, "resetsAt": "2999-01-01T00:00:00.000Z"}, + "weekly": {"percent": 6, "resetsAt": "2026-10-05T00:00:00.000Z"}, + "monthly": {"percent": 29, "resetsAt": "2026-10-11T15:38:40.000Z"}, +}}) +result = collector.collect_limits(True) +check( + "OpenCode collector maps every Zen usage window", + [(w["label"], w["percent"]) for w in result["limits"]] == [("Rolling (5h)", 0.16), ("Weekly (7-day)", 0.06), ("Monthly", 0.29)] + and result["tierLabel"] == "Go" and result["usageStatusText"] == "", + json.dumps(result), +) + +# A panel that opens and shuts repeatedly must not probe every time. +collector.urllib.request.urlopen = answer(error=RuntimeError("probe should have been skipped")) +result = collector.collect_limits(False) +check( + "OpenCode collector reuses a fresh probe cache without a request", + [(w["label"], w["percent"]) for w in result["limits"]] == [("Rolling (5h)", 0.16), ("Weekly (7-day)", 0.06), ("Monthly", 0.29)] + and result["usageStatusText"] == "", + json.dumps(result), +) + +collector.urllib.request.urlopen = answer({"usage": {"rolling": {"percent": 50, "resetsAt": "2999-01-01T00:00:00.000Z"}}}) +result = collector.collect_limits(True) +check( + "OpenCode collector --force probes past the cache", + [(w["label"], w["percent"]) for w in result["limits"]] == [("Rolling (5h)", 0.5)], + json.dumps(result), +) + +clear_cache() +collector.urllib.request.urlopen = answer({"usage": {}}) +result = collector.collect_limits(True) +check( + "OpenCode collector reports a payload without windows", + result["limits"] == [] and result["usageStatusText"] == "OpenCode limits unavailable", + json.dumps(result), +) + +clear_cache() +collector.urllib.request.urlopen = answer(error=urllib.error.HTTPError("https://x", 401, "Unauthorized", {}, None)) +result = collector.collect_limits(True) +check( + "OpenCode collector asks for a fresh key on 401", + result["limits"] == [] and result["usageStatusText"] == "OpenCode limits unavailable" and "rejected the saved key" in result["authHelpText"], + json.dumps(result), +) + +clear_cache() +collector.urllib.request.urlopen = answer(error=urllib.error.HTTPError("https://x", 403, "Forbidden", {}, None)) +result = collector.collect_limits(True) +check( + "OpenCode collector names a missing Go subscription on 403", + "no OpenCode Go subscription" in result["authHelpText"], + json.dumps(result), +) + +clear_cache() +collector.urllib.request.urlopen = answer(error=urllib.error.HTTPError("https://x", 500, "Server Error", {}, None)) +result = collector.collect_limits(True) +check( + "OpenCode collector shows a non-auth status code", + "returned status 500" in result["authHelpText"], + json.dumps(result), +) + +clear_cache() +collector.urllib.request.urlopen = answer(error=urllib.error.URLError("no route")) +result = collector.collect_limits(True) +check( + "OpenCode collector advises a retry when no server answered", + result["retryAdvised"] is True and result["usageStatusText"] == "OpenCode limits unavailable", + json.dumps(result), +) + +clear_key() +collector.urllib.request.urlopen = answer({"usage": {"rolling": {"percent": 1}}}) +result = collector.collect_limits(True) +check( + "OpenCode collector falls back to local stats without a key", + result["limits"] == [] and result["usageStatusText"] == "Local usage only" and result["authHelpText"] != "", + json.dumps(result), +) + +# A cached percentage is only good until its window resets. +collector.cache_root().mkdir(parents=True, exist_ok=True) +limits_cache.write_text(json.dumps({ + "fetchedAtMs": 0, + "limits": [ + {"label": "Rolling (5h)", "percent": 0.9, "resetsAt": "2020-01-01T00:00:00.000Z"}, + {"label": "Monthly", "percent": 0.4, "resetsAt": "2999-01-01T00:00:00.000Z"}, + ], +})) +result = collector.collect_limits(True) +check( + "OpenCode collector drops a cached window that has already reset", + [w["label"] for w in result["limits"]] == ["Monthly"], + json.dumps(result), +) + +# The key can come from opencode's own credential store or, when a machine +# codes through pi, from pi's. +opencode_home = os.path.join(home, "opencode-data") +os.makedirs(os.path.join(opencode_home, "opencode"), exist_ok=True) +with open(os.path.join(opencode_home, "opencode", "auth.json"), "w") as handle: + json.dump({"opencode": {"type": "api_key", "key": "from_opencode"}}, handle) +os.environ["XDG_DATA_HOME"] = opencode_home +check("OpenCode collector reads opencode's own credential", collector.subscription_key() == "from_opencode") + +clear_key() +pi_home = os.path.join(home, "pi") +os.makedirs(os.path.join(pi_home, ".pi", "agent"), exist_ok=True) +with open(os.path.join(pi_home, ".pi", "agent", "auth.json"), "w") as handle: + json.dump({"opencode-go": {"type": "api_key", "key": "from_pi"}}, handle) +os.environ["HOME"] = pi_home +check("OpenCode collector reads pi's credential as a last resort", collector.subscription_key() == "from_pi") +PY diff --git a/test/shell.d/agent-usage-opencode-scanner-test.sh b/test/shell.d/agent-usage-opencode-scanner-test.sh new file mode 100644 index 00000000000..6acbe998cbb --- /dev/null +++ b/test/shell.d/agent-usage-opencode-scanner-test.sh @@ -0,0 +1,158 @@ +#!/bin/bash + +source "$(dirname "$0")/base-test.sh" + +require_command jq +require_command python3 + +timestamp="$(date +%Y-%m-%d)T12:00:00Z" + +# pi and omp can spend an OpenCode subscription without opencode's CLI ever +# running, so its own database cannot be the only source. +PI_HOME=$(mktemp -d) +trap 'rm -rf "$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" <"$CACHE_HOME/.pi/agent/sessions/project/pi.jsonl" < Date: Thu, 1 Oct 2026 21:57:36 -0400 Subject: [PATCH 008/546] 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 009/546] 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 010/546] 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 011/546] 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 012/546] 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 f6e1900bc6d42c0924dbc1be32b49db17804f953 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Thu, 1 Oct 2026 22:01:59 -0400 Subject: [PATCH 013/546] Read OpenCode V2 and Console sign-ins, and dim stale OpenCode limits Fold in the pieces of @felixzsh's OpenCode Go collector from #12651: read both message stores, V2's session_message first, deduplicated by message id and leaving out fork copies; honor OPENCODE_DB; take Go keys from opencode's credential table and fall back to the Console session, asked on /inference/go/v1/usage with its x-opencode-org-id; accept the alternate Usage payload; read a rate-limited window as full; and cache limits per credential, handing a refused one over to the next. Kept limits now carry limitsStale and limitsFetchedAt so the panel dims them, the Go tier follows the limits, and a machine with no OpenCode usage gives an empty record the panel skips. Only the Zen and Go providers count, so opencode runs on Anthropic or OpenAI stay with the claude and codex records. Co-Authored-By: felixzsh Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-opencode | 418 ++++++++++++------ shell/plugins/agents/README.md | 23 +- shell/plugins/agents/manifest.json | 3 +- .../agent-usage-opencode-limits-test.sh | 120 ++++- .../agent-usage-opencode-scanner-test.sh | 71 ++- 5 files changed, 482 insertions(+), 153 deletions(-) diff --git a/bin/omarchy-agent-usage-opencode b/bin/omarchy-agent-usage-opencode index 10eb1970214..9b1ceb49c74 100755 --- a/bin/omarchy-agent-usage-opencode +++ b/bin/omarchy-agent-usage-opencode @@ -5,8 +5,11 @@ """Collect OpenCode usage into one display-ready JSON record. Local stats come from pi/omp sessions that ran through an opencode provider -and from opencode's own message database. The agents panel only ever reads -the JSON this prints. +and from opencode's own message database, both of its stores. Only the Zen +and Go providers count: opencode runs on Anthropic or OpenAI belong to the +claude and codex records. The limits are the Go plan's rolling, weekly, and +monthly meters from OpenCode's usage endpoint. The agents panel only ever +reads the JSON this prints. """ import argparse @@ -24,15 +27,21 @@ import urllib.error import urllib.request from datetime import datetime, timedelta, timezone from pathlib import Path +from typing import NamedTuple AGENT_ID = "opencode" AGENT_NAME = "OpenCode" # The subscription meters live on the Zen "go" gateway; the plain gateway has # no usage endpoint. Limits are per-account, so they come from the same key -# regardless of which provider id a given run used. -DEFAULT_USAGE_URL = "https://opencode.ai/zen/go/v1/usage" +# regardless of which provider id a given run used. A Console session asks +# Console's own endpoint for the same numbers. +API_BASE_URL = "https://opencode.ai" +USAGE_PATH = "/zen/go/v1/usage" +CONSOLE_USAGE_PATH = "/inference/go/v1/usage" +CONSOLE_ORG_HEADER = "x-opencode-org-id" AUTH_HELP = "Run `opencode providers login`, set OPENCODE_API_KEY, or sign in with pi to read subscription limits." +CONSOLE_EXPIRED_HELP = "OpenCode's Console session has expired. Run `opencode auth login` to refresh it." # A panel opened and closed repeatedly must not become one request per flick. PROBE_MIN_INTERVAL_SECONDS = 15 @@ -235,6 +244,25 @@ def scan_pi_sessions(): proc.kill() +def data_home(): + return Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) + + +def opencode_db(): + # OPENCODE_DB points opencode at another database, as the CLI itself reads + # it; a relative name lives beside the default one. Only a build whose + # channel was literally "v2" wrote opencode-v2.db, so it is a fallback for + # those installs alone. + root = data_home() / "opencode" + configured = os.environ.get("OPENCODE_DB", "").strip() + if configured: + path = Path(configured).expanduser() + return path if path.is_absolute() else root / path + normal = root / "opencode.db" + legacy = root / "opencode-v2.db" + return legacy if not normal.is_file() and legacy.is_file() else normal + + def table_columns(conn, table): try: return [row[1] for row in conn.execute(f"PRAGMA table_info({table})")] @@ -242,14 +270,22 @@ def table_columns(conn, table): return [] -def store_row(session_id, entry, created): - # Shared by both message tables: same JSON shape, different containers. - if not isinstance(entry, dict) or entry.get("role") != "assistant": - return False +def store_row(message_id, session_id, entry, row_type, created, seen): + # Shared by both message stores. The original `message` table keeps the + # provider and model at the top level; the V2 `session_message` table nests + # them under `model` and keeps the role in its `type` column, with a + # compaction's tokens spent like any reply's. + if not isinstance(entry, dict): + return + if (entry.get("role") or row_type) not in ("assistant", "compaction"): + return + model = entry.get("model") + provider = entry.get("providerID") or (model.get("providerID") if isinstance(model, dict) else "") # Exact match: opencode provider ids are free-form, and a custom - # "opencode-local" gateway is not this subscription. - if str(entry.get("providerID") or "") not in PROVIDERS: - return False + # "opencode-local" gateway is not this subscription. Anthropic and OpenAI + # runs never match, so the claude and codex collectors keep those. + if str(provider or "") not in PROVIDERS: + return tokens = entry.get("tokens") or {} cache = tokens.get("cache") or {} input_tokens = number(tokens.get("input")) @@ -258,21 +294,52 @@ def store_row(session_id, entry, created): 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): - return False + return + # A migrated message can sit in both stores under the same id. + if message_id in seen: + return + seen.add(message_id) day = local_day((entry.get("time") or {}).get("created") or created) - model = model_name(entry.get("model") or entry.get("modelID")) - add_usage(day, "opencode:" + str(session_id), model, input_tokens, output_tokens, cache_read, cache_write) - return True - - -def scan_message_rows(conn): - # Current opencode writes one row per chat message here. The LIKE gates - # skip rows whose JSON cannot hold the key/value pairs before any parse. - # json_extract is the authority after them: it RAISES on malformed JSON, so - # each call sits inside a json_valid CASE rather than a bare AND term. - found = 0 - for session_id, raw in conn.execute( - "SELECT session_id, data FROM message" + add_usage(day, "opencode:" + str(session_id), model_name(entry.get("modelID") or model), input_tokens, output_tokens, cache_read, cache_write) + + +def scan_session_message_rows(conn, tables, seen): + # The V2 store. The LIKE gate skips rows whose JSON cannot name one of the + # providers before any parse. + columns = set(table_columns(conn, "session_message")) + if not {"id", "session_id", "type", "data"} <= columns: + return + created = "m.time_created" if "time_created" in columns else "NULL" + source = "session_message m" + fork_filter = "" + if "time_created" in columns and "session_v2" in tables and {"id", "fork_session_id", "time_created"} <= set(table_columns(conn, "session_v2")): + # A fork copies its source session's messages under new ids; only what + # came after the fork was spent in it. + source += " LEFT JOIN session_v2 s ON s.id = m.session_id" + fork_filter = " AND (s.fork_session_id IS NULL OR m.time_created >= s.time_created)" + for message_id, session_id, row_type, time_created, raw in conn.execute( + f"SELECT m.id, m.session_id, m.type, {created}, m.data FROM {source}" + " WHERE m.type IN ('assistant', 'compaction')" + " AND m.data LIKE '%\"providerID\"%:%\"opencode%'" + fork_filter + ): + try: + entry = json.loads(raw) + except (TypeError, ValueError): + continue + store_row(str(message_id), session_id, entry, str(row_type or ""), time_created, seen) + + +def scan_message_rows(conn, seen): + # The original store, one row per chat message. The LIKE gates skip rows + # whose JSON cannot hold the key/value pairs before any parse. json_extract + # is the authority after them: it RAISES on malformed JSON, so each call + # sits inside a json_valid CASE rather than a bare AND term. + columns = set(table_columns(conn, "message")) + if not {"session_id", "data"} <= columns: + return + message_id = "id" if "id" in columns else "rowid" + for row_id, session_id, raw in conn.execute( + f"SELECT {message_id}, session_id, data FROM message" " WHERE data LIKE '%\"role\"%:%\"assistant\"%'" " AND (data LIKE '%\"providerID\"%:%\"opencode\"%' OR data LIKE '%\"providerID\"%:%\"opencode-go\"%')" " AND CASE WHEN json_valid(data) THEN json_extract(data, '$.role') END = 'assistant'" @@ -280,32 +347,9 @@ def scan_message_rows(conn): ): try: entry = json.loads(raw) - except ValueError: - continue - if store_row(session_id, entry, None): - found += 1 - return found - - -def scan_session_message_rows(conn): - # Some builds move the same message JSON into session_message with the role - # in a `type` column. It is empty on the current release, so this only runs - # when `message` held no assistant rows; that keeps the two stores from - # double counting a session if both ever carry it. - columns = table_columns(conn, "session_message") - if "type" not in columns or "data" not in columns: - return 0 - found = 0 - for session_id, row_type, created, raw in conn.execute("SELECT session_id, type, time_created, data FROM session_message"): - try: - entry = json.loads(raw) - except ValueError: + except (TypeError, ValueError): continue - if isinstance(entry, dict) and not entry.get("role"): - entry["role"] = str(row_type or "") - if store_row(session_id, entry, created): - found += 1 - return found + store_row(str(row_id), session_id, entry, "", None, seen) def scan_opencode_db(): @@ -315,7 +359,7 @@ def scan_opencode_db(): # 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 = Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) / "opencode" / "opencode.db" + db = opencode_db() if not db.is_file(): return True try: @@ -324,11 +368,14 @@ def scan_opencode_db(): return False try: conn.execute("PRAGMA query_only = ON") - # Read the current store first, the migrated one only when `message` held - # nothing, so a store that has both cannot count a session twice. - found = scan_message_rows(conn) if table_columns(conn, "message") else 0 - if found == 0: - scan_session_message_rows(conn) + tables = {name for (name,) in conn.execute("SELECT name FROM sqlite_master WHERE type = 'table'")} + # Both stores, V2 first, deduplicated by message id: a build that + # migrated keeps reading, and a message in both counts once. + seen = set() + if "session_message" in tables: + scan_session_message_rows(conn, tables, seen) + if "message" in tables: + scan_message_rows(conn, seen) except sqlite3.Error: # Transient lock, schema migration, corruption: the numbers stop here, # incomplete. @@ -345,7 +392,7 @@ def cache_root(): def scan_cache_paths(): - db = Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) / "opencode" / "opencode.db" + db = opencode_db() # The digest covers every data path the scan reads: the opencode DB and, # via Path.home(), the pi/omp session roots. digest = hashlib.sha1((str(Path.home()) + "\n" + str(db)).encode("utf-8")).hexdigest()[:16] @@ -473,6 +520,15 @@ def _cached_local_stats(max_age): return stats +class Credential(NamedTuple): + """One way of asking OpenCode for the Go meters.""" + + token: str + url: str + headers: dict + expired: bool = False + + def read_auth_key(path): try: parsed = json.loads(path.read_text(encoding="utf-8")) @@ -490,37 +546,120 @@ def read_auth_key(path): return "" -def subscription_key(): - # An explicit key wins, then opencode's own credential store, then pi's. pi - # drives the same Zen subscription and records it under the same provider - # ids, so a machine that codes through pi never signs the opencode CLI in. - explicit = str(os.environ.get("OPENCODE_API_KEY", "")).strip() - if explicit: - return explicit - data_home = Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) - key = read_auth_key(data_home / "opencode" / "auth.json") - if key: - return key - return read_auth_key(Path.home() / ".pi" / "agent" / "auth.json") - - -def probe_limits(key): - url = os.environ.get("OPENCODE_USAGE_URL") or DEFAULT_USAGE_URL - request = urllib.request.Request(url, headers={ - "Authorization": "Bearer " + key, +def read_store_credentials(db): + """The active credentials opencode V2 keeps in its database: Go keys, and + the Console session it signs the opencode integration in with.""" + keys, console = [], None + if not db.is_file(): + return keys, console + try: + conn = sqlite3.connect(db.resolve().as_uri() + "?mode=ro", uri=True, timeout=2) + except sqlite3.Error: + return keys, console + try: + conn.execute("PRAGMA query_only = ON") + rows = conn.execute( + "SELECT integration_id, value FROM credential WHERE active = 1 AND integration_id IN ('opencode-go', 'opencode')" + " ORDER BY integration_id = 'opencode-go' DESC, time_updated DESC" + ).fetchall() + except sqlite3.Error: + return keys, console + finally: + conn.close() + for _, raw in rows: + try: + value = json.loads(raw) + except (TypeError, ValueError): + continue + if not isinstance(value, dict): + continue + if value.get("type") == "key" and str(value.get("key") or "").strip(): + keys.append(str(value["key"]).strip()) + elif value.get("type") == "oauth" and console is None and str(value.get("access") or "").strip(): + # Console traffic is scoped to a workspace, and a wrong one is a hard + # 403, so the header only goes out when the id is there. + metadata = value.get("metadata") + org = str(metadata.get("orgID") or "").strip() if isinstance(metadata, dict) else "" + expired = 0 < number(value.get("expires")) <= time.time() * 1000 + console = Credential(str(value["access"]).strip(), API_BASE_URL + CONSOLE_USAGE_PATH, {CONSOLE_ORG_HEADER: org} if org else {}, expired) + return keys, console + + +def credential_candidates(): + # An explicit key wins, then opencode's own credential stores, then pi's: + # pi drives the same Zen subscription and records it under the same + # provider ids, so a machine that codes through pi never signs the opencode + # CLI in. The Console session is the implicit one, so it comes last. Which + # one the endpoint accepts isn't knowable from here, so a refused one only + # hands over to the next. + store_keys, console = read_store_credentials(opencode_db()) + tokens = [str(os.environ.get("OPENCODE_API_KEY", "")).strip()] + tokens += store_keys + tokens.append(read_auth_key(data_home() / "opencode" / "auth.json")) + tokens.append(read_auth_key(Path.home() / ".pi" / "agent" / "auth.json")) + candidates = [] + for token in tokens: + if token and token not in [c.token for c in candidates]: + candidates.append(Credential(token, API_BASE_URL + USAGE_PATH, {})) + if console is not None: + candidates.append(console) + return candidates + + +def reset_from_seconds(value): + seconds = number(value) + if seconds <= 0: + return "" + return (datetime.now(timezone.utc) + timedelta(seconds=seconds)).isoformat() + + +# The live endpoint nests the windows under "usage" with "percent" and +# "resetsAt"; another deployment spells them "Usage" with +# "usagePercent" and "resetInSec". Both are read, so the collector doesn't +# depend on which one answers. +def parse_usage(payload): + usage = payload.get("usage") if isinstance(payload, dict) else None + usage = usage if isinstance(usage, dict) else (payload if isinstance(payload, dict) else {}) + limits = [] + for name, label in (("rolling", "Rolling (5h)"), ("weekly", "Weekly (7-day)"), ("monthly", "Monthly")): + window = usage.get(name) + if not isinstance(window, dict): + window = payload.get(name + "Usage") if isinstance(payload, dict) else None + if not isinstance(window, dict): + continue + percent = fraction(window.get("percent", window.get("usagePercent"))) + if percent < 0: + continue + # A rate-limited window is spent until it resets, whatever it reads. + if str(window.get("status") or "") == "rate-limited": + percent = 1.0 + limits.append({ + "label": label, + "percent": max(0.0, min(1.0, percent)), + "resetsAt": str(window.get("resetsAt") or "").strip() or reset_from_seconds(window.get("resetInSec")), + }) + return limits + + +def probe_limits(credential): + request = urllib.request.Request(credential.url, headers={ + "Authorization": "Bearer " + credential.token, "Accept": "application/json", # Cloudflare answers the default urllib signature with a 1010 before the # request reaches the API, so identify the collector instead. "User-Agent": "omarchy-agent-usage-opencode/1.0", + **credential.headers, }) 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: + # A refused credential may just be the wrong one for this account, so the + # next way in still gets its turn. if error.code == 401: - return {"ok": False, "helpText": "OpenCode rejected the saved key. Reconnect OpenCode Go with /connect in the TUI, or set OPENCODE_API_KEY."} + return {"ok": False, "refused": True, "helpText": "OpenCode rejected the saved key. Reconnect OpenCode Go with /connect in the TUI, or set OPENCODE_API_KEY."} if error.code == 403: - return {"ok": False, "helpText": "The saved key has no OpenCode Go subscription behind it. Local OpenCode stats are still shown."} + return {"ok": False, "refused": True, "helpText": "The saved key has no OpenCode Go subscription behind it. Local OpenCode stats are still shown."} return {"ok": False, "helpText": f"OpenCode's usage endpoint returned status {error.code}. Local OpenCode stats are still shown."} except Exception: # A transport failure reached no server at all — no route, no DNS. Any @@ -528,23 +667,9 @@ def probe_limits(key): # pestering; this is not. return {"ok": False, "transport": True, "helpText": "Couldn't reach OpenCode's usage endpoint. Retrying shortly. Local OpenCode stats are still shown."} - usage = payload.get("usage") if isinstance(payload, dict) else None - if not isinstance(usage, dict): + if not isinstance(payload, dict) or not payload: return {"ok": False, "helpText": "OpenCode's usage endpoint returned no data. Local OpenCode stats are still shown."} - - limits = [] - for name, label in (("rolling", "Rolling (5h)"), ("weekly", "Weekly (7-day)"), ("monthly", "Monthly")): - window = usage.get(name) - if not isinstance(window, dict): - continue - percent = fraction(window.get("percent")) - if percent < 0: - continue - limits.append({ - "label": label, - "percent": max(0.0, min(1.0, percent)), - "resetsAt": str(window.get("resetsAt") or ""), - }) + limits = parse_usage(payload) if not limits: return {"ok": False, "helpText": "OpenCode's usage endpoint returned no quota windows. Local OpenCode stats are still shown."} return {"ok": True, "limits": limits} @@ -574,48 +699,75 @@ def open_cached_limits(cached): return [entry for entry in entries if isinstance(entry, dict) and limit_window_open(entry, now)] +# One cache per credential, so a machine signed in to another account never +# shows the last one's meters. +def limits_cache_path(token): + return cache_root() / f"opencode-limits-{hashlib.sha256(token.encode('utf-8')).hexdigest()[:16]}.json" + + def collect_limits(force): - result = {"limits": [], "tierLabel": "", "usageStatusText": "Local usage only", "authHelpText": "", "retryAdvised": False} - probe_cache = cache_root() / "opencode-limits.json" - cached = read_fresh_json(probe_cache, float("inf")) or {} - fallback = open_cached_limits(cached) - - key = subscription_key() - if not key: - result["limits"] = fallback + result = {"limits": [], "tierLabel": "", "usageStatusText": "Local usage only", "authHelpText": "", "retryAdvised": False, "stale": False, "fetchedAtMs": 0} + candidates = credential_candidates() + if not candidates: result["authHelpText"] = AUTH_HELP return result - # A panel opened and closed repeatedly must not become one request per - # flick, so a recent answer is reused; --force skips the wait. - fetched_at = number(cached.get("fetchedAtMs")) / 1000.0 - if fallback and not force and 0 <= time.time() - fetched_at < PROBE_MIN_INTERVAL_SECONDS: - result["limits"] = fallback - result["tierLabel"] = "Go" - result["usageStatusText"] = "" + # Limits only ever show for a Go plan, so a tier follows them. + def answer(limits, fetched_at, stale): + result.update(limits=limits, tierLabel="Go" if limits else "", usageStatusText="", stale=stale, fetchedAtMs=fetched_at if limits else 0) return result - probe = probe_limits(key) - if probe.get("ok"): - result["limits"] = probe["limits"] - result["tierLabel"] = "Go" - result["usageStatusText"] = "" - try: - write_json(probe_cache, {"fetchedAtMs": round(time.time() * 1000), "limits": probe["limits"]}) - except Exception: - pass + refused = None + for candidate in candidates: + cache = limits_cache_path(candidate.token) + cached = read_fresh_json(cache, float("inf")) or {} + fallback = open_cached_limits(cached) + fetched_at = number(cached.get("fetchedAtMs")) + + # A panel opened and closed repeatedly must not become one request per + # flick, so a recent answer is reused; --force skips the wait. + if fallback and not force and 0 <= time.time() - fetched_at / 1000 < PROBE_MIN_INTERVAL_SECONDS: + return answer(fallback, fetched_at, False) + + probe = probe_limits(candidate) + if probe.get("ok"): + fetched_at = round(time.time() * 1000) + try: + write_json(cache, {"fetchedAtMs": fetched_at, "limits": probe["limits"]}) + except Exception: + pass + return answer(probe["limits"], fetched_at, False) + + if probe.get("refused"): + # Keep the first refusal that still has open windows, so a later + # credential without a cache can't blank them. A lapsed Console + # session is what to fix, whichever refusal is kept. + expired = candidate.expired or bool(refused and refused[3]) + if refused is None or (not refused[1] and fallback): + refused = (probe, fallback, fetched_at, expired) + else: + refused = refused[:3] + (expired,) + continue + + # Last good numbers, dimmed as stale, beat an empty limits section while + # the probe is down. + result["retryAdvised"] = bool(probe.get("transport")) + answer(fallback, fetched_at, True) + if not fallback: + result["usageStatusText"] = "OpenCode limits unavailable" + result["authHelpText"] = probe["helpText"] return result - if probe.get("transport"): - result["retryAdvised"] = True - # Last good numbers beat an empty limits section while the probe is down. - if fallback: - result["limits"] = fallback - result["tierLabel"] = "Go" - result["usageStatusText"] = "" + # Every way in was refused: an account problem, not an outage. Keep any + # still-open numbers, but say what went wrong. + probe, fallback, fetched_at, expired = refused + answer(fallback, fetched_at, True) + if expired: + result["usageStatusText"] = "OpenCode Console session expired" + result["authHelpText"] = CONSOLE_EXPIRED_HELP else: - result["usageStatusText"] = "OpenCode limits unavailable" - result["authHelpText"] = probe.get("helpText", "") + result["usageStatusText"] = "Sign-in rejected" if fallback else "OpenCode limits unavailable" + result["authHelpText"] = probe["helpText"] return result @@ -633,19 +785,25 @@ def main(): max_age = 0 if args.force else (LIMITS_ONLY_REUSE_SECONDS if args.limits_only else SCAN_REUSE_SECONDS) stats = cached_local_stats(max_age) limits = collect_limits(args.force) - retry_advised = bool(limits.pop("retryAdvised", False)) + # Nobody uses OpenCode's subscription here: no usage and no limits leave + # an empty tier, a record the panel skips. record = { "schemaVersion": 1, "id": AGENT_ID, "name": AGENT_NAME, "updatedAt": datetime.now(timezone.utc).isoformat(), - "ready": True, + "ready": number(stats.get("totalPrompts")) > 0 or bool(limits["limits"]), "hasLocalStats": True, + "tierLabel": limits["tierLabel"], + "usageStatusText": limits["usageStatusText"], + "authHelpText": limits["authHelpText"], + "limits": limits["limits"], + "limitsStale": limits["stale"], + "limitsFetchedAt": limits["fetchedAtMs"], } record.update(stats) - record.update(limits) - if retry_advised: + if limits["retryAdvised"]: record["retryAdvised"] = True print(json.dumps(record, separators=(",", ":"))) diff --git a/shell/plugins/agents/README.md b/shell/plugins/agents/README.md index 039f998f910..dd21382d710 100644 --- a/shell/plugins/agents/README.md +++ b/shell/plugins/agents/README.md @@ -74,7 +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 | | `cursor` | Cursor's dashboard RPCs: included, auto-model, named-model, and on-demand meters | the same dashboard RPCs, one call per day for the last week | -| `opencode` | OpenCode Zen's Go usage endpoint (rolling + weekly + monthly) | pi and omp sessions on an opencode provider, plus opencode's own message store (`message`, falling back to `session_message`) | +| `opencode` | OpenCode Go's usage endpoint (rolling + weekly + monthly), with an API key or the Console session | pi and omp sessions on the Zen and Go providers, plus opencode's own `message` and `session_message` stores | | `fireworks` | Estimated prepaid balance: configured funding minus rated account costs | Fireworks billing API, grouped by day and model for the last 30 days | When `~/.local/state/omarchy/agents/accounts/.json` @@ -123,11 +123,22 @@ past 100% while Cursor's own meters do not. The on-demand row is the exception and is a real ratio: dollars spent against the on-demand limit the account sets. -OpenCode limits read `OPENCODE_API_KEY` first, then opencode's own credential -store, then pi's. A recent probe answer is reused for 15 seconds, so opening -the panel repeatedly does not repeat the request, and cached windows are -dropped once they reset. Without a key the panel still shows local token -usage. +### OpenCode + +One record covers the OpenCode subscription, Zen and Go together: usage on +the `opencode` and `opencode-go` providers. opencode sessions on an Anthropic +or OpenAI provider stay in the `claude` and `codex` records, so nothing counts +twice. Both of opencode's message stores are read (`message` and V2's +`session_message`, deduplicated by message id, leaving out the copies a fork +makes), from the database `OPENCODE_DB` names when set. + +The Go meters read `OPENCODE_API_KEY` first, then the keys opencode keeps in +its database and `auth.json`, then pi's, and last the OpenCode Console +session, which asks Console's own endpoint with its workspace id. A refused +credential hands over to the next one. A recent answer is reused for 15 +seconds, each credential caches its own, a failed check keeps the last good +windows, dimmed, until they reset, and a rate-limited window reads as full. +Without a Go plan the panel still shows local token usage. ### Fireworks balance diff --git a/shell/plugins/agents/manifest.json b/shell/plugins/agents/manifest.json index 60af1e5a2fd..ea8b873b1ad 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, Cursor, and Fireworks usage, limits, and pace in a native Omarchy bar panel.", + "description": "Claude Code, Codex, Grok, Cursor, OpenCode, and Fireworks usage, limits, and pace in a native Omarchy bar panel.", "kinds": ["bar-widget"], "activation": "on-demand", "entryPoints": { @@ -23,6 +23,7 @@ "codex": { "enabled": true }, "grok": { "enabled": true }, "cursor": { "enabled": true }, + "opencode": { "enabled": true }, "fireworks": { "enabled": true } }, "refreshIntervalSec": 900, diff --git a/test/shell.d/agent-usage-opencode-limits-test.sh b/test/shell.d/agent-usage-opencode-limits-test.sh index 1f4b8ae1173..860f581ea64 100644 --- a/test/shell.d/agent-usage-opencode-limits-test.sh +++ b/test/shell.d/agent-usage-opencode-limits-test.sh @@ -4,6 +4,8 @@ source "$(dirname "$0")/base-test.sh" require_command python3 +unset OPENCODE_API_KEY OPENCODE_DB + # collect_limits reaches the Zen gateway, so the reader that interprets its # answer is exercised on its own: the collector loads as a module, and a fake # urlopen stands in for the response. @@ -27,7 +29,10 @@ loader.exec_module(collector) home = os.environ["TEST_HOME"] os.environ["XDG_CACHE_HOME"] = os.path.join(home, "cache") -limits_cache = pathlib.Path(os.environ["XDG_CACHE_HOME"]) / "omarchy" / "agent-usage" / "opencode-limits.json" +# Never the real sign-ins: every credential store lives under the test home. +os.environ["HOME"] = os.path.join(home, "empty") +os.environ["XDG_DATA_HOME"] = os.path.join(home, "empty-data") +cache_dir = pathlib.Path(os.environ["XDG_CACHE_HOME"]) / "omarchy" / "agent-usage" def check(description, condition, detail=""): @@ -49,12 +54,13 @@ def answer(payload=None, error=None): def clear_cache(): - limits_cache.unlink(missing_ok=True) + for cache in cache_dir.glob("opencode-limits-*.json"): + cache.unlink() def clear_key(): os.environ.pop("OPENCODE_API_KEY", None) - os.environ.pop("XDG_DATA_HOME", None) + os.environ["XDG_DATA_HOME"] = os.path.join(home, "empty-data") os.environ["HOME"] = os.path.join(home, "empty") os.makedirs(os.environ["HOME"], exist_ok=True) clear_cache() @@ -63,8 +69,8 @@ def clear_key(): os.environ["OPENCODE_API_KEY"] = "zen_test" collector.urllib.request.urlopen = answer({"usage": { "rolling": {"percent": 16, "resetsAt": "2999-01-01T00:00:00.000Z"}, - "weekly": {"percent": 6, "resetsAt": "2026-10-05T00:00:00.000Z"}, - "monthly": {"percent": 29, "resetsAt": "2026-10-11T15:38:40.000Z"}, + "weekly": {"percent": 6, "resetsAt": "2999-01-05T00:00:00.000Z"}, + "monthly": {"percent": 29, "resetsAt": "2999-01-11T15:38:40.000Z"}, }}) result = collector.collect_limits(True) check( @@ -146,30 +152,64 @@ check( json.dumps(result), ) -# A cached percentage is only good until its window resets. +# A failed probe keeps the last good numbers, marked stale with when they +# were measured, but only until each window resets. +os.environ["OPENCODE_API_KEY"] = "zen_test" collector.cache_root().mkdir(parents=True, exist_ok=True) -limits_cache.write_text(json.dumps({ - "fetchedAtMs": 0, +collector.limits_cache_path("zen_test").write_text(json.dumps({ + "fetchedAtMs": 1700000000000, "limits": [ {"label": "Rolling (5h)", "percent": 0.9, "resetsAt": "2020-01-01T00:00:00.000Z"}, {"label": "Monthly", "percent": 0.4, "resetsAt": "2999-01-01T00:00:00.000Z"}, ], })) +collector.urllib.request.urlopen = answer(error=urllib.error.HTTPError("https://x", 500, "Server Error", {}, None)) +result = collector.collect_limits(True) +check( + "OpenCode collector keeps open cached windows, stale, after a failed probe", + [w["label"] for w in result["limits"]] == ["Monthly"] and result["stale"] is True + and result["fetchedAtMs"] == 1700000000000 and result["tierLabel"] == "Go" and result["usageStatusText"] == "", + json.dumps(result), +) + +collector.urllib.request.urlopen = answer({"usage": {"monthly": {"percent": 41, "resetsAt": "2999-01-01T00:00:00.000Z"}}}) +result = collector.collect_limits(True) +check( + "OpenCode collector marks fresh limits live", + result["stale"] is False and result["fetchedAtMs"] > 1700000000000, + json.dumps(result), +) + +# Another sign-in never shows this one's numbers. +os.environ["OPENCODE_API_KEY"] = "zen_other" +collector.urllib.request.urlopen = answer(error=urllib.error.URLError("no route")) +result = collector.collect_limits(True) +check("OpenCode collector caches limits per credential", result["limits"] == [] and result["stale"] is True, json.dumps(result)) + +# Another deployment spells the windows Usage, and a rate-limited +# window is spent whatever it reads. +clear_cache() +collector.urllib.request.urlopen = answer({ + "rollingUsage": {"usagePercent": 40, "resetInSec": 3600, "status": "rate-limited"}, + "weeklyUsage": {"usagePercent": "12.5", "resetInSec": 0}, +}) result = collector.collect_limits(True) check( - "OpenCode collector drops a cached window that has already reset", - [w["label"] for w in result["limits"]] == ["Monthly"], + "OpenCode collector reads the alternate payload shape", + [(w["label"], w["percent"]) for w in result["limits"]] == [("Rolling (5h)", 1.0), ("Weekly (7-day)", 0.125)] + and result["limits"][0]["resetsAt"] != "" and result["limits"][1]["resetsAt"] == "", json.dumps(result), ) -# The key can come from opencode's own credential store or, when a machine +# The key can come from opencode's own credential stores or, when a machine # codes through pi, from pi's. +clear_key() opencode_home = os.path.join(home, "opencode-data") os.makedirs(os.path.join(opencode_home, "opencode"), exist_ok=True) with open(os.path.join(opencode_home, "opencode", "auth.json"), "w") as handle: json.dump({"opencode": {"type": "api_key", "key": "from_opencode"}}, handle) os.environ["XDG_DATA_HOME"] = opencode_home -check("OpenCode collector reads opencode's own credential", collector.subscription_key() == "from_opencode") +check("OpenCode collector reads opencode's own credential", [c.token for c in collector.credential_candidates()] == ["from_opencode"]) clear_key() pi_home = os.path.join(home, "pi") @@ -177,5 +217,59 @@ os.makedirs(os.path.join(pi_home, ".pi", "agent"), exist_ok=True) with open(os.path.join(pi_home, ".pi", "agent", "auth.json"), "w") as handle: json.dump({"opencode-go": {"type": "api_key", "key": "from_pi"}}, handle) os.environ["HOME"] = pi_home -check("OpenCode collector reads pi's credential as a last resort", collector.subscription_key() == "from_pi") +check("OpenCode collector reads pi's credential as a last resort", [c.token for c in collector.credential_candidates()] == ["from_pi"]) + +# OpenCode V2 keeps its credentials in the database: a Go key first, and the +# Console session last, asked on Console's endpoint with its workspace. +import sqlite3 +import time +clear_key() +v2_home = os.path.join(home, "v2-data") +os.makedirs(os.path.join(v2_home, "opencode"), exist_ok=True) +conn = sqlite3.connect(os.path.join(v2_home, "opencode", "opencode.db")) +conn.execute("CREATE TABLE credential (id text, integration_id text, active integer, time_updated integer, value text)") +conn.executemany("INSERT INTO credential VALUES (?, ?, ?, ?, ?)", [ + ("c1", "opencode-go", 1, 2, json.dumps({"type": "key", "key": "store_key"})), + ("c2", "opencode-go", 0, 3, json.dumps({"type": "key", "key": "inactive_key"})), + ("c3", "opencode", 1, 1, json.dumps({"type": "oauth", "access": "console_token", "expires": int(time.time() * 1000) + 60000, "metadata": {"orgID": "org_1"}})), +]) +conn.commit() +conn.close() +os.environ["XDG_DATA_HOME"] = v2_home +candidates = collector.credential_candidates() +check( + "OpenCode collector reads V2 credentials, the Console session last", + [(c.token, c.url.endswith("/inference/go/v1/usage"), c.headers) for c in candidates] + == [("store_key", False, {}), ("console_token", True, {"x-opencode-org-id": "org_1"})], + repr(candidates), +) + +asked = [] +def refuse_key(request, timeout=None): + asked.append((request.full_url, request.get_header("Authorization"), request.get_header("X-opencode-org-id"))) + if "store_key" in request.get_header("Authorization"): + raise urllib.error.HTTPError(request.full_url, 401, "Unauthorized", {}, None) + return io.BytesIO(json.dumps({"usage": {"weekly": {"percent": 30, "resetsAt": "2999-01-01T00:00:00.000Z"}}}).encode()) +collector.urllib.request.urlopen = refuse_key +result = collector.collect_limits(True) +check( + "OpenCode collector hands a refused key over to the Console session", + [(w["label"], w["percent"]) for w in result["limits"]] == [("Weekly (7-day)", 0.3)] + and asked[-1] == ("https://opencode.ai/inference/go/v1/usage", "Bearer console_token", "org_1"), + json.dumps({"result": result, "asked": asked}), +) + +# Every way in refused, with the Console session expired, says so. +conn = sqlite3.connect(os.path.join(v2_home, "opencode", "opencode.db")) +conn.execute("UPDATE credential SET value = ? WHERE id = 'c3'", (json.dumps({"type": "oauth", "access": "console_token", "expires": 1000}),)) +conn.commit() +conn.close() +clear_cache() +collector.urllib.request.urlopen = answer(error=urllib.error.HTTPError("https://x", 401, "Unauthorized", {}, None)) +result = collector.collect_limits(True) +check( + "OpenCode collector names an expired Console session", + result["usageStatusText"] == "OpenCode Console session expired" and result["limits"] == [], + json.dumps(result), +) PY diff --git a/test/shell.d/agent-usage-opencode-scanner-test.sh b/test/shell.d/agent-usage-opencode-scanner-test.sh index 6acbe998cbb..80942b3c957 100644 --- a/test/shell.d/agent-usage-opencode-scanner-test.sh +++ b/test/shell.d/agent-usage-opencode-scanner-test.sh @@ -5,6 +5,8 @@ source "$(dirname "$0")/base-test.sh" require_command jq require_command python3 +unset OPENCODE_API_KEY OPENCODE_DB + timestamp="$(date +%Y-%m-%d)T12:00:00Z" # pi and omp can spend an OpenCode subscription without opencode's CLI ever @@ -67,6 +69,7 @@ conn.executemany("INSERT INTO message VALUES (?, ?, ?, ?, ?)", [ message("msg_1", "opencode", "deepseek-v4-flash", input=80, output=40, reasoning=5, read=30), message("msg_2", "opencode-go", "minimax-m3", input=10, output=2), message("msg_3", "anthropic", "claude-opus-5", input=999, output=999), + message("msg_7", "openai", "gpt-5.5", input=999, output=999), message("msg_4", "opencode-local", "gpt-nope", input=999, output=999), message("msg_5", "opencode", "deepseek-v4-flash", role="user", input=999, output=999), ]) @@ -81,7 +84,7 @@ result=$(HOME="$OPENCODE_HOME" XDG_CACHE_HOME="$OPENCODE_HOME/.cache" XDG_DATA_H [[ $(jq -r '.todayTotalTokens' <<<"$result") == "167" ]] || fail "OpenCode collector counts opencode database usage, reasoning included" "$result" [[ $(jq -c '.modelUsage' <<<"$result") == '{"deepseek-v4-flash":{"inputTokens":80,"outputTokens":45,"cacheReadInputTokens":30,"cacheCreationInputTokens":0},"minimax-m3":{"inputTokens":10,"outputTokens":2,"cacheReadInputTokens":0,"cacheCreationInputTokens":0}}' ]] || - fail "OpenCode collector ignores prefix-colliding providers, user messages, and malformed rows" "$result" + fail "OpenCode collector leaves Anthropic and OpenAI runs to their own records, and ignores prefix-colliding providers, user messages, and malformed rows" "$result" [[ $(jq -r '.totalSessions' <<<"$result") == "1" ]] || fail "OpenCode collector counts a database session once" "$result" pass "OpenCode collector counts opencode database usage" @@ -111,8 +114,7 @@ result=$(HOME="$CACHE_HOME" XDG_CACHE_HOME="$CACHE_HOME/.cache" XDG_DATA_HOME="$ pass "OpenCode collector caches and reuses local stats" # Some builds move the same message JSON into session_message with the role in -# a type column. It must be read when `message` has no assistant rows, and -# ignored when `message` does, so a session is never counted twice. +# a type column, and it reads on its own. SESSION_HOME=$(mktemp -d) trap 'rm -rf "$PI_HOME" "$OPENCODE_HOME" "$CACHE_HOME" "$SESSION_HOME"' EXIT @@ -156,3 +158,66 @@ result=$(HOME="$SESSION_HOME" XDG_CACHE_HOME="$SESSION_HOME/.cache" XDG_DATA_HOM [[ $(jq -r '.totalSessions' <<<"$result") == "1" ]] || fail "OpenCode collector counts a fallback session once" "$result" pass "OpenCode collector falls back to the session_message store" + +# OpenCode V2 nests the provider under `model` and spends tokens on +# compactions too. A message migrated into both stores counts once, and a +# fork's copies of its source's messages don't count again. OPENCODE_DB +# points at the database the way opencode itself reads it. +V2_HOME=$(mktemp -d) +trap 'rm -rf "$PI_HOME" "$OPENCODE_HOME" "$CACHE_HOME" "$SESSION_HOME" "$V2_HOME"' EXIT + +python3 - "$V2_HOME/elsewhere/opencode-dev.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)") +conn.execute("CREATE TABLE session_message (id text PRIMARY KEY, session_id text NOT NULL, type text NOT NULL, time_created integer NOT NULL, data text NOT NULL)") +conn.execute("CREATE TABLE session_v2 (id text PRIMARY KEY, fork_session_id text, time_created integer NOT NULL)") +now_ms = int(time.time() * 1000) +forked_ms = now_ms - 1000 + +def v2(id, session, kind, provider, model, input, created=now_ms): + return (id, session, kind, created, json.dumps({ + "model": {"providerID": provider, "id": model}, + "tokens": {"input": input, "output": 0, "reasoning": 0, "cache": {"read": 0, "write": 0}}, + "time": {"created": created}, + })) + +conn.executemany("INSERT INTO session_v2 VALUES (?, ?, ?)", [("ses_a", None, now_ms - 5000), ("ses_b", "ses_a", forked_ms)]) +conn.executemany("INSERT INTO session_message VALUES (?, ?, ?, ?, ?)", [ + v2("msg_1", "ses_a", "assistant", "opencode-go", "kimi-k3", 100, now_ms - 4000), + v2("msg_2", "ses_a", "compaction", "opencode-go", "kimi-k3", 10, now_ms - 3000), + v2("msg_3", "ses_a", "assistant", "anthropic", "claude-opus-5", 999), + v2("msg_4", "ses_b", "assistant", "opencode-go", "kimi-k3", 100, forked_ms - 4000), + v2("msg_5", "ses_b", "assistant", "opencode-go", "kimi-k3", 1), +]) +conn.execute("INSERT INTO message VALUES ('msg_1', 'ses_a', ?, ?, ?)", (now_ms, now_ms, json.dumps({ + "role": "assistant", "providerID": "opencode-go", "modelID": "kimi-k3", + "tokens": {"input": 100, "output": 0, "reasoning": 0, "cache": {"read": 0, "write": 0}}, + "time": {"created": now_ms - 4000}, +}))) +conn.commit() +conn.close() +PY + +result=$(HOME="$V2_HOME" XDG_CACHE_HOME="$V2_HOME/.cache" XDG_DATA_HOME="$V2_HOME/.local/share" \ + OPENCODE_DB="$V2_HOME/elsewhere/opencode-dev.db" PATH="$V2_HOME/bin:$PATH" "$ROOT/bin/omarchy-agent-usage-opencode") +[[ $(jq -c '{todayTotalTokens, totalPrompts, totalSessions}' <<<"$result") == '{"todayTotalTokens":111,"totalPrompts":3,"totalSessions":2}' ]] || + fail "OpenCode collector reads V2 records once each, without fork copies" "$result" +pass "OpenCode collector reads V2 records once each, without fork copies" + +# A machine nobody uses OpenCode's subscription on gives a record the panel +# skips: no plan, no limits, no usage. +EMPTY_HOME=$(mktemp -d) +trap 'rm -rf "$PI_HOME" "$OPENCODE_HOME" "$CACHE_HOME" "$SESSION_HOME" "$V2_HOME" "$EMPTY_HOME"' EXIT +result=$(HOME="$EMPTY_HOME" XDG_CACHE_HOME="$EMPTY_HOME/.cache" XDG_DATA_HOME="$EMPTY_HOME/.local/share" \ + OPENCODE_API_KEY= PATH="$EMPTY_HOME/bin:$PATH" "$ROOT/bin/omarchy-agent-usage-opencode") +[[ $(jq -c '{ready, tierLabel, limits, totalPrompts, limitsStale}' <<<"$result") == '{"ready":false,"tierLabel":"","limits":[],"totalPrompts":0,"limitsStale":false}' ]] || + fail "OpenCode collector gives an empty record where nobody uses it" "$result" +pass "OpenCode collector gives an empty record where nobody uses it" From 53d4a6491a5962cf0194fccfbdf774b2a0ef2263 Mon Sep 17 00:00:00 2001 From: Michael D Date: Thu, 1 Oct 2026 22:02:18 -0400 Subject: [PATCH 014/546] 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 015/546] 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 cda42911130595206c2d9f720658f9c5ffabf6ca Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Marek=20=C5=A0neberger?= Date: Wed, 2 Sep 2026 15:42:41 +0200 Subject: [PATCH 016/546] Add a GitHub Copilot collector to the agents panel Read the Copilot CLI's session store (~/.copilot/session-store.db), whose assistant_usage_events table records every billed request with its day, model, session, token split, and cost, and fall back to the session transcripts for CLI versions whose store predates it: user-typed prompts, and tokens and cost from the cumulative shutdown metrics, split over the days the replies landed on. Input already includes cache reads and writes and output already includes reasoning, so neither counts twice. The monthly meter is the account-wide AI credit allowance from GitHub's internal quota endpoint, read as credits, with a local estimate labeled "(est.)" when the endpoint can't be asked. --- bin/omarchy-agent-usage-copilot | 996 ++++++++++++++++++ shell/plugins/agents/README.md | 44 + shell/plugins/agents/assets/copilot-light.svg | 1 + shell/plugins/agents/assets/copilot.svg | 1 + test/shell.d/agent-usage-copilot-test.sh | 294 ++++++ 5 files changed, 1336 insertions(+) create mode 100755 bin/omarchy-agent-usage-copilot create mode 100644 shell/plugins/agents/assets/copilot-light.svg create mode 100644 shell/plugins/agents/assets/copilot.svg create mode 100755 test/shell.d/agent-usage-copilot-test.sh diff --git a/bin/omarchy-agent-usage-copilot b/bin/omarchy-agent-usage-copilot new file mode 100755 index 00000000000..1a8eb52b3c0 --- /dev/null +++ b/bin/omarchy-agent-usage-copilot @@ -0,0 +1,996 @@ +#!/usr/bin/python3 +# omarchy:summary=Print the GitHub Copilot usage record as JSON +# omarchy:args=[--force] [--limits-only] +# omarchy:hidden=true +"""Collect GitHub Copilot usage into one display-ready JSON record. + +GitHub exposes no supported personal-account usage API — the Copilot REST +endpoints are organization- or enterprise-scoped — but the Copilot CLI keeps +two local sources this collector reads. The primary one is the session store +(~/.copilot/session-store.db): its assistant_usage_events table records every +billed request with day, model, session, token split, and cost in nano AI +units (1e9 nano units = 1 AI credit at 1 cent). The fallback, for CLI versions +whose store predates that table, is the session transcript directory +(~/.copilot/session-state//events.jsonl), whose shutdown events carry the +same per-model token split and cost. The record prefers whichever source +produced numbers, never both. + +The monthly meter is an allowance of AI credits. GitHub's own account-wide +figure is authoritative when the internal quota endpoint answers (the same one +the editor plugins use); otherwise the collector falls back to this machine's +CLI cost against a plan allowance from configuration, because the included +amount is not discoverable locally. All remote reads fail soft — an offline +machine or a missing token degrades to the local estimate instead of erroring. + +Configuration lives in ~/.config/omarchy/agents/copilot.json: + + { + "plan": "pro", + "monthlyCredits": 1500, + "remote": true + } + +plan is one of free, student, pro, pro+, max, business, enterprise and picks +the default allowance (Pro 1500, Pro+ 7000, Max 20000, Business 1900, +Enterprise 3900 credits per month; existing seat-based customers may see +higher promotional amounts through September 2026). monthlyCredits overrides +the plan table, for AI-credit plans without a published number or for a +personal budget. The plan is not discoverable locally, so without plan or +monthlyCredits there is no local estimate and no plan label. remote (default +true) turns the account-wide quota probe on or off; the probe needs a token +from COPILOT_QUOTA_TOKEN, GH_TOKEN, GITHUB_TOKEN, or `gh auth token`. The +allowance window resets on the 1st of each month at 00:00 UTC, GitHub's own +boundary. + +The agents panel only ever reads the JSON this prints. +""" + +from __future__ import annotations + +import argparse +import calendar +import datetime as dt +import fcntl +import hashlib +import json +import os +import shutil +import sqlite3 +import subprocess +import sys +import tempfile +import time +import urllib.error +import urllib.request +from pathlib import Path +from typing import Any + +AGENT_ID = "copilot" +AGENT_NAME = "GitHub Copilot" +AUTH_HELP = ( + "Run `copilot` in a terminal and sign in to GitHub to start recording sessions." +) + +ULTRA_SECONDS = 1_000_000_000 +QUOTA_URL = "https://api.github.com/copilot_internal/user" +REMOTE_TIMEOUT = 6.0 + +# Plan identifiers as the quota endpoint spells them. +API_PLAN_KEYS = { + "individual": "pro", + "individual_pro": "pro+", + "business": "business", + "enterprise": "enterprise", + "free": "free", + "max": "max", +} + +# Monthly AI credit allowances from GitHub's usage-based billing tables. +# Seat-based plans pool credits at the billing entity, so an individual's +# practical ceiling is set by their admin rather than this figure — +# monthlyCredits in the config overrides it. +PLAN_ALLOWANCES = { + "free": 0, + "student": 0, + "pro": 1_500, + "pro+": 7_000, + "max": 20_000, + "business": 1_900, + "enterprise": 3_900, +} + +PLAN_LABELS = { + "free": "Free", + "student": "Student", + "pro": "Pro", + "pro+": "Pro+", + "max": "Max", + "business": "Business", + "enterprise": "Enterprise", +} + +# The session store is mutable and grows without bound, so it is size-capped +# before it is snapshotted. Every scan aggregates as it streams, so the record's +# size depends on days and models, never on the number of rows. +MAX_DB_BYTES = 512 * 1024 * 1024 + +SCAN_REUSE_SECONDS = 20 +LIMITS_ONLY_REUSE_SECONDS = 900 + + +def number(value: Any) -> int: + try: + n = float(value or 0) + return round(n) if n == n else 0 + except (TypeError, ValueError): + return 0 + + +def expand_path(value: str) -> Path: + return Path(os.path.expandvars(os.path.expanduser(value))).resolve() + + +def config_path() -> Path: + config_home = Path(os.environ.get("XDG_CONFIG_HOME") or (Path.home() / ".config")) + return config_home / "omarchy" / "agents" / "copilot.json" + + +def read_config() -> dict[str, Any]: + try: + parsed = json.loads(config_path().read_text(encoding="utf-8")) + return parsed if isinstance(parsed, dict) else {} + except (OSError, ValueError): + return {} + + +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 + + +TOKEN_KEYS = ( + "inputTokens", + "outputTokens", + "cacheReadInputTokens", + "cacheCreationInputTokens", +) + + +def empty_bucket() -> dict[str, int]: + return dict.fromkeys(TOKEN_KEYS, 0) + + +def local_date_string() -> str: + return dt.datetime.now().date().isoformat() + + +def recent_date_strings() -> list[str]: + today = dt.datetime.now().date() + return [ + (today - dt.timedelta(days=offset)).isoformat() for offset in range(6, -1, -1) + ] + + +def month_start(now: dt.datetime) -> dt.datetime: + return now.replace(day=1, hour=0, minute=0, second=0, microsecond=0) + + +def next_month_start(now: dt.datetime) -> dt.datetime: + days = calendar.monthrange(now.year, now.month)[1] + return month_start(now) + dt.timedelta(days=days) + + +def positive_float(value: Any) -> float: + try: + n = float(value or 0) + except (TypeError, ValueError): + return 0.0 + return n if n > 0 else 0.0 + + +def credits_from_nano(nano: Any) -> float: + return positive_float(nano) / ULTRA_SECONDS + + +def token_split( + prompt_tokens: Any, output_tokens: Any, cache_read: Any, cache_write: Any +) -> dict[str, int]: + """Map Copilot's token counts onto the panel's four disjoint buckets. + + Copilot's input count is the whole prompt, cache reads and writes included, + and its output count already includes reasoning tokens. The buckets are + disjoint, like the other collectors', so the cached part comes off input + and reasoning is never added a second time. + """ + read = number(cache_read) + write = number(cache_write) + return { + "inputTokens": max(0, number(prompt_tokens) - read - write), + "outputTokens": number(output_tokens), + "cacheReadInputTokens": read, + "cacheCreationInputTokens": write, + } + + +def clean_name(value: Any) -> str: + return str(value or "").strip() + + +def parse_event_time(value: Any) -> dt.datetime | None: + """Parse a store or transcript timestamp into timezone-aware local time. + + Accepts ISO-8601 (with or without Z) and epoch seconds or milliseconds. + Everything else collapses to None; the caller falls back to the local date. + """ + if value is None: + return None + if isinstance(value, (int, float)): + try: + seconds = ( + float(value) / 1000.0 if float(value) > 10_000_000_000 else float(value) + ) + return dt.datetime.fromtimestamp(seconds, dt.timezone.utc).astimezone() + except (OverflowError, OSError, ValueError): + return None + raw = str(value).strip() + if not raw: + return None + if raw.isdigit(): + return parse_event_time(int(raw)) + try: + parsed = dt.datetime.fromisoformat(raw.replace("Z", "+00:00")) + except ValueError: + return None + if parsed.tzinfo is None: + parsed = parsed.replace(tzinfo=dt.timezone.utc) + return parsed.astimezone() + + +def local_day(value: Any) -> str: + parsed = parse_event_time(value) + return parsed.date().isoformat() if parsed else local_date_string() + + +def utc_month_key(value: Any) -> str: + parsed = parse_event_time(value) + if parsed is None: + return "" + return parsed.astimezone(dt.timezone.utc).strftime("%Y-%m") + + +# ------------------------------------------------------------------ database + + +def open_store(path: Path) -> tuple[sqlite3.Connection | None, Path | None]: + """Open a private snapshot of the session store. + + The CLI writes to this database continuously in WAL mode. Raw-copying the + main file and its -wal sibling cannot be made safe: a checkpoint between + the two copies pairs an old database with a newer WAL generation. SQLite's + online backup API takes the read lock it needs to produce a consistent + copy instead, without ever blocking the writer. The caller must remove the + returned temp dir once the connection is closed. + """ + if not path.exists(): + return None, None + try: + if path.stat().st_size > MAX_DB_BYTES: + return None, None + except OSError: + return None, None + + tmpdir = None + source = None + try: + tmpdir = Path(tempfile.mkdtemp(prefix="omarchy-copilot-")) + source = sqlite3.connect(f"file:{path}?mode=ro", uri=True, timeout=5.0) + target = sqlite3.connect(tmpdir / "snapshot.db") + with target: + source.backup(target) + target.row_factory = sqlite3.Row + return target, tmpdir + except (OSError, sqlite3.Error): + if tmpdir is not None: + shutil.rmtree(tmpdir, ignore_errors=True) + return None, None + finally: + if source is not None: + try: + source.close() + except sqlite3.Error: + pass + + +def has_table(conn: sqlite3.Connection, name: str) -> bool: + try: + row = conn.execute( + "SELECT 1 FROM sqlite_master WHERE type='table' AND name=? LIMIT 1", (name,) + ).fetchone() + return row is not None + except sqlite3.Error: + return False + + +def blank_stats() -> dict[str, Any]: + return { + "todayPrompts": 0, + "todaySessions": 0, + "todayTotalTokens": 0, + "todayTokensByModel": {}, + "recentDays": [ + {"date": day, "messageCount": 0} for day in recent_date_strings() + ], + "modelUsage": {}, + "totalPrompts": 0, + "totalSessions": 0, + "activeDays": 0, + "activeDates": [], + "monthCredits": 0.0, + } + + +# --------------------------------------------------------------- aggregate + + +class UsageTally: + """Running aggregate shared by both scans. + + It grows with days, sessions, and models, never with raw events, so a scan + can stream an entire store or transcript directory through it. + """ + + def __init__(self) -> None: + self.today = local_date_string() + self.recent = { + day: {"date": day, "messageCount": 0} for day in recent_date_strings() + } + self.current_utc_month = dt.datetime.now(dt.timezone.utc).strftime("%Y-%m") + self.sessions: set[str] = set() + self.active_dates: set[str] = set() + self.today_sessions: set[str] = set() + self.today_prompts = 0 + self.today_total_tokens = 0 + self.today_tokens_by_model: dict[str, int] = {} + self.total_prompts = 0 + self.model_usage: dict[str, dict[str, int]] = {} + self.month_credits = 0.0 + + def prompt(self, session_key: str, day: str) -> None: + self.total_prompts += 1 + self.sessions.add(session_key) + self.active_dates.add(day) + if day == self.today: + self.today_prompts += 1 + self.today_sessions.add(session_key) + + def usage(self, model: str, day: str, tokens: dict[str, int]) -> None: + total = sum(tokens[key] for key in TOKEN_KEYS) + if total <= 0: + return + bucket = self.model_usage.setdefault(model, empty_bucket()) + for key in TOKEN_KEYS: + bucket[key] += tokens[key] + if day in self.recent: + self.recent[day]["messageCount"] += total + if day == self.today: + self.today_total_tokens += total + self.today_tokens_by_model[model] = ( + self.today_tokens_by_model.get(model, 0) + total + ) + + def spend(self, month_key: str, credits: float) -> None: + # The allowance window is the current UTC month, matching the reset + # boundary; older spend belongs to windows that are already over. + if month_key == self.current_utc_month: + self.month_credits += credits + + def stats(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": [self.recent[day] for day in recent_date_strings()], + "modelUsage": self.model_usage, + "totalPrompts": self.total_prompts, + "totalSessions": len(self.sessions), + "activeDays": len(self.active_dates), + "activeDates": sorted(self.active_dates), + "monthCredits": round(self.month_credits, 6), + } + + +# ---------------------------------------------------------- store scan +# +# The session store's assistant_usage_events table records one row per billed +# assistant request with day, model, session, the token split, and the cost in +# total_nano_aiu. Rows stream through the tally rather than being fetched, so +# lifetime totals cover the whole store at constant memory, and timestamps go +# through the same lenient parser the transcript fallback uses. + + +def scan_store(conn: sqlite3.Connection) -> dict[str, Any]: + tally = UsageTally() + try: + rows = conn.execute( + "SELECT session_id, model, created_at, total_nano_aiu, input_tokens," + " output_tokens, cache_read_tokens, cache_write_tokens" + " FROM assistant_usage_events" + ) + for row in rows: + session_id = clean_name(row["session_id"]) + day = local_day(row["created_at"]) + tally.prompt(session_id or "unknown", day) + tally.spend( + utc_month_key(row["created_at"]), + credits_from_nano(row["total_nano_aiu"]), + ) + tally.usage( + clean_name(row["model"]) or "copilot", + day, + token_split( + row["input_tokens"], + row["output_tokens"], + row["cache_read_tokens"], + row["cache_write_tokens"], + ), + ) + except sqlite3.Error: + return blank_stats() + return tally.stats() + + +# ------------------------------------------------------ transcript fallback +# +# A CLI whose store predates the assistant_usage_events table still writes one +# events.jsonl per session, one {"type", "timestamp", "data"} event per line: +# +# - user.message is a prompt. data.source marks messages a skill, a subagent, +# or the CLI itself injected; those are not the user's prompts. +# - session.shutdown carries data.modelMetrics: per-model token usage and cost +# (totalNanoAiu), cumulative for the session across resumes. The growth since +# the previous shutdown is the usage of the segment that just ended, spread +# over the days that segment's replies landed on. +# - assistant.message carries data.outputTokens per reply. It dates a segment's +# usage, and stands in for a segment still running with no shutdown yet. + + +def is_prompt(event_type: Any, data: dict[str, Any]) -> bool: + return event_type == "user.message" and clean_name(data.get("source")) in ( + "", + "user", + ) + + +def shutdown_metrics(data: dict[str, Any]) -> dict[str, dict[str, float]]: + metrics = data.get("modelMetrics") + if not isinstance(metrics, dict): + return {} + totals: dict[str, dict[str, float]] = {} + for model, entry in metrics.items(): + if not isinstance(entry, dict): + continue + usage = entry.get("usage") + if not isinstance(usage, dict): + usage = {} + totals[clean_name(model) or "copilot"] = { + **token_split( + usage.get("inputTokens"), + usage.get("outputTokens"), + usage.get("cacheReadTokens"), + usage.get("cacheWriteTokens"), + ), + "credits": credits_from_nano(entry.get("totalNanoAiu")), + } + return totals + + +def apportion(total: int, weights: list[int]) -> list[int]: + """Split a whole total in proportion to positive weights, summing exactly.""" + whole = sum(weights) + shares: list[int] = [] + given = 0 + running = 0 + for weight in weights: + running += weight + share = total * running // whole - given + shares.append(share) + given += share + return shares + + +Reply = tuple[str, str, str, int] # model, local day, UTC month, output tokens + + +def spread_segment( + tally: UsageTally, + model: str, + delta: dict[str, int], + credits: float, + replies: list[Reply], + shutdown_slot: tuple[str, str], +) -> None: + """Book one model's usage for a finished segment on the days it happened. + + A segment can span midnight, and its shutdown only reports totals. Output + splits by each day's reply output, which the transcript records exactly. + Input, cache, and cost split by reply count: every reply is one model call + that resends the context, so calls, not output, track what was read. + """ + slots: dict[tuple[str, str], list[int]] = {} + for reply_model, day, month_key, output in replies: + if reply_model == model: + slot = slots.setdefault((day, month_key), [0, 0]) + slot[0] += 1 + slot[1] += output + if not slots: + slots[shutdown_slot] = [1, 0] + + calls = [counts[0] for counts in slots.values()] + outputs = [counts[1] for counts in slots.values()] + if sum(outputs) <= 0: + outputs = calls + split = { + key: apportion(delta[key], outputs if key == "outputTokens" else calls) + for key in TOKEN_KEYS + } + all_calls = sum(calls) + for index, (day, month_key) in enumerate(slots): + tally.usage(model, day, {key: split[key][index] for key in TOKEN_KEYS}) + tally.spend(month_key, credits * calls[index] / all_calls) + + +def scan_transcript(event_file: Path, session_key: str, tally: UsageTally) -> None: + try: + session_mtime = event_file.stat().st_mtime + except OSError: + session_mtime = 0.0 + + # Session totals as of the last shutdown, and the replies since then. + reported: dict[str, dict[str, float]] = {} + replies: list[Reply] = [] + + with event_file.open("r", encoding="utf-8", errors="replace") as handle: + for raw in handle: + try: + entry = json.loads(raw) + except ValueError: + continue + if not isinstance(entry, dict): + continue + event_type = entry.get("type") + data = entry.get("data") + if not isinstance(data, dict): + continue + + parsed = parse_event_time(entry.get("timestamp")) + if parsed is None and session_mtime > 0: + parsed = dt.datetime.fromtimestamp( + session_mtime, dt.timezone.utc + ).astimezone() + day = parsed.date().isoformat() if parsed else tally.today + month_key = ( + parsed.astimezone(dt.timezone.utc).strftime("%Y-%m") if parsed else "" + ) + + if is_prompt(event_type, data): + tally.prompt(session_key, day) + elif event_type == "assistant.message": + model = clean_name(data.get("model")) or "copilot" + replies.append( + (model, day, month_key, number(data.get("outputTokens"))) + ) + elif event_type == "session.shutdown": + for model, totals in shutdown_metrics(data).items(): + before = reported.get(model, {}) + spread_segment( + tally, + model, + { + key: max(0, int(totals[key] - before.get(key, 0))) + for key in TOKEN_KEYS + }, + max(0.0, totals["credits"] - before.get("credits", 0.0)), + replies, + (day, month_key), + ) + reported[model] = totals + # The shutdown totals already count these replies. + replies.clear() + + for model, day, _month_key, output in replies: + tally.usage(model, day, {**empty_bucket(), "outputTokens": output}) + + +def scan_transcripts(session_root: Path) -> dict[str, Any]: + if not session_root.is_dir(): + return blank_stats() + tally = UsageTally() + for session_dir in sorted(p for p in session_root.iterdir() if p.is_dir()): + event_file = session_dir / "events.jsonl" + if not event_file.is_file(): + continue + try: + scan_transcript(event_file, session_dir.name, tally) + except OSError: + continue + return tally.stats() + + +# ------------------------------------------------------------------- cache + + +def scan_cache_paths(home_dir: Path) -> tuple[Path, Path]: + digest = hashlib.sha1(str(home_dir).encode("utf-8")).hexdigest()[:16] + root = cache_root() + return root / f"copilot-scan-{digest}.json", root / f"copilot-scan-{digest}.lock" + + +def read_fresh_json(path: Path, max_age_seconds: float) -> dict[str, Any] | None: + if max_age_seconds <= 0 or not path.exists(): + return None + try: + 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: 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=(",", ":")) + "\n") + tmp.chmod(0o644) + tmp.replace(path) + except BaseException: + tmp.unlink(missing_ok=True) + raise + + +def read_cached_stats( + cache_file: Path, max_age_seconds: float, today: str +) -> 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 + 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": local_date_string(), "stats": stats}, + ) + except Exception as exc: + print( + f"omarchy-agent-usage-copilot: could not write usage cache ({exc})", + file=sys.stderr, + ) + + +def cached_scan(home_dir: Path, max_age_seconds: float) -> dict[str, Any]: + try: + cache_file, lock_file = scan_cache_paths(home_dir) + today = local_date_string() + cached = read_cached_stats(cache_file, max_age_seconds, today) + if cached is not None: + return cached + with lock_file.open("w") as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + cached = read_cached_stats(cache_file, max_age_seconds, today) + if cached is not None: + return cached + stats = direct_scan(home_dir) + write_cached_stats(cache_file, stats) + return stats + except Exception as exc: + print( + f"omarchy-agent-usage-copilot: cache unavailable ({exc}); scanning directly", + file=sys.stderr, + ) + return direct_scan(home_dir) + + +def direct_scan(home_dir: Path) -> dict[str, Any]: + conn, tmpdir = open_store(home_dir / "session-store.db") + if conn is None: + return scan_transcripts(home_dir / "session-state") + try: + if has_table(conn, "assistant_usage_events"): + stats = scan_store(conn) + if number(stats.get("totalPrompts")) > 0: + return stats + # A store without the usage table, or one that is still empty, falls back + # to the transcripts exactly like a store-less machine. + return scan_transcripts(home_dir / "session-state") + finally: + conn.close() + if tmpdir is not None: + shutil.rmtree(tmpdir, ignore_errors=True) + + +# ------------------------------------------------------------ remote quota + + +def github_token() -> str: + """Find a GitHub token without ever storing or logging one. + + Environment first, so a user can scope exactly what this sees, then the + `gh` CLI's own credential store, which is where a machine that has signed + in already keeps it. + """ + for var in ("COPILOT_QUOTA_TOKEN", "GH_TOKEN", "GITHUB_TOKEN"): + value = os.environ.get(var) + if value and value.strip(): + return value.strip() + exe = shutil.which("gh") + if not exe: + return "" + try: + result = subprocess.run( + [exe, "auth", "token"], + capture_output=True, + text=True, + timeout=REMOTE_TIMEOUT, + check=False, + ) + except (OSError, subprocess.SubprocessError): + return "" + return result.stdout.strip() if result.returncode == 0 else "" + + +def api_plan(raw: Any) -> tuple[str, str]: + key = API_PLAN_KEYS.get(str(raw or "").strip().lower(), "") + if not key: + return "", "" + return key, PLAN_LABELS.get(key, "Copilot") + + +def fetch_remote_quota() -> dict[str, Any]: + """Ask GitHub for the account-wide credit balance. + + Every failure is soft: a missing token, an offline machine, or a changed + payload leaves the record on its local estimate rather than failing the + whole collection. + """ + token = github_token() + if not token: + return {"available": False, "reason": "no-token"} + + request = urllib.request.Request( + QUOTA_URL, + headers={ + "Authorization": f"Bearer {token}", + "Accept": "application/json", + "User-Agent": "omarchy-agent-usage", + }, + ) + try: + with urllib.request.urlopen(request, timeout=REMOTE_TIMEOUT) as response: + payload = json.loads( + response.read(2_000_000).decode("utf-8", errors="replace") + ) + except urllib.error.HTTPError as exc: + return {"available": False, "reason": f"http-{exc.code}"} + except (urllib.error.URLError, OSError, ValueError, TimeoutError): + return {"available": False, "reason": "unreachable"} + + if not isinstance(payload, dict): + return {"available": False, "reason": "bad-payload"} + + snapshots = payload.get("quota_snapshots") + snapshot = ( + snapshots.get("premium_interactions") if isinstance(snapshots, dict) else None + ) + if not isinstance(snapshot, dict): + return {"available": False, "reason": "no-quota"} + + plan_key, plan_label = api_plan(payload.get("copilot_plan")) + orgs = payload.get("organization_login_list") + # Under token-based billing the snapshot counts AI credits directly; only + # the local store and transcripts count in nano units. + return { + "available": True, + "unlimited": bool(snapshot.get("unlimited")), + "entitlement": positive_float(snapshot.get("entitlement")), + "used": positive_float(snapshot.get("credits_used")), + "resetsAt": str(payload.get("quota_reset_date_utc") or ""), + "plan": plan_key, + "planLabel": plan_label, + "org": str(orgs[0]) if isinstance(orgs, list) and orgs else "", + "login": clean_name(payload.get("login")), + } + + +# ---------------------------------------------------------------- allowance + + +def resolve_plan(config: dict[str, Any]) -> tuple[str, float]: + """Return the configured plan key and its monthly allowance in credits. + + The plan is not discoverable locally, so a missing or unrecognized plan + stays unknown (key "") instead of defaulting to a tier the user may not + have; only monthlyCredits can then give the estimate an allowance. + """ + raw = str(config.get("plan") or "").strip().lower() + key = raw.replace(" ", "").replace("copilot", "") + if key in ("proplus", "pro-plus"): + key = "pro+" + if key not in PLAN_LABELS: + key = "" + + override = config.get("monthlyCredits") + if isinstance(override, (int, float)) and override > 0: + return key, float(override) + return key, float(PLAN_ALLOWANCES.get(key, 0)) + + +def allowance_limits( + month_credits: float, + config: dict[str, Any], + remote: dict[str, Any] | None, + now: dt.datetime, + has_stats: bool, +) -> dict[str, Any]: + plan_key, plan_allowance = resolve_plan(config) + plan_label = PLAN_LABELS.get(plan_key, "") + estimated = True + allow = plan_allowance + used = month_credits + resets_at = next_month_start(now).isoformat() + reason = "" + + live = remote if isinstance(remote, dict) and remote.get("available") else None + if live and not live.get("unlimited") and float(live.get("entitlement") or 0) > 0: + allow = float(live["entitlement"]) + used = float(live.get("used") or 0.0) + estimated = False + if live.get("plan"): + plan_key = str(live["plan"]) + plan_label = str(live["planLabel"] or PLAN_LABELS.get(plan_key, "Copilot")) + raw_reset = str(live.get("resetsAt") or "") + if raw_reset: + parsed = parse_event_time(raw_reset) + if parsed is not None and parsed > now: + resets_at = parsed.isoformat() + elif ( + remote is not None + and remote.get("available") is False + and remote.get("reason") != "no-token" + ): + reason = "Copilot limits unavailable" + + limits = [] + if allow > 0: + # A plan-default allowance must not advertise a meter on a machine the + # CLI has never touched — the local estimate needs local numbers to + # mean anything. The live account figure stands on its own. + if estimated and not has_stats: + pass + else: + percent = min(1.0, used / allow) + limits.append( + { + "label": "Monthly allowance (est.)" + if estimated + else "Monthly allowance", + "title": "Monthly allowance (est.)" + if estimated + else "Monthly allowance", + "percent": percent, + "resetsAt": resets_at, + } + ) + + return { + "limits": limits, + "tierLabel": plan_label, + "usageStatusText": reason, + "authHelpText": AUTH_HELP, + } + + +# ------------------------------------------------------------------ record + + +def main() -> int: + parser = argparse.ArgumentParser() + # --force rescans everything and rewrites the cache. --limits-only is kept + # for CLI compatibility with the panel's refreshLimits() call: only the + # limits probe must be fresh, so it may reuse a scan for far longer than a + # normal run, whose short window exists purely to dedup concurrent + # collector runs. + parser.add_argument("--force", action="store_true") + parser.add_argument("--limits-only", action="store_true") + parser.add_argument( + "--copilot-home", default=os.environ.get("COPILOT_HOME") or "~/.copilot" + ) + args = parser.parse_args() + + home_dir = expand_path(args.copilot_home) + max_age = ( + 0 + if args.force + else (LIMITS_ONLY_REUSE_SECONDS if args.limits_only else SCAN_REUSE_SECONDS) + ) + stats = cached_scan(home_dir, max_age) + + config = read_config() + now = dt.datetime.now(dt.timezone.utc) + + remote = {"available": False, "reason": "disabled"} + if config.get("remote", True) is not False: + remote = fetch_remote_quota() + + has_stats = ( + number(stats.get("totalPrompts")) > 0 + or float(stats.get("monthCredits") or 0.0) > 0 + ) + allowance = allowance_limits( + float(stats.get("monthCredits") or 0.0), config, remote, now, has_stats + ) + + record: dict[str, Any] = { + "schemaVersion": 1, + "id": AGENT_ID, + "name": AGENT_NAME, + "updatedAt": now.isoformat(), + "ready": number(stats.get("totalPrompts")) > 0 or len(allowance["limits"]) > 0, + "hasLocalStats": True, + "tierLabel": allowance["tierLabel"], + "usageStatusText": allowance["usageStatusText"], + "authHelpText": allowance["authHelpText"], + "limits": allowance["limits"], + } + for key in ( + "todayPrompts", + "todaySessions", + "todayTotalTokens", + "totalPrompts", + "totalSessions", + "activeDays", + ): + record[key] = number(stats.get(key)) + record["todayTokensByModel"] = stats.get("todayTokensByModel") or {} + record["recentDays"] = stats.get("recentDays") or [] + record["modelUsage"] = stats.get("modelUsage") or {} + record["activeDates"] = stats.get("activeDates") or [] + 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 dd21382d710..bccf9d6a95f 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. | `cursor` | Cursor's dashboard RPCs: included, auto-model, named-model, and on-demand meters | the same dashboard RPCs, one call per day for the last week | | `opencode` | OpenCode Go's usage endpoint (rolling + weekly + monthly), with an API key or the Console session | pi and omp sessions on the Zen and Go providers, plus opencode's own `message` and `session_message` stores | | `fireworks` | Estimated prepaid balance: configured funding minus rated account costs | Fireworks billing API, grouped by day and model for the last 30 days | +| `copilot` | Account-wide AI credit allowance (GitHub quota endpoint, or estimated locally) | the Copilot CLI session store's `assistant_usage_events`, with `~/.copilot/session-state` transcripts as fallback | When `~/.local/state/omarchy/agents/accounts/.json` registers more than one account, the `claude`, `codex`, and `grok` records @@ -140,6 +141,49 @@ seconds, each credential caches its own, a failed check keeps the last good windows, dimmed, until they reset, and a rate-limited window reads as full. Without a Go plan the panel still shows local token usage. +### Copilot allowance + +Copilot reads the CLI's session store under `~/.copilot`, honoring +`COPILOT_HOME`. + +GitHub exposes no supported personal-account usage API — its Copilot REST +endpoints are organization- or enterprise-scoped — so the collector reads what +the Copilot CLI keeps locally. Its session store (`~/.copilot/session-store.db`) +records every billed request in an `assistant_usage_events` table: day, model, +session, the token split, and the cost in nano AI units (1e9 nano units = 1 +AI credit ≈ $0.01). The panel's token sections come from that table; CLI +versions whose store predates it are served by the session transcripts under +`~/.copilot/session-state` instead, whose `session.shutdown` events carry the +same per-model token split and cost. Neither source is network-dependent. + +The monthly meter is an allowance of AI credits. The account-wide figure +comes from GitHub's internal quota endpoint (the same one the editor plugins +ask), which counts every machine, IDE, and agent on the entitlement — not just +this machine's CLI. It needs a token from `COPILOT_QUOTA_TOKEN`, `GH_TOKEN`, +`GITHUB_TOKEN`, or `gh auth token`; every failure is soft and falls back to +the local estimate. The included allowance is not discoverable locally, so it +is configured in `~/.config/omarchy/agents/copilot.json`: + +```json +{ + "plan": "pro", + "monthlyCredits": 1500, + "remote": true +} +``` + +`plan` is one of `free`, `student`, `pro`, `pro+`, `max`, `business`, +`enterprise` and picks the default allowance (Pro 1500, Pro+ 7000, Max 20000, +Business 1900, Enterprise 3900 credits per month; existing seat-based +customers may see promotional amounts through September 2026). +`monthlyCredits` overrides the plan table, for AI-credit plans without a +published number or for a personal budget. With neither set there is no local +estimate and no plan label, because the plan cannot be read from disk. `remote` +(default true) turns the account-wide quota probe on or off. The allowance +window resets on the 1st of each month at 00:00 UTC, GitHub's own boundary. A +limit labeled "(est.)" is the local estimate — this machine's CLI spend only; a +plain "Monthly allowance" is the live account figure. + ### Fireworks balance The collector first asks the account's `:getBalance` endpoint for the real diff --git a/shell/plugins/agents/assets/copilot-light.svg b/shell/plugins/agents/assets/copilot-light.svg new file mode 100644 index 00000000000..717e3d72142 --- /dev/null +++ b/shell/plugins/agents/assets/copilot-light.svg @@ -0,0 +1 @@ +GitHub Copilot \ No newline at end of file diff --git a/shell/plugins/agents/assets/copilot.svg b/shell/plugins/agents/assets/copilot.svg new file mode 100644 index 00000000000..4edbc713f39 --- /dev/null +++ b/shell/plugins/agents/assets/copilot.svg @@ -0,0 +1 @@ +GitHub Copilot \ No newline at end of file diff --git a/test/shell.d/agent-usage-copilot-test.sh b/test/shell.d/agent-usage-copilot-test.sh new file mode 100755 index 00000000000..bac95a6838a --- /dev/null +++ b/test/shell.d/agent-usage-copilot-test.sh @@ -0,0 +1,294 @@ +#!/bin/bash + +source "$(dirname "$0")/base-test.sh" + +require_command jq +require_command python3 + +# Fixtures are stamped in UTC and checked against the local day and the UTC +# allowance month; pinning the zone keeps both on the same date at any hour. +export TZ=UTC + +TEST_HOME=$(mktemp -d) +EMPTY_HOME=$(mktemp -d) +TRANSCRIPT_HOME=$(mktemp -d) +MIDNIGHT_HOME=$(mktemp -d) +HISTORY_HOME=$(mktemp -d) +trap 'rm -rf "$TEST_HOME" "$EMPTY_HOME" "$TRANSCRIPT_HOME" "$MIDNIGHT_HOME" "$HISTORY_HOME"' EXIT + +for home in "$TEST_HOME" "$EMPTY_HOME" "$TRANSCRIPT_HOME" "$MIDNIGHT_HOME" "$HISTORY_HOME"; do + mkdir -p "$home/.copilot/session-state" "$home/.config/omarchy/agents" "$home/.cache" + printf '{"plan": "pro", "remote": false}' >"$home/.config/omarchy/agents/copilot.json" +done +rm -rf "$EMPTY_HOME/.copilot" + +# Every collector run below re-derives "today" and the allowance month from +# the clock, and compares them with the stamps taken here. Starting too close to +# UTC midnight (under TZ=UTC also the month boundary) lets them diverge mid-run, +# so wait for the new day instead. +seconds_to_midnight=$(( 86400 - $(date -u +%s) % 86400 )) +if (( seconds_to_midnight < 120 )); then + sleep "$seconds_to_midnight" +fi + +now=$(date -u +%Y-%m-%dT%H:%M:%SZ) +yesterday=$(date -u -d '1 day ago' +%Y-%m-%dT%H:%M:%SZ) +today=$(date +%F) + +run_collector_in() { + local home=$1 + shift + HOME="$home" COPILOT_HOME="$home/.copilot" XDG_CONFIG_HOME="$home/.config" \ + XDG_CACHE_HOME="$home/.cache" "$ROOT/bin/omarchy-agent-usage-copilot" --force "$@" +} + +# The session store carries one row per billed assistant request. As in the +# real CLI, input_tokens is the whole prompt (cache reads and writes included) +# and output_tokens already includes reasoning_tokens. +python3 - "$TEST_HOME/.copilot/session-store.db" "$now" <<'PY' +import sqlite3 +import sys + +db = sys.argv[1] +stamp = sys.argv[2] +conn = sqlite3.connect(db) +conn.execute("CREATE TABLE assistant_usage_events (id INTEGER PRIMARY KEY AUTOINCREMENT, session_id TEXT, created_at TEXT, model TEXT, total_nano_aiu INTEGER, input_tokens INTEGER, output_tokens INTEGER, cache_read_tokens INTEGER, cache_write_tokens INTEGER, reasoning_tokens INTEGER, duration_ms INTEGER)") +conn.executemany( + "INSERT INTO assistant_usage_events (session_id, created_at, model, total_nano_aiu, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens, reasoning_tokens, duration_ms) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + [ + ("sess-a", stamp, "claude-sonnet-4-5", 500_000_000, 1600, 300, 400, 0, 50, 3000), + ("sess-a", stamp, "claude-sonnet-4-5", 500_000_000, 900, 200, 100, 0, 0, 2000), + ("sess-b", stamp, "gpt-5.2", 500_000_000, 1900, 700, 100, 300, 100, 5000), + ], +) +conn.commit() +conn.close() +PY + +result=$(run_collector_in "$TEST_HOME") + +[[ $(jq -r '.id' <<<"$result") == "copilot" ]] || + fail "Copilot collector identifies itself" "$result" +pass "Copilot collector identifies itself" + +[[ $(jq -r '.ready' <<<"$result") == "true" ]] || + fail "Copilot collector is ready with usage" "$result" +pass "Copilot collector is ready with usage" + +[[ $(jq -r '.todayPrompts' <<<"$result") == "3" ]] || + fail "Copilot collector counts billed requests" "$result" +pass "Copilot collector counts billed requests" + +[[ $(jq -r '.todaySessions' <<<"$result") == "2" ]] || + fail "Copilot collector counts distinct sessions" "$result" +pass "Copilot collector counts distinct sessions" + +[[ $(jq -r '.todayTotalTokens' <<<"$result") == "5600" ]] || + fail "Copilot collector sums each token once" "$result" +pass "Copilot collector sums each token once" + +[[ $(jq -c '.modelUsage["claude-sonnet-4-5"]' <<<"$result") == '{"inputTokens":2000,"outputTokens":500,"cacheReadInputTokens":500,"cacheCreationInputTokens":0}' ]] || + fail "Copilot collector keeps reasoning inside output" "$result" +pass "Copilot collector keeps reasoning inside output" + +[[ $(jq -c '.modelUsage["gpt-5.2"]' <<<"$result") == '{"inputTokens":1500,"outputTokens":700,"cacheReadInputTokens":100,"cacheCreationInputTokens":300}' ]] || + fail "Copilot collector splits cached tokens off the prompt" "$result" +pass "Copilot collector splits cached tokens off the prompt" + +[[ $(jq -r '.recentDays[-1].date' <<<"$result") == "$today" ]] || + fail "Copilot collector reports today's day row" "$result" +pass "Copilot collector reports today's day row" + +[[ $(jq -r '.tierLabel' <<<"$result") == "Pro" ]] || + fail "Copilot collector exposes the configured plan label" "$result" +pass "Copilot collector exposes the configured plan label" + +[[ $(jq -r '.limits[0].label' <<<"$result") == "Monthly allowance (est.)" ]] || + fail "Copilot collector labels the local estimate" "$result" +pass "Copilot collector labels the local estimate" + +[[ $(jq -r '.limits[0].percent' <<<"$result") == "0.001" ]] || + fail "Copilot collector divides month credits by allowance" "$result" +pass "Copilot collector divides month credits by allowance" + +[[ $(jq -r '.limits[0].resetsAt' <<<"$result") != "" ]] || + fail "Copilot collector sets an allowance reset time" "$result" +pass "Copilot collector sets an allowance reset time" + +# The plan cannot be read from disk, so an unconfigured one is not presented as +# a tier: no label and no estimated meter, while the local stats still show. +printf '{"remote": false}' >"$TEST_HOME/.config/omarchy/agents/copilot.json" +result=$(run_collector_in "$TEST_HOME") + +[[ $(jq -r '.ready' <<<"$result") == "true" && $(jq -r '.tierLabel' <<<"$result") == "" && $(jq -c '.limits' <<<"$result") == "[]" ]] || + fail "Copilot collector leaves an unconfigured plan unknown" "$result" +pass "Copilot collector leaves an unconfigured plan unknown" + +printf '{"monthlyCredits": 3000, "remote": false}' >"$TEST_HOME/.config/omarchy/agents/copilot.json" +result=$(run_collector_in "$TEST_HOME") + +[[ $(jq -r '.tierLabel' <<<"$result") == "" && $(jq -r '.limits[0].percent' <<<"$result") == "0.0005" ]] || + fail "Copilot collector estimates from monthlyCredits without a plan" "$result" +pass "Copilot collector estimates from monthlyCredits without a plan" + +# The live quota answers in AI credits, account-wide; it wins over the local +# estimate and supplies the plan and reset date. +printf '{"plan": "pro", "remote": true}' >"$TEST_HOME/.config/omarchy/agents/copilot.json" +result=$(HOME="$TEST_HOME" COPILOT_HOME="$TEST_HOME/.copilot" XDG_CONFIG_HOME="$TEST_HOME/.config" \ + XDG_CACHE_HOME="$TEST_HOME/.cache" COPILOT_QUOTA_TOKEN="test-token" \ + python3 - "$ROOT/bin/omarchy-agent-usage-copilot" <<'PY' +import datetime as dt +import importlib.machinery +import importlib.util +import io +import json +import 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) + +reset = (dt.datetime.now(dt.timezone.utc) + dt.timedelta(days=10)).replace(microsecond=0) +payload = { + "copilot_plan": "business", + "quota_reset_date_utc": reset.strftime("%Y-%m-%dT%H:%M:%S.000Z"), + "quota_snapshots": { + "premium_interactions": { + "unlimited": False, + "token_based_billing": True, + "entitlement": 1900, + "credits_used": 475, + "remaining": 1425, + } + }, +} +collector.urllib.request.urlopen = lambda request, timeout=None: io.BytesIO(json.dumps(payload).encode()) + +out = io.StringIO() +sys.argv = ["omarchy-agent-usage-copilot", "--force"] +sys.stdout, real_stdout = out, sys.stdout +collector.main() +sys.stdout = real_stdout +print(json.dumps({"record": json.loads(out.getvalue()), "reset": reset.isoformat()})) +PY +) + +[[ $(jq -r '.record.limits[0].label' <<<"$result") == "Monthly allowance" && $(jq -r '.record.limits[0].percent' <<<"$result") == "0.25" ]] || + fail "Copilot collector reads the live allowance in credits" "$result" +pass "Copilot collector reads the live allowance in credits" + +[[ $(jq -r '.record.tierLabel' <<<"$result") == "Business" ]] || + fail "Copilot collector takes the plan from the live quota" "$result" +pass "Copilot collector takes the plan from the live quota" + +[[ $(jq -r '.record.limits[0].resetsAt' <<<"$result") == $(jq -r '.reset' <<<"$result") ]] || + fail "Copilot collector takes the reset time from the live quota" "$result" +pass "Copilot collector takes the reset time from the live quota" + +# Lifetime totals cover the whole store, however many rows it holds. +python3 - "$HISTORY_HOME/.copilot/session-store.db" "$now" <<'PY' +import datetime as dt +import sqlite3 +import sys + +db = sys.argv[1] +stamp = sys.argv[2] +old = (dt.datetime.now(dt.timezone.utc) - dt.timedelta(days=100)).strftime("%Y-%m-%dT%H:%M:%SZ") +conn = sqlite3.connect(db) +conn.execute("CREATE TABLE assistant_usage_events (id INTEGER PRIMARY KEY AUTOINCREMENT, session_id TEXT, created_at TEXT, model TEXT, total_nano_aiu INTEGER, input_tokens INTEGER, output_tokens INTEGER, cache_read_tokens INTEGER, cache_write_tokens INTEGER, reasoning_tokens INTEGER)") +rows = [(f"old-{n % 7}", old, "gpt-5.2", 0, 10, 1, 0, 0, 0) for n in range(25_000)] +rows.append(("new", stamp, "gpt-5.2", 0, 10, 1, 0, 0, 0)) +conn.executemany( + "INSERT INTO assistant_usage_events (session_id, created_at, model, total_nano_aiu, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens, reasoning_tokens) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)", + rows, +) +conn.commit() +conn.close() +PY + +result=$(run_collector_in "$HISTORY_HOME") + +[[ $(jq -r '.totalPrompts' <<<"$result") == "25001" && $(jq -r '.totalSessions' <<<"$result") == "8" && $(jq -r '.activeDays' <<<"$result") == "2" ]] || + fail "Copilot collector counts lifetime totals from the full store" "$(jq -c '{totalPrompts, totalSessions, activeDays}' <<<"$result")" +pass "Copilot collector counts lifetime totals from the full store" + +# Without a usage table the transcripts stand in. Shutdown metrics are +# cumulative across resumes and already include cached and reasoning tokens; +# replies after the last shutdown count on their output tokens alone. +mkdir -p "$TRANSCRIPT_HOME/.copilot/session-state/sess-t" +jq -nc --arg ts "$now" ' + def ev($type; $data): {type: $type, data: $data, id: "x", timestamp: $ts}; + def metrics($in; $out; $read; $write; $nano): { + "claude-sonnet-4.6": { + requests: {count: 1, cost: 1}, + usage: {inputTokens: $in, outputTokens: $out, cacheReadTokens: $read, cacheWriteTokens: $write, reasoningTokens: 20}, + totalNanoAiu: $nano + } + }; + ev("session.start"; {sessionId: "sess-t"}), + ev("user.message"; {content: "first prompt"}), + ev("user.message"; {content: "injected", source: "skill-demo"}), + ev("assistant.message"; {model: "claude-sonnet-4.6", outputTokens: 100}), + ev("session.shutdown"; {shutdownType: "routine", modelMetrics: metrics(1000; 200; 600; 100; 2000000000)}), + ev("session.resume"; {}), + ev("user.message"; {content: "second prompt", source: "user"}), + ev("assistant.message"; {model: "claude-sonnet-4.6", outputTokens: 50}), + ev("session.shutdown"; {shutdownType: "routine", modelMetrics: metrics(1500; 300; 900; 100; 3000000000)}), + ev("user.message"; {content: "third prompt"}), + ev("assistant.message"; {model: "claude-sonnet-4.6", outputTokens: 40}) +' >"$TRANSCRIPT_HOME/.copilot/session-state/sess-t/events.jsonl" + +result=$(run_collector_in "$TRANSCRIPT_HOME") + +[[ $(jq -r '.ready' <<<"$result") == "true" && $(jq -r '.todayPrompts' <<<"$result") == "3" && $(jq -r '.todaySessions' <<<"$result") == "1" ]] || + fail "Copilot transcript fallback counts the user's prompts" "$result" +pass "Copilot transcript fallback counts the user's prompts" + +[[ $(jq -c '.modelUsage["claude-sonnet-4.6"]' <<<"$result") == '{"inputTokens":500,"outputTokens":340,"cacheReadInputTokens":900,"cacheCreationInputTokens":100}' ]] || + fail "Copilot transcript fallback reads shutdown usage once" "$result" +pass "Copilot transcript fallback reads shutdown usage once" + +[[ $(jq -r '.todayTotalTokens' <<<"$result") == "1840" ]] || + fail "Copilot transcript fallback sums today's tokens" "$result" +pass "Copilot transcript fallback sums today's tokens" + +[[ $(jq -r '.limits[0].percent' <<<"$result") == "0.002" ]] || + fail "Copilot transcript fallback estimates spend from shutdown cost" "$result" +pass "Copilot transcript fallback estimates spend from shutdown cost" + +# A segment that spans midnight books its shutdown totals on the days its +# replies landed on: output by each day's reply output (40 / 60), everything +# else by reply count (2 / 1). +mkdir -p "$MIDNIGHT_HOME/.copilot/session-state/sess-m" +jq -nc --arg ts "$now" --arg before "$yesterday" ' + def ev($type; $data; $at): {type: $type, data: $data, id: "x", timestamp: $at}; + ev("user.message"; {content: "late prompt"}; $before), + ev("assistant.message"; {model: "claude-sonnet-4.6", outputTokens: 30}; $before), + ev("assistant.message"; {model: "claude-sonnet-4.6", outputTokens: 10}; $before), + ev("assistant.message"; {model: "claude-sonnet-4.6", outputTokens: 60}; $ts), + ev("session.shutdown"; {shutdownType: "routine", modelMetrics: { + "claude-sonnet-4.6": { + usage: {inputTokens: 1300, outputTokens: 100, cacheReadTokens: 600, cacheWriteTokens: 100, reasoningTokens: 0}, + totalNanoAiu: 0 + } + }}; $ts) +' >"$MIDNIGHT_HOME/.copilot/session-state/sess-m/events.jsonl" + +result=$(run_collector_in "$MIDNIGHT_HOME") + +[[ $(jq -r '.recentDays[-2].messageCount' <<<"$result") == "906" && $(jq -r '.recentDays[-1].messageCount' <<<"$result") == "494" && $(jq -r '.todayTotalTokens' <<<"$result") == "494" ]] || + fail "Copilot transcript fallback books a midnight-spanning segment on both days" "$(jq -c '{recentDays, todayTotalTokens}' <<<"$result")" +pass "Copilot transcript fallback books a midnight-spanning segment on both days" + +[[ $(jq -c '.modelUsage["claude-sonnet-4.6"]' <<<"$result") == '{"inputTokens":600,"outputTokens":100,"cacheReadInputTokens":600,"cacheCreationInputTokens":100}' ]] || + fail "Copilot transcript fallback keeps a split segment's totals whole" "$result" +pass "Copilot transcript fallback keeps a split segment's totals whole" + +# A machine that never ran the CLI reports nothing and stays hidden: no meter, +# no ready. +result=$(run_collector_in "$EMPTY_HOME") + +[[ $(jq -r '.ready' <<<"$result") == "false" && $(jq -c '.limits' <<<"$result") == "[]" ]] || + fail "Copilot collector stays hidden on a clean machine" "$result" +pass "Copilot collector stays hidden on a clean machine" From c5e5bffe2dbeced52aae9efe7c133b62dddce9f7 Mon Sep 17 00:00:00 2001 From: Azonnali Date: Thu, 1 Oct 2026 22:03:37 -0400 Subject: [PATCH 017/546] 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 ca0021e5302d2a9a94a31e0de92bb1d84c46f069 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Thu, 1 Oct 2026 22:04:12 -0400 Subject: [PATCH 018/546] Indent the Copilot collector with two spaces Match the rest of the repo's Python. Whitespace only. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-copilot | 1428 +++++++++++++++---------------- 1 file changed, 714 insertions(+), 714 deletions(-) diff --git a/bin/omarchy-agent-usage-copilot b/bin/omarchy-agent-usage-copilot index 1a8eb52b3c0..e060cb31a95 100755 --- a/bin/omarchy-agent-usage-copilot +++ b/bin/omarchy-agent-usage-copilot @@ -24,11 +24,11 @@ machine or a missing token degrades to the local estimate instead of erroring. Configuration lives in ~/.config/omarchy/agents/copilot.json: - { + { "plan": "pro", "monthlyCredits": 1500, "remote": true - } + } plan is one of free, student, pro, pro+, max, business, enterprise and picks the default allowance (Pro 1500, Pro+ 7000, Max 20000, Business 1900, @@ -68,7 +68,7 @@ from typing import Any AGENT_ID = "copilot" AGENT_NAME = "GitHub Copilot" AUTH_HELP = ( - "Run `copilot` in a terminal and sign in to GitHub to start recording sessions." + "Run `copilot` in a terminal and sign in to GitHub to start recording sessions." ) ULTRA_SECONDS = 1_000_000_000 @@ -77,12 +77,12 @@ REMOTE_TIMEOUT = 6.0 # Plan identifiers as the quota endpoint spells them. API_PLAN_KEYS = { - "individual": "pro", - "individual_pro": "pro+", - "business": "business", - "enterprise": "enterprise", - "free": "free", - "max": "max", + "individual": "pro", + "individual_pro": "pro+", + "business": "business", + "enterprise": "enterprise", + "free": "free", + "max": "max", } # Monthly AI credit allowances from GitHub's usage-based billing tables. @@ -90,23 +90,23 @@ API_PLAN_KEYS = { # practical ceiling is set by their admin rather than this figure — # monthlyCredits in the config overrides it. PLAN_ALLOWANCES = { - "free": 0, - "student": 0, - "pro": 1_500, - "pro+": 7_000, - "max": 20_000, - "business": 1_900, - "enterprise": 3_900, + "free": 0, + "student": 0, + "pro": 1_500, + "pro+": 7_000, + "max": 20_000, + "business": 1_900, + "enterprise": 3_900, } PLAN_LABELS = { - "free": "Free", - "student": "Student", - "pro": "Pro", - "pro+": "Pro+", - "max": "Max", - "business": "Business", - "enterprise": "Enterprise", + "free": "Free", + "student": "Student", + "pro": "Pro", + "pro+": "Pro+", + "max": "Max", + "business": "Business", + "enterprise": "Enterprise", } # The session store is mutable and grows without bound, so it is size-capped @@ -119,290 +119,290 @@ LIMITS_ONLY_REUSE_SECONDS = 900 def number(value: Any) -> int: - try: - n = float(value or 0) - return round(n) if n == n else 0 - except (TypeError, ValueError): - return 0 + try: + n = float(value or 0) + return round(n) if n == n else 0 + except (TypeError, ValueError): + return 0 def expand_path(value: str) -> Path: - return Path(os.path.expandvars(os.path.expanduser(value))).resolve() + return Path(os.path.expandvars(os.path.expanduser(value))).resolve() def config_path() -> Path: - config_home = Path(os.environ.get("XDG_CONFIG_HOME") or (Path.home() / ".config")) - return config_home / "omarchy" / "agents" / "copilot.json" + config_home = Path(os.environ.get("XDG_CONFIG_HOME") or (Path.home() / ".config")) + return config_home / "omarchy" / "agents" / "copilot.json" def read_config() -> dict[str, Any]: - try: - parsed = json.loads(config_path().read_text(encoding="utf-8")) - return parsed if isinstance(parsed, dict) else {} - except (OSError, ValueError): - return {} + try: + parsed = json.loads(config_path().read_text(encoding="utf-8")) + return parsed if isinstance(parsed, dict) else {} + except (OSError, ValueError): + return {} 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 TOKEN_KEYS = ( - "inputTokens", - "outputTokens", - "cacheReadInputTokens", - "cacheCreationInputTokens", + "inputTokens", + "outputTokens", + "cacheReadInputTokens", + "cacheCreationInputTokens", ) def empty_bucket() -> dict[str, int]: - return dict.fromkeys(TOKEN_KEYS, 0) + return dict.fromkeys(TOKEN_KEYS, 0) def local_date_string() -> str: - return dt.datetime.now().date().isoformat() + return dt.datetime.now().date().isoformat() def recent_date_strings() -> list[str]: - today = dt.datetime.now().date() - return [ - (today - dt.timedelta(days=offset)).isoformat() for offset in range(6, -1, -1) - ] + today = dt.datetime.now().date() + return [ + (today - dt.timedelta(days=offset)).isoformat() for offset in range(6, -1, -1) + ] def month_start(now: dt.datetime) -> dt.datetime: - return now.replace(day=1, hour=0, minute=0, second=0, microsecond=0) + return now.replace(day=1, hour=0, minute=0, second=0, microsecond=0) def next_month_start(now: dt.datetime) -> dt.datetime: - days = calendar.monthrange(now.year, now.month)[1] - return month_start(now) + dt.timedelta(days=days) + days = calendar.monthrange(now.year, now.month)[1] + return month_start(now) + dt.timedelta(days=days) def positive_float(value: Any) -> float: - try: - n = float(value or 0) - except (TypeError, ValueError): - return 0.0 - return n if n > 0 else 0.0 + try: + n = float(value or 0) + except (TypeError, ValueError): + return 0.0 + return n if n > 0 else 0.0 def credits_from_nano(nano: Any) -> float: - return positive_float(nano) / ULTRA_SECONDS + return positive_float(nano) / ULTRA_SECONDS def token_split( - prompt_tokens: Any, output_tokens: Any, cache_read: Any, cache_write: Any + prompt_tokens: Any, output_tokens: Any, cache_read: Any, cache_write: Any ) -> dict[str, int]: - """Map Copilot's token counts onto the panel's four disjoint buckets. - - Copilot's input count is the whole prompt, cache reads and writes included, - and its output count already includes reasoning tokens. The buckets are - disjoint, like the other collectors', so the cached part comes off input - and reasoning is never added a second time. - """ - read = number(cache_read) - write = number(cache_write) - return { - "inputTokens": max(0, number(prompt_tokens) - read - write), - "outputTokens": number(output_tokens), - "cacheReadInputTokens": read, - "cacheCreationInputTokens": write, - } + """Map Copilot's token counts onto the panel's four disjoint buckets. + + Copilot's input count is the whole prompt, cache reads and writes included, + and its output count already includes reasoning tokens. The buckets are + disjoint, like the other collectors', so the cached part comes off input + and reasoning is never added a second time. + """ + read = number(cache_read) + write = number(cache_write) + return { + "inputTokens": max(0, number(prompt_tokens) - read - write), + "outputTokens": number(output_tokens), + "cacheReadInputTokens": read, + "cacheCreationInputTokens": write, + } def clean_name(value: Any) -> str: - return str(value or "").strip() + return str(value or "").strip() def parse_event_time(value: Any) -> dt.datetime | None: - """Parse a store or transcript timestamp into timezone-aware local time. - - Accepts ISO-8601 (with or without Z) and epoch seconds or milliseconds. - Everything else collapses to None; the caller falls back to the local date. - """ - if value is None: - return None - if isinstance(value, (int, float)): - try: - seconds = ( - float(value) / 1000.0 if float(value) > 10_000_000_000 else float(value) - ) - return dt.datetime.fromtimestamp(seconds, dt.timezone.utc).astimezone() - except (OverflowError, OSError, ValueError): - return None - raw = str(value).strip() - if not raw: - return None - if raw.isdigit(): - return parse_event_time(int(raw)) + """Parse a store or transcript timestamp into timezone-aware local time. + + Accepts ISO-8601 (with or without Z) and epoch seconds or milliseconds. + Everything else collapses to None; the caller falls back to the local date. + """ + if value is None: + return None + if isinstance(value, (int, float)): try: - parsed = dt.datetime.fromisoformat(raw.replace("Z", "+00:00")) - except ValueError: - return None - if parsed.tzinfo is None: - parsed = parsed.replace(tzinfo=dt.timezone.utc) - return parsed.astimezone() + seconds = ( + float(value) / 1000.0 if float(value) > 10_000_000_000 else float(value) + ) + return dt.datetime.fromtimestamp(seconds, dt.timezone.utc).astimezone() + except (OverflowError, OSError, ValueError): + return None + raw = str(value).strip() + if not raw: + return None + if raw.isdigit(): + return parse_event_time(int(raw)) + try: + parsed = dt.datetime.fromisoformat(raw.replace("Z", "+00:00")) + except ValueError: + return None + if parsed.tzinfo is None: + parsed = parsed.replace(tzinfo=dt.timezone.utc) + return parsed.astimezone() def local_day(value: Any) -> str: - parsed = parse_event_time(value) - return parsed.date().isoformat() if parsed else local_date_string() + parsed = parse_event_time(value) + return parsed.date().isoformat() if parsed else local_date_string() def utc_month_key(value: Any) -> str: - parsed = parse_event_time(value) - if parsed is None: - return "" - return parsed.astimezone(dt.timezone.utc).strftime("%Y-%m") + parsed = parse_event_time(value) + if parsed is None: + return "" + return parsed.astimezone(dt.timezone.utc).strftime("%Y-%m") # ------------------------------------------------------------------ database def open_store(path: Path) -> tuple[sqlite3.Connection | None, Path | None]: - """Open a private snapshot of the session store. - - The CLI writes to this database continuously in WAL mode. Raw-copying the - main file and its -wal sibling cannot be made safe: a checkpoint between - the two copies pairs an old database with a newer WAL generation. SQLite's - online backup API takes the read lock it needs to produce a consistent - copy instead, without ever blocking the writer. The caller must remove the - returned temp dir once the connection is closed. - """ - if not path.exists(): - return None, None - try: - if path.stat().st_size > MAX_DB_BYTES: - return None, None - except OSError: - return None, None - - tmpdir = None - source = None - try: - tmpdir = Path(tempfile.mkdtemp(prefix="omarchy-copilot-")) - source = sqlite3.connect(f"file:{path}?mode=ro", uri=True, timeout=5.0) - target = sqlite3.connect(tmpdir / "snapshot.db") - with target: - source.backup(target) - target.row_factory = sqlite3.Row - return target, tmpdir - except (OSError, sqlite3.Error): - if tmpdir is not None: - shutil.rmtree(tmpdir, ignore_errors=True) - return None, None - finally: - if source is not None: - try: - source.close() - except sqlite3.Error: - pass + """Open a private snapshot of the session store. + + The CLI writes to this database continuously in WAL mode. Raw-copying the + main file and its -wal sibling cannot be made safe: a checkpoint between + the two copies pairs an old database with a newer WAL generation. SQLite's + online backup API takes the read lock it needs to produce a consistent + copy instead, without ever blocking the writer. The caller must remove the + returned temp dir once the connection is closed. + """ + if not path.exists(): + return None, None + try: + if path.stat().st_size > MAX_DB_BYTES: + return None, None + except OSError: + return None, None + + tmpdir = None + source = None + try: + tmpdir = Path(tempfile.mkdtemp(prefix="omarchy-copilot-")) + source = sqlite3.connect(f"file:{path}?mode=ro", uri=True, timeout=5.0) + target = sqlite3.connect(tmpdir / "snapshot.db") + with target: + source.backup(target) + target.row_factory = sqlite3.Row + return target, tmpdir + except (OSError, sqlite3.Error): + if tmpdir is not None: + shutil.rmtree(tmpdir, ignore_errors=True) + return None, None + finally: + if source is not None: + try: + source.close() + except sqlite3.Error: + pass def has_table(conn: sqlite3.Connection, name: str) -> bool: - try: - row = conn.execute( - "SELECT 1 FROM sqlite_master WHERE type='table' AND name=? LIMIT 1", (name,) - ).fetchone() - return row is not None - except sqlite3.Error: - return False + try: + row = conn.execute( + "SELECT 1 FROM sqlite_master WHERE type='table' AND name=? LIMIT 1", (name,) + ).fetchone() + return row is not None + except sqlite3.Error: + return False def blank_stats() -> dict[str, Any]: - return { - "todayPrompts": 0, - "todaySessions": 0, - "todayTotalTokens": 0, - "todayTokensByModel": {}, - "recentDays": [ - {"date": day, "messageCount": 0} for day in recent_date_strings() - ], - "modelUsage": {}, - "totalPrompts": 0, - "totalSessions": 0, - "activeDays": 0, - "activeDates": [], - "monthCredits": 0.0, - } + return { + "todayPrompts": 0, + "todaySessions": 0, + "todayTotalTokens": 0, + "todayTokensByModel": {}, + "recentDays": [ + {"date": day, "messageCount": 0} for day in recent_date_strings() + ], + "modelUsage": {}, + "totalPrompts": 0, + "totalSessions": 0, + "activeDays": 0, + "activeDates": [], + "monthCredits": 0.0, + } # --------------------------------------------------------------- aggregate class UsageTally: - """Running aggregate shared by both scans. + """Running aggregate shared by both scans. - It grows with days, sessions, and models, never with raw events, so a scan - can stream an entire store or transcript directory through it. - """ + It grows with days, sessions, and models, never with raw events, so a scan + can stream an entire store or transcript directory through it. + """ - def __init__(self) -> None: - self.today = local_date_string() - self.recent = { - day: {"date": day, "messageCount": 0} for day in recent_date_strings() - } - self.current_utc_month = dt.datetime.now(dt.timezone.utc).strftime("%Y-%m") - self.sessions: set[str] = set() - self.active_dates: set[str] = set() - self.today_sessions: set[str] = set() - self.today_prompts = 0 - self.today_total_tokens = 0 - self.today_tokens_by_model: dict[str, int] = {} - self.total_prompts = 0 - self.model_usage: dict[str, dict[str, int]] = {} - self.month_credits = 0.0 - - def prompt(self, session_key: str, day: str) -> None: - self.total_prompts += 1 - self.sessions.add(session_key) - self.active_dates.add(day) - if day == self.today: - self.today_prompts += 1 - self.today_sessions.add(session_key) - - def usage(self, model: str, day: str, tokens: dict[str, int]) -> None: - total = sum(tokens[key] for key in TOKEN_KEYS) - if total <= 0: - return - bucket = self.model_usage.setdefault(model, empty_bucket()) - for key in TOKEN_KEYS: - bucket[key] += tokens[key] - if day in self.recent: - self.recent[day]["messageCount"] += total - if day == self.today: - self.today_total_tokens += total - self.today_tokens_by_model[model] = ( - self.today_tokens_by_model.get(model, 0) + total - ) - - def spend(self, month_key: str, credits: float) -> None: - # The allowance window is the current UTC month, matching the reset - # boundary; older spend belongs to windows that are already over. - if month_key == self.current_utc_month: - self.month_credits += credits - - def stats(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": [self.recent[day] for day in recent_date_strings()], - "modelUsage": self.model_usage, - "totalPrompts": self.total_prompts, - "totalSessions": len(self.sessions), - "activeDays": len(self.active_dates), - "activeDates": sorted(self.active_dates), - "monthCredits": round(self.month_credits, 6), - } + def __init__(self) -> None: + self.today = local_date_string() + self.recent = { + day: {"date": day, "messageCount": 0} for day in recent_date_strings() + } + self.current_utc_month = dt.datetime.now(dt.timezone.utc).strftime("%Y-%m") + self.sessions: set[str] = set() + self.active_dates: set[str] = set() + self.today_sessions: set[str] = set() + self.today_prompts = 0 + self.today_total_tokens = 0 + self.today_tokens_by_model: dict[str, int] = {} + self.total_prompts = 0 + self.model_usage: dict[str, dict[str, int]] = {} + self.month_credits = 0.0 + + def prompt(self, session_key: str, day: str) -> None: + self.total_prompts += 1 + self.sessions.add(session_key) + self.active_dates.add(day) + if day == self.today: + self.today_prompts += 1 + self.today_sessions.add(session_key) + + def usage(self, model: str, day: str, tokens: dict[str, int]) -> None: + total = sum(tokens[key] for key in TOKEN_KEYS) + if total <= 0: + return + bucket = self.model_usage.setdefault(model, empty_bucket()) + for key in TOKEN_KEYS: + bucket[key] += tokens[key] + if day in self.recent: + self.recent[day]["messageCount"] += total + if day == self.today: + self.today_total_tokens += total + self.today_tokens_by_model[model] = ( + self.today_tokens_by_model.get(model, 0) + total + ) + + def spend(self, month_key: str, credits: float) -> None: + # The allowance window is the current UTC month, matching the reset + # boundary; older spend belongs to windows that are already over. + if month_key == self.current_utc_month: + self.month_credits += credits + + def stats(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": [self.recent[day] for day in recent_date_strings()], + "modelUsage": self.model_usage, + "totalPrompts": self.total_prompts, + "totalSessions": len(self.sessions), + "activeDays": len(self.active_dates), + "activeDates": sorted(self.active_dates), + "monthCredits": round(self.month_credits, 6), + } # ---------------------------------------------------------- store scan @@ -415,34 +415,34 @@ class UsageTally: def scan_store(conn: sqlite3.Connection) -> dict[str, Any]: - tally = UsageTally() - try: - rows = conn.execute( - "SELECT session_id, model, created_at, total_nano_aiu, input_tokens," - " output_tokens, cache_read_tokens, cache_write_tokens" - " FROM assistant_usage_events" - ) - for row in rows: - session_id = clean_name(row["session_id"]) - day = local_day(row["created_at"]) - tally.prompt(session_id or "unknown", day) - tally.spend( - utc_month_key(row["created_at"]), - credits_from_nano(row["total_nano_aiu"]), - ) - tally.usage( - clean_name(row["model"]) or "copilot", - day, - token_split( - row["input_tokens"], - row["output_tokens"], - row["cache_read_tokens"], - row["cache_write_tokens"], - ), - ) - except sqlite3.Error: - return blank_stats() - return tally.stats() + tally = UsageTally() + try: + rows = conn.execute( + "SELECT session_id, model, created_at, total_nano_aiu, input_tokens," + " output_tokens, cache_read_tokens, cache_write_tokens" + " FROM assistant_usage_events" + ) + for row in rows: + session_id = clean_name(row["session_id"]) + day = local_day(row["created_at"]) + tally.prompt(session_id or "unknown", day) + tally.spend( + utc_month_key(row["created_at"]), + credits_from_nano(row["total_nano_aiu"]), + ) + tally.usage( + clean_name(row["model"]) or "copilot", + day, + token_split( + row["input_tokens"], + row["output_tokens"], + row["cache_read_tokens"], + row["cache_write_tokens"], + ), + ) + except sqlite3.Error: + return blank_stats() + return tally.stats() # ------------------------------------------------------ transcript fallback @@ -461,536 +461,536 @@ def scan_store(conn: sqlite3.Connection) -> dict[str, Any]: def is_prompt(event_type: Any, data: dict[str, Any]) -> bool: - return event_type == "user.message" and clean_name(data.get("source")) in ( - "", - "user", - ) + return event_type == "user.message" and clean_name(data.get("source")) in ( + "", + "user", + ) def shutdown_metrics(data: dict[str, Any]) -> dict[str, dict[str, float]]: - metrics = data.get("modelMetrics") - if not isinstance(metrics, dict): - return {} - totals: dict[str, dict[str, float]] = {} - for model, entry in metrics.items(): - if not isinstance(entry, dict): - continue - usage = entry.get("usage") - if not isinstance(usage, dict): - usage = {} - totals[clean_name(model) or "copilot"] = { - **token_split( - usage.get("inputTokens"), - usage.get("outputTokens"), - usage.get("cacheReadTokens"), - usage.get("cacheWriteTokens"), - ), - "credits": credits_from_nano(entry.get("totalNanoAiu")), - } - return totals + metrics = data.get("modelMetrics") + if not isinstance(metrics, dict): + return {} + totals: dict[str, dict[str, float]] = {} + for model, entry in metrics.items(): + if not isinstance(entry, dict): + continue + usage = entry.get("usage") + if not isinstance(usage, dict): + usage = {} + totals[clean_name(model) or "copilot"] = { + **token_split( + usage.get("inputTokens"), + usage.get("outputTokens"), + usage.get("cacheReadTokens"), + usage.get("cacheWriteTokens"), + ), + "credits": credits_from_nano(entry.get("totalNanoAiu")), + } + return totals def apportion(total: int, weights: list[int]) -> list[int]: - """Split a whole total in proportion to positive weights, summing exactly.""" - whole = sum(weights) - shares: list[int] = [] - given = 0 - running = 0 - for weight in weights: - running += weight - share = total * running // whole - given - shares.append(share) - given += share - return shares + """Split a whole total in proportion to positive weights, summing exactly.""" + whole = sum(weights) + shares: list[int] = [] + given = 0 + running = 0 + for weight in weights: + running += weight + share = total * running // whole - given + shares.append(share) + given += share + return shares Reply = tuple[str, str, str, int] # model, local day, UTC month, output tokens def spread_segment( - tally: UsageTally, - model: str, - delta: dict[str, int], - credits: float, - replies: list[Reply], - shutdown_slot: tuple[str, str], + tally: UsageTally, + model: str, + delta: dict[str, int], + credits: float, + replies: list[Reply], + shutdown_slot: tuple[str, str], ) -> None: - """Book one model's usage for a finished segment on the days it happened. - - A segment can span midnight, and its shutdown only reports totals. Output - splits by each day's reply output, which the transcript records exactly. - Input, cache, and cost split by reply count: every reply is one model call - that resends the context, so calls, not output, track what was read. - """ - slots: dict[tuple[str, str], list[int]] = {} - for reply_model, day, month_key, output in replies: - if reply_model == model: - slot = slots.setdefault((day, month_key), [0, 0]) - slot[0] += 1 - slot[1] += output - if not slots: - slots[shutdown_slot] = [1, 0] - - calls = [counts[0] for counts in slots.values()] - outputs = [counts[1] for counts in slots.values()] - if sum(outputs) <= 0: - outputs = calls - split = { - key: apportion(delta[key], outputs if key == "outputTokens" else calls) - for key in TOKEN_KEYS - } - all_calls = sum(calls) - for index, (day, month_key) in enumerate(slots): - tally.usage(model, day, {key: split[key][index] for key in TOKEN_KEYS}) - tally.spend(month_key, credits * calls[index] / all_calls) + """Book one model's usage for a finished segment on the days it happened. + + A segment can span midnight, and its shutdown only reports totals. Output + splits by each day's reply output, which the transcript records exactly. + Input, cache, and cost split by reply count: every reply is one model call + that resends the context, so calls, not output, track what was read. + """ + slots: dict[tuple[str, str], list[int]] = {} + for reply_model, day, month_key, output in replies: + if reply_model == model: + slot = slots.setdefault((day, month_key), [0, 0]) + slot[0] += 1 + slot[1] += output + if not slots: + slots[shutdown_slot] = [1, 0] + + calls = [counts[0] for counts in slots.values()] + outputs = [counts[1] for counts in slots.values()] + if sum(outputs) <= 0: + outputs = calls + split = { + key: apportion(delta[key], outputs if key == "outputTokens" else calls) + for key in TOKEN_KEYS + } + all_calls = sum(calls) + for index, (day, month_key) in enumerate(slots): + tally.usage(model, day, {key: split[key][index] for key in TOKEN_KEYS}) + tally.spend(month_key, credits * calls[index] / all_calls) def scan_transcript(event_file: Path, session_key: str, tally: UsageTally) -> None: - try: - session_mtime = event_file.stat().st_mtime - except OSError: - session_mtime = 0.0 - - # Session totals as of the last shutdown, and the replies since then. - reported: dict[str, dict[str, float]] = {} - replies: list[Reply] = [] - - with event_file.open("r", encoding="utf-8", errors="replace") as handle: - for raw in handle: - try: - entry = json.loads(raw) - except ValueError: - continue - if not isinstance(entry, dict): - continue - event_type = entry.get("type") - data = entry.get("data") - if not isinstance(data, dict): - continue - - parsed = parse_event_time(entry.get("timestamp")) - if parsed is None and session_mtime > 0: - parsed = dt.datetime.fromtimestamp( - session_mtime, dt.timezone.utc - ).astimezone() - day = parsed.date().isoformat() if parsed else tally.today - month_key = ( - parsed.astimezone(dt.timezone.utc).strftime("%Y-%m") if parsed else "" - ) - - if is_prompt(event_type, data): - tally.prompt(session_key, day) - elif event_type == "assistant.message": - model = clean_name(data.get("model")) or "copilot" - replies.append( - (model, day, month_key, number(data.get("outputTokens"))) - ) - elif event_type == "session.shutdown": - for model, totals in shutdown_metrics(data).items(): - before = reported.get(model, {}) - spread_segment( - tally, - model, - { - key: max(0, int(totals[key] - before.get(key, 0))) - for key in TOKEN_KEYS - }, - max(0.0, totals["credits"] - before.get("credits", 0.0)), - replies, - (day, month_key), - ) - reported[model] = totals - # The shutdown totals already count these replies. - replies.clear() - - for model, day, _month_key, output in replies: - tally.usage(model, day, {**empty_bucket(), "outputTokens": output}) + try: + session_mtime = event_file.stat().st_mtime + except OSError: + session_mtime = 0.0 + + # Session totals as of the last shutdown, and the replies since then. + reported: dict[str, dict[str, float]] = {} + replies: list[Reply] = [] + + with event_file.open("r", encoding="utf-8", errors="replace") as handle: + for raw in handle: + try: + entry = json.loads(raw) + except ValueError: + continue + if not isinstance(entry, dict): + continue + event_type = entry.get("type") + data = entry.get("data") + if not isinstance(data, dict): + continue + + parsed = parse_event_time(entry.get("timestamp")) + if parsed is None and session_mtime > 0: + parsed = dt.datetime.fromtimestamp( + session_mtime, dt.timezone.utc + ).astimezone() + day = parsed.date().isoformat() if parsed else tally.today + month_key = ( + parsed.astimezone(dt.timezone.utc).strftime("%Y-%m") if parsed else "" + ) + + if is_prompt(event_type, data): + tally.prompt(session_key, day) + elif event_type == "assistant.message": + model = clean_name(data.get("model")) or "copilot" + replies.append( + (model, day, month_key, number(data.get("outputTokens"))) + ) + elif event_type == "session.shutdown": + for model, totals in shutdown_metrics(data).items(): + before = reported.get(model, {}) + spread_segment( + tally, + model, + { + key: max(0, int(totals[key] - before.get(key, 0))) + for key in TOKEN_KEYS + }, + max(0.0, totals["credits"] - before.get("credits", 0.0)), + replies, + (day, month_key), + ) + reported[model] = totals + # The shutdown totals already count these replies. + replies.clear() + + for model, day, _month_key, output in replies: + tally.usage(model, day, {**empty_bucket(), "outputTokens": output}) def scan_transcripts(session_root: Path) -> dict[str, Any]: - if not session_root.is_dir(): - return blank_stats() - tally = UsageTally() - for session_dir in sorted(p for p in session_root.iterdir() if p.is_dir()): - event_file = session_dir / "events.jsonl" - if not event_file.is_file(): - continue - try: - scan_transcript(event_file, session_dir.name, tally) - except OSError: - continue - return tally.stats() + if not session_root.is_dir(): + return blank_stats() + tally = UsageTally() + for session_dir in sorted(p for p in session_root.iterdir() if p.is_dir()): + event_file = session_dir / "events.jsonl" + if not event_file.is_file(): + continue + try: + scan_transcript(event_file, session_dir.name, tally) + except OSError: + continue + return tally.stats() # ------------------------------------------------------------------- cache def scan_cache_paths(home_dir: Path) -> tuple[Path, Path]: - digest = hashlib.sha1(str(home_dir).encode("utf-8")).hexdigest()[:16] - root = cache_root() - return root / f"copilot-scan-{digest}.json", root / f"copilot-scan-{digest}.lock" + digest = hashlib.sha1(str(home_dir).encode("utf-8")).hexdigest()[:16] + root = cache_root() + return root / f"copilot-scan-{digest}.json", root / f"copilot-scan-{digest}.lock" def read_fresh_json(path: Path, max_age_seconds: float) -> dict[str, Any] | None: - if max_age_seconds <= 0 or not path.exists(): - return None - try: - 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 + if max_age_seconds <= 0 or not path.exists(): + return None + try: + 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: 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=(",", ":")) + "\n") - tmp.chmod(0o644) - tmp.replace(path) - except BaseException: - tmp.unlink(missing_ok=True) - raise + # 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=(",", ":")) + "\n") + tmp.chmod(0o644) + tmp.replace(path) + except BaseException: + tmp.unlink(missing_ok=True) + raise def read_cached_stats( - cache_file: Path, max_age_seconds: float, today: str + cache_file: Path, max_age_seconds: float, today: str ) -> 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 - 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 + cached = read_fresh_json(cache_file, max_age_seconds) + if not isinstance(cached, dict) or cached.get("schemaVersion") != 1: + return None + 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": local_date_string(), "stats": stats}, - ) - except Exception as exc: - print( - f"omarchy-agent-usage-copilot: could not write usage cache ({exc})", - file=sys.stderr, - ) + try: + write_json( + cache_file, + {"schemaVersion": 1, "scanDate": local_date_string(), "stats": stats}, + ) + except Exception as exc: + print( + f"omarchy-agent-usage-copilot: could not write usage cache ({exc})", + file=sys.stderr, + ) def cached_scan(home_dir: Path, max_age_seconds: float) -> dict[str, Any]: - try: - cache_file, lock_file = scan_cache_paths(home_dir) - today = local_date_string() - cached = read_cached_stats(cache_file, max_age_seconds, today) - if cached is not None: - return cached - with lock_file.open("w") as lock: - fcntl.flock(lock, fcntl.LOCK_EX) - cached = read_cached_stats(cache_file, max_age_seconds, today) - if cached is not None: - return cached - stats = direct_scan(home_dir) - write_cached_stats(cache_file, stats) - return stats - except Exception as exc: - print( - f"omarchy-agent-usage-copilot: cache unavailable ({exc}); scanning directly", - file=sys.stderr, - ) - return direct_scan(home_dir) + try: + cache_file, lock_file = scan_cache_paths(home_dir) + today = local_date_string() + cached = read_cached_stats(cache_file, max_age_seconds, today) + if cached is not None: + return cached + with lock_file.open("w") as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + cached = read_cached_stats(cache_file, max_age_seconds, today) + if cached is not None: + return cached + stats = direct_scan(home_dir) + write_cached_stats(cache_file, stats) + return stats + except Exception as exc: + print( + f"omarchy-agent-usage-copilot: cache unavailable ({exc}); scanning directly", + file=sys.stderr, + ) + return direct_scan(home_dir) def direct_scan(home_dir: Path) -> dict[str, Any]: - conn, tmpdir = open_store(home_dir / "session-store.db") - if conn is None: - return scan_transcripts(home_dir / "session-state") - try: - if has_table(conn, "assistant_usage_events"): - stats = scan_store(conn) - if number(stats.get("totalPrompts")) > 0: - return stats - # A store without the usage table, or one that is still empty, falls back - # to the transcripts exactly like a store-less machine. - return scan_transcripts(home_dir / "session-state") - finally: - conn.close() - if tmpdir is not None: - shutil.rmtree(tmpdir, ignore_errors=True) + conn, tmpdir = open_store(home_dir / "session-store.db") + if conn is None: + return scan_transcripts(home_dir / "session-state") + try: + if has_table(conn, "assistant_usage_events"): + stats = scan_store(conn) + if number(stats.get("totalPrompts")) > 0: + return stats + # A store without the usage table, or one that is still empty, falls back + # to the transcripts exactly like a store-less machine. + return scan_transcripts(home_dir / "session-state") + finally: + conn.close() + if tmpdir is not None: + shutil.rmtree(tmpdir, ignore_errors=True) # ------------------------------------------------------------ remote quota def github_token() -> str: - """Find a GitHub token without ever storing or logging one. - - Environment first, so a user can scope exactly what this sees, then the - `gh` CLI's own credential store, which is where a machine that has signed - in already keeps it. - """ - for var in ("COPILOT_QUOTA_TOKEN", "GH_TOKEN", "GITHUB_TOKEN"): - value = os.environ.get(var) - if value and value.strip(): - return value.strip() - exe = shutil.which("gh") - if not exe: - return "" - try: - result = subprocess.run( - [exe, "auth", "token"], - capture_output=True, - text=True, - timeout=REMOTE_TIMEOUT, - check=False, - ) - except (OSError, subprocess.SubprocessError): - return "" - return result.stdout.strip() if result.returncode == 0 else "" + """Find a GitHub token without ever storing or logging one. + + Environment first, so a user can scope exactly what this sees, then the + `gh` CLI's own credential store, which is where a machine that has signed + in already keeps it. + """ + for var in ("COPILOT_QUOTA_TOKEN", "GH_TOKEN", "GITHUB_TOKEN"): + value = os.environ.get(var) + if value and value.strip(): + return value.strip() + exe = shutil.which("gh") + if not exe: + return "" + try: + result = subprocess.run( + [exe, "auth", "token"], + capture_output=True, + text=True, + timeout=REMOTE_TIMEOUT, + check=False, + ) + except (OSError, subprocess.SubprocessError): + return "" + return result.stdout.strip() if result.returncode == 0 else "" def api_plan(raw: Any) -> tuple[str, str]: - key = API_PLAN_KEYS.get(str(raw or "").strip().lower(), "") - if not key: - return "", "" - return key, PLAN_LABELS.get(key, "Copilot") + key = API_PLAN_KEYS.get(str(raw or "").strip().lower(), "") + if not key: + return "", "" + return key, PLAN_LABELS.get(key, "Copilot") def fetch_remote_quota() -> dict[str, Any]: - """Ask GitHub for the account-wide credit balance. - - Every failure is soft: a missing token, an offline machine, or a changed - payload leaves the record on its local estimate rather than failing the - whole collection. - """ - token = github_token() - if not token: - return {"available": False, "reason": "no-token"} - - request = urllib.request.Request( - QUOTA_URL, - headers={ - "Authorization": f"Bearer {token}", - "Accept": "application/json", - "User-Agent": "omarchy-agent-usage", - }, - ) - try: - with urllib.request.urlopen(request, timeout=REMOTE_TIMEOUT) as response: - payload = json.loads( - response.read(2_000_000).decode("utf-8", errors="replace") - ) - except urllib.error.HTTPError as exc: - return {"available": False, "reason": f"http-{exc.code}"} - except (urllib.error.URLError, OSError, ValueError, TimeoutError): - return {"available": False, "reason": "unreachable"} - - if not isinstance(payload, dict): - return {"available": False, "reason": "bad-payload"} - - snapshots = payload.get("quota_snapshots") - snapshot = ( - snapshots.get("premium_interactions") if isinstance(snapshots, dict) else None - ) - if not isinstance(snapshot, dict): - return {"available": False, "reason": "no-quota"} - - plan_key, plan_label = api_plan(payload.get("copilot_plan")) - orgs = payload.get("organization_login_list") - # Under token-based billing the snapshot counts AI credits directly; only - # the local store and transcripts count in nano units. - return { - "available": True, - "unlimited": bool(snapshot.get("unlimited")), - "entitlement": positive_float(snapshot.get("entitlement")), - "used": positive_float(snapshot.get("credits_used")), - "resetsAt": str(payload.get("quota_reset_date_utc") or ""), - "plan": plan_key, - "planLabel": plan_label, - "org": str(orgs[0]) if isinstance(orgs, list) and orgs else "", - "login": clean_name(payload.get("login")), - } + """Ask GitHub for the account-wide credit balance. + + Every failure is soft: a missing token, an offline machine, or a changed + payload leaves the record on its local estimate rather than failing the + whole collection. + """ + token = github_token() + if not token: + return {"available": False, "reason": "no-token"} + + request = urllib.request.Request( + QUOTA_URL, + headers={ + "Authorization": f"Bearer {token}", + "Accept": "application/json", + "User-Agent": "omarchy-agent-usage", + }, + ) + try: + with urllib.request.urlopen(request, timeout=REMOTE_TIMEOUT) as response: + payload = json.loads( + response.read(2_000_000).decode("utf-8", errors="replace") + ) + except urllib.error.HTTPError as exc: + return {"available": False, "reason": f"http-{exc.code}"} + except (urllib.error.URLError, OSError, ValueError, TimeoutError): + return {"available": False, "reason": "unreachable"} + + if not isinstance(payload, dict): + return {"available": False, "reason": "bad-payload"} + + snapshots = payload.get("quota_snapshots") + snapshot = ( + snapshots.get("premium_interactions") if isinstance(snapshots, dict) else None + ) + if not isinstance(snapshot, dict): + return {"available": False, "reason": "no-quota"} + + plan_key, plan_label = api_plan(payload.get("copilot_plan")) + orgs = payload.get("organization_login_list") + # Under token-based billing the snapshot counts AI credits directly; only + # the local store and transcripts count in nano units. + return { + "available": True, + "unlimited": bool(snapshot.get("unlimited")), + "entitlement": positive_float(snapshot.get("entitlement")), + "used": positive_float(snapshot.get("credits_used")), + "resetsAt": str(payload.get("quota_reset_date_utc") or ""), + "plan": plan_key, + "planLabel": plan_label, + "org": str(orgs[0]) if isinstance(orgs, list) and orgs else "", + "login": clean_name(payload.get("login")), + } # ---------------------------------------------------------------- allowance def resolve_plan(config: dict[str, Any]) -> tuple[str, float]: - """Return the configured plan key and its monthly allowance in credits. + """Return the configured plan key and its monthly allowance in credits. - The plan is not discoverable locally, so a missing or unrecognized plan - stays unknown (key "") instead of defaulting to a tier the user may not - have; only monthlyCredits can then give the estimate an allowance. - """ - raw = str(config.get("plan") or "").strip().lower() - key = raw.replace(" ", "").replace("copilot", "") - if key in ("proplus", "pro-plus"): - key = "pro+" - if key not in PLAN_LABELS: - key = "" + The plan is not discoverable locally, so a missing or unrecognized plan + stays unknown (key "") instead of defaulting to a tier the user may not + have; only monthlyCredits can then give the estimate an allowance. + """ + raw = str(config.get("plan") or "").strip().lower() + key = raw.replace(" ", "").replace("copilot", "") + if key in ("proplus", "pro-plus"): + key = "pro+" + if key not in PLAN_LABELS: + key = "" - override = config.get("monthlyCredits") - if isinstance(override, (int, float)) and override > 0: - return key, float(override) - return key, float(PLAN_ALLOWANCES.get(key, 0)) + override = config.get("monthlyCredits") + if isinstance(override, (int, float)) and override > 0: + return key, float(override) + return key, float(PLAN_ALLOWANCES.get(key, 0)) def allowance_limits( - month_credits: float, - config: dict[str, Any], - remote: dict[str, Any] | None, - now: dt.datetime, - has_stats: bool, + month_credits: float, + config: dict[str, Any], + remote: dict[str, Any] | None, + now: dt.datetime, + has_stats: bool, ) -> dict[str, Any]: - plan_key, plan_allowance = resolve_plan(config) - plan_label = PLAN_LABELS.get(plan_key, "") - estimated = True - allow = plan_allowance - used = month_credits - resets_at = next_month_start(now).isoformat() - reason = "" - - live = remote if isinstance(remote, dict) and remote.get("available") else None - if live and not live.get("unlimited") and float(live.get("entitlement") or 0) > 0: - allow = float(live["entitlement"]) - used = float(live.get("used") or 0.0) - estimated = False - if live.get("plan"): - plan_key = str(live["plan"]) - plan_label = str(live["planLabel"] or PLAN_LABELS.get(plan_key, "Copilot")) - raw_reset = str(live.get("resetsAt") or "") - if raw_reset: - parsed = parse_event_time(raw_reset) - if parsed is not None and parsed > now: - resets_at = parsed.isoformat() - elif ( - remote is not None - and remote.get("available") is False - and remote.get("reason") != "no-token" - ): - reason = "Copilot limits unavailable" - - limits = [] - if allow > 0: - # A plan-default allowance must not advertise a meter on a machine the - # CLI has never touched — the local estimate needs local numbers to - # mean anything. The live account figure stands on its own. - if estimated and not has_stats: - pass - else: - percent = min(1.0, used / allow) - limits.append( - { - "label": "Monthly allowance (est.)" - if estimated - else "Monthly allowance", - "title": "Monthly allowance (est.)" - if estimated - else "Monthly allowance", - "percent": percent, - "resetsAt": resets_at, - } - ) + plan_key, plan_allowance = resolve_plan(config) + plan_label = PLAN_LABELS.get(plan_key, "") + estimated = True + allow = plan_allowance + used = month_credits + resets_at = next_month_start(now).isoformat() + reason = "" + + live = remote if isinstance(remote, dict) and remote.get("available") else None + if live and not live.get("unlimited") and float(live.get("entitlement") or 0) > 0: + allow = float(live["entitlement"]) + used = float(live.get("used") or 0.0) + estimated = False + if live.get("plan"): + plan_key = str(live["plan"]) + plan_label = str(live["planLabel"] or PLAN_LABELS.get(plan_key, "Copilot")) + raw_reset = str(live.get("resetsAt") or "") + if raw_reset: + parsed = parse_event_time(raw_reset) + if parsed is not None and parsed > now: + resets_at = parsed.isoformat() + elif ( + remote is not None + and remote.get("available") is False + and remote.get("reason") != "no-token" + ): + reason = "Copilot limits unavailable" + + limits = [] + if allow > 0: + # A plan-default allowance must not advertise a meter on a machine the + # CLI has never touched — the local estimate needs local numbers to + # mean anything. The live account figure stands on its own. + if estimated and not has_stats: + pass + else: + percent = min(1.0, used / allow) + limits.append( + { + "label": "Monthly allowance (est.)" + if estimated + else "Monthly allowance", + "title": "Monthly allowance (est.)" + if estimated + else "Monthly allowance", + "percent": percent, + "resetsAt": resets_at, + } + ) - return { - "limits": limits, - "tierLabel": plan_label, - "usageStatusText": reason, - "authHelpText": AUTH_HELP, - } + return { + "limits": limits, + "tierLabel": plan_label, + "usageStatusText": reason, + "authHelpText": AUTH_HELP, + } # ------------------------------------------------------------------ record def main() -> int: - parser = argparse.ArgumentParser() - # --force rescans everything and rewrites the cache. --limits-only is kept - # for CLI compatibility with the panel's refreshLimits() call: only the - # limits probe must be fresh, so it may reuse a scan for far longer than a - # normal run, whose short window exists purely to dedup concurrent - # collector runs. - parser.add_argument("--force", action="store_true") - parser.add_argument("--limits-only", action="store_true") - parser.add_argument( - "--copilot-home", default=os.environ.get("COPILOT_HOME") or "~/.copilot" - ) - args = parser.parse_args() - - home_dir = expand_path(args.copilot_home) - max_age = ( - 0 - if args.force - else (LIMITS_ONLY_REUSE_SECONDS if args.limits_only else SCAN_REUSE_SECONDS) - ) - stats = cached_scan(home_dir, max_age) - - config = read_config() - now = dt.datetime.now(dt.timezone.utc) - - remote = {"available": False, "reason": "disabled"} - if config.get("remote", True) is not False: - remote = fetch_remote_quota() - - has_stats = ( - number(stats.get("totalPrompts")) > 0 - or float(stats.get("monthCredits") or 0.0) > 0 - ) - allowance = allowance_limits( - float(stats.get("monthCredits") or 0.0), config, remote, now, has_stats - ) - - record: dict[str, Any] = { - "schemaVersion": 1, - "id": AGENT_ID, - "name": AGENT_NAME, - "updatedAt": now.isoformat(), - "ready": number(stats.get("totalPrompts")) > 0 or len(allowance["limits"]) > 0, - "hasLocalStats": True, - "tierLabel": allowance["tierLabel"], - "usageStatusText": allowance["usageStatusText"], - "authHelpText": allowance["authHelpText"], - "limits": allowance["limits"], - } - for key in ( - "todayPrompts", - "todaySessions", - "todayTotalTokens", - "totalPrompts", - "totalSessions", - "activeDays", - ): - record[key] = number(stats.get(key)) - record["todayTokensByModel"] = stats.get("todayTokensByModel") or {} - record["recentDays"] = stats.get("recentDays") or [] - record["modelUsage"] = stats.get("modelUsage") or {} - record["activeDates"] = stats.get("activeDates") or [] - print(json.dumps(record, separators=(",", ":"))) - return 0 + parser = argparse.ArgumentParser() + # --force rescans everything and rewrites the cache. --limits-only is kept + # for CLI compatibility with the panel's refreshLimits() call: only the + # limits probe must be fresh, so it may reuse a scan for far longer than a + # normal run, whose short window exists purely to dedup concurrent + # collector runs. + parser.add_argument("--force", action="store_true") + parser.add_argument("--limits-only", action="store_true") + parser.add_argument( + "--copilot-home", default=os.environ.get("COPILOT_HOME") or "~/.copilot" + ) + args = parser.parse_args() + + home_dir = expand_path(args.copilot_home) + max_age = ( + 0 + if args.force + else (LIMITS_ONLY_REUSE_SECONDS if args.limits_only else SCAN_REUSE_SECONDS) + ) + stats = cached_scan(home_dir, max_age) + + config = read_config() + now = dt.datetime.now(dt.timezone.utc) + + remote = {"available": False, "reason": "disabled"} + if config.get("remote", True) is not False: + remote = fetch_remote_quota() + + has_stats = ( + number(stats.get("totalPrompts")) > 0 + or float(stats.get("monthCredits") or 0.0) > 0 + ) + allowance = allowance_limits( + float(stats.get("monthCredits") or 0.0), config, remote, now, has_stats + ) + + record: dict[str, Any] = { + "schemaVersion": 1, + "id": AGENT_ID, + "name": AGENT_NAME, + "updatedAt": now.isoformat(), + "ready": number(stats.get("totalPrompts")) > 0 or len(allowance["limits"]) > 0, + "hasLocalStats": True, + "tierLabel": allowance["tierLabel"], + "usageStatusText": allowance["usageStatusText"], + "authHelpText": allowance["authHelpText"], + "limits": allowance["limits"], + } + for key in ( + "todayPrompts", + "todaySessions", + "todayTotalTokens", + "totalPrompts", + "totalSessions", + "activeDays", + ): + record[key] = number(stats.get(key)) + record["todayTokensByModel"] = stats.get("todayTokensByModel") or {} + record["recentDays"] = stats.get("recentDays") or [] + record["modelUsage"] = stats.get("modelUsage") or {} + record["activeDates"] = stats.get("activeDates") or [] + print(json.dumps(record, separators=(",", ":"))) + return 0 if __name__ == "__main__": - raise SystemExit(main()) + raise SystemExit(main()) From 8e45c3ce8741d16c43dd7e613f144e08e66c0278 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Fri, 2 Oct 2026 06:14:32 -0400 Subject: [PATCH 019/546] 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 020/546] 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 021/546] 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 022/546] 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 51188a64fd2f7bbc066f6a175497d1e1d60613c1 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Fri, 2 Oct 2026 06:16:40 -0400 Subject: [PATCH 023/546] Take the Copilot plan from GitHub and keep its last allowance The live copilot_internal/user answer already names the plan, the allowance, and the reset date, so drop the hardcoded plan-allowance table and the plan setting. A monthlyCredits budget stays as the one fallback, labeled "(est.)", for machines that can't ask GitHub. Cache each token's last good quota and, when a check fails, show it marked limitsStale with its limitsFetchedAt so the panel dims it. Only ask GitHub where the Copilot CLI has been used (or the config says to), so a machine with gh signed in but no Copilot never sends its token and stays an empty record. Mention Copilot in the manifest and the manual. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-copilot | 272 +++++++++-------------- manual/17-ai.md | 2 +- shell/plugins/agents/README.md | 52 ++--- shell/plugins/agents/manifest.json | 3 +- test/shell.d/agent-usage-copilot-test.sh | 69 ++++-- 5 files changed, 175 insertions(+), 223 deletions(-) diff --git a/bin/omarchy-agent-usage-copilot b/bin/omarchy-agent-usage-copilot index e060cb31a95..de8467c1e94 100755 --- a/bin/omarchy-agent-usage-copilot +++ b/bin/omarchy-agent-usage-copilot @@ -15,32 +15,21 @@ whose store predates that table, is the session transcript directory same per-model token split and cost. The record prefers whichever source produced numbers, never both. -The monthly meter is an allowance of AI credits. GitHub's own account-wide -figure is authoritative when the internal quota endpoint answers (the same one -the editor plugins use); otherwise the collector falls back to this machine's -CLI cost against a plan allowance from configuration, because the included -amount is not discoverable locally. All remote reads fail soft — an offline -machine or a missing token degrades to the local estimate instead of erroring. - -Configuration lives in ~/.config/omarchy/agents/copilot.json: +The monthly meter is an allowance of AI credits: GitHub's own account-wide +figure from the internal quota endpoint (the same one the editor plugins +use), asked with a token from COPILOT_QUOTA_TOKEN, GH_TOKEN, GITHUB_TOKEN, or +`gh auth token`. It names the plan, the allowance, and when it resets, and a +failed check keeps the last answer, marked stale. Only when it can't be asked +at all does a monthlyCredits budget from ~/.config/omarchy/agents/copilot.json +give a local estimate, labeled as one, of this machine's CLI spend: { - "plan": "pro", - "monthlyCredits": 1500, - "remote": true + "monthlyCredits": 1500, + "remote": true } -plan is one of free, student, pro, pro+, max, business, enterprise and picks -the default allowance (Pro 1500, Pro+ 7000, Max 20000, Business 1900, -Enterprise 3900 credits per month; existing seat-based customers may see -higher promotional amounts through September 2026). monthlyCredits overrides -the plan table, for AI-credit plans without a published number or for a -personal budget. The plan is not discoverable locally, so without plan or -monthlyCredits there is no local estimate and no plan label. remote (default -true) turns the account-wide quota probe on or off; the probe needs a token -from COPILOT_QUOTA_TOKEN, GH_TOKEN, GITHUB_TOKEN, or `gh auth token`. The -allowance window resets on the 1st of each month at 00:00 UTC, GitHub's own -boundary. +remote turns the quota probe off when false, and on even without a Copilot +CLI home when true; by default it runs only where the CLI has been used. The agents panel only ever reads the JSON this prints. """ @@ -75,38 +64,14 @@ ULTRA_SECONDS = 1_000_000_000 QUOTA_URL = "https://api.github.com/copilot_internal/user" REMOTE_TIMEOUT = 6.0 -# Plan identifiers as the quota endpoint spells them. -API_PLAN_KEYS = { - "individual": "pro", - "individual_pro": "pro+", - "business": "business", - "enterprise": "enterprise", - "free": "free", - "max": "max", -} - -# Monthly AI credit allowances from GitHub's usage-based billing tables. -# Seat-based plans pool credits at the billing entity, so an individual's -# practical ceiling is set by their admin rather than this figure — -# monthlyCredits in the config overrides it. -PLAN_ALLOWANCES = { - "free": 0, - "student": 0, - "pro": 1_500, - "pro+": 7_000, - "max": 20_000, - "business": 1_900, - "enterprise": 3_900, -} - -PLAN_LABELS = { - "free": "Free", - "student": "Student", - "pro": "Pro", - "pro+": "Pro+", - "max": "Max", +# Plan names as the quota endpoint spells them. +API_PLAN_LABELS = { + "individual": "Pro", + "individual_pro": "Pro+", "business": "Business", "enterprise": "Enterprise", + "free": "Free", + "max": "Max", } # The session store is mutable and grows without bound, so it is size-capped @@ -771,24 +736,8 @@ def github_token() -> str: return result.stdout.strip() if result.returncode == 0 else "" -def api_plan(raw: Any) -> tuple[str, str]: - key = API_PLAN_KEYS.get(str(raw or "").strip().lower(), "") - if not key: - return "", "" - return key, PLAN_LABELS.get(key, "Copilot") - - -def fetch_remote_quota() -> dict[str, Any]: - """Ask GitHub for the account-wide credit balance. - - Every failure is soft: a missing token, an offline machine, or a changed - payload leaves the record on its local estimate rather than failing the - whole collection. - """ - token = github_token() - if not token: - return {"available": False, "reason": "no-token"} - +def ask_quota(token: str) -> dict[str, Any]: + """Ask GitHub for the account-wide credit balance.""" request = urllib.request.Request( QUOTA_URL, headers={ @@ -799,9 +748,7 @@ def fetch_remote_quota() -> dict[str, Any]: ) try: with urllib.request.urlopen(request, timeout=REMOTE_TIMEOUT) as response: - payload = json.loads( - response.read(2_000_000).decode("utf-8", errors="replace") - ) + payload = json.loads(response.read(2_000_000).decode("utf-8", errors="replace")) except urllib.error.HTTPError as exc: return {"available": False, "reason": f"http-{exc.code}"} except (urllib.error.URLError, OSError, ValueError, TimeoutError): @@ -811,14 +758,10 @@ def fetch_remote_quota() -> dict[str, Any]: return {"available": False, "reason": "bad-payload"} snapshots = payload.get("quota_snapshots") - snapshot = ( - snapshots.get("premium_interactions") if isinstance(snapshots, dict) else None - ) + snapshot = snapshots.get("premium_interactions") if isinstance(snapshots, dict) else None if not isinstance(snapshot, dict): return {"available": False, "reason": "no-quota"} - plan_key, plan_label = api_plan(payload.get("copilot_plan")) - orgs = payload.get("organization_login_list") # Under token-based billing the snapshot counts AI credits directly; only # the local store and transcripts count in nano units. return { @@ -827,34 +770,41 @@ def fetch_remote_quota() -> dict[str, Any]: "entitlement": positive_float(snapshot.get("entitlement")), "used": positive_float(snapshot.get("credits_used")), "resetsAt": str(payload.get("quota_reset_date_utc") or ""), - "plan": plan_key, - "planLabel": plan_label, - "org": str(orgs[0]) if isinstance(orgs, list) and orgs else "", - "login": clean_name(payload.get("login")), + "planLabel": API_PLAN_LABELS.get(str(payload.get("copilot_plan") or "").strip().lower(), ""), } -# ---------------------------------------------------------------- allowance +def fetch_remote_quota() -> dict[str, Any]: + """The live quota, or the last good one marked stale. + + Every failure is soft: a missing token, an offline machine, or a changed + payload keeps the last answer this token got, dimmed, or leaves the record + on its local stats rather than failing the whole collection. The cache is + keyed by the token, so a machine signed in to another account never shows + the last one's allowance. + """ + token = github_token() + if not token: + return {"available": False, "reason": "no-token"} + cache = cache_root() / f"copilot-quota-{hashlib.sha256(token.encode('utf-8')).hexdigest()[:16]}.json" + quota = ask_quota(token) + if quota.get("available"): + quota.update(stale=False, fetchedAtMs=round(time.time() * 1000)) + try: + write_json(cache, quota) + except OSError: + pass + return quota -def resolve_plan(config: dict[str, Any]) -> tuple[str, float]: - """Return the configured plan key and its monthly allowance in credits. + cached = read_fresh_json(cache, float("inf")) or {} + if cached.get("available"): + cached.update(stale=True, failure=quota.get("reason", "")) + return cached + return quota - The plan is not discoverable locally, so a missing or unrecognized plan - stays unknown (key "") instead of defaulting to a tier the user may not - have; only monthlyCredits can then give the estimate an allowance. - """ - raw = str(config.get("plan") or "").strip().lower() - key = raw.replace(" ", "").replace("copilot", "") - if key in ("proplus", "pro-plus"): - key = "pro+" - if key not in PLAN_LABELS: - key = "" - override = config.get("monthlyCredits") - if isinstance(override, (int, float)) and override > 0: - return key, float(override) - return key, float(PLAN_ALLOWANCES.get(key, 0)) +# ---------------------------------------------------------------- allowance def allowance_limits( @@ -864,62 +814,53 @@ def allowance_limits( now: dt.datetime, has_stats: bool, ) -> dict[str, Any]: - plan_key, plan_allowance = resolve_plan(config) - plan_label = PLAN_LABELS.get(plan_key, "") - estimated = True - allow = plan_allowance - used = month_credits - resets_at = next_month_start(now).isoformat() - reason = "" + """The monthly allowance meter. - live = remote if isinstance(remote, dict) and remote.get("available") else None - if live and not live.get("unlimited") and float(live.get("entitlement") or 0) > 0: - allow = float(live["entitlement"]) - used = float(live.get("used") or 0.0) - estimated = False - if live.get("plan"): - plan_key = str(live["plan"]) - plan_label = str(live["planLabel"] or PLAN_LABELS.get(plan_key, "Copilot")) - raw_reset = str(live.get("resetsAt") or "") - if raw_reset: - parsed = parse_event_time(raw_reset) - if parsed is not None and parsed > now: - resets_at = parsed.isoformat() - elif ( - remote is not None - and remote.get("available") is False - and remote.get("reason") != "no-token" - ): - reason = "Copilot limits unavailable" - - limits = [] - if allow > 0: - # A plan-default allowance must not advertise a meter on a machine the - # CLI has never touched — the local estimate needs local numbers to - # mean anything. The live account figure stands on its own. - if estimated and not has_stats: - pass - else: - percent = min(1.0, used / allow) - limits.append( - { - "label": "Monthly allowance (est.)" - if estimated - else "Monthly allowance", - "title": "Monthly allowance (est.)" - if estimated - else "Monthly allowance", - "percent": percent, - "resetsAt": resets_at, - } - ) - - return { - "limits": limits, - "tierLabel": plan_label, - "usageStatusText": reason, + The live figure from the quota endpoint counts every machine, editor, and + agent on the entitlement, plan included. When it can't be asked, a + monthlyCredits budget from the configuration gives a local estimate of this + machine's CLI spend, labeled as such; without one there is no meter. + """ + result = { + "limits": [], + "tierLabel": "", + "usageStatusText": "", "authHelpText": AUTH_HELP, + "limitsStale": False, + "limitsFetchedAt": 0, } + resets_at = next_month_start(now).isoformat() + + live = remote if isinstance(remote, dict) and remote.get("available") else None + if live: + result["tierLabel"] = str(live.get("planLabel") or "") + if live.get("unlimited") or float(live.get("entitlement") or 0) <= 0: + return result + parsed = parse_event_time(str(live.get("resetsAt") or "")) + if parsed is not None and parsed > now: + resets_at = parsed.isoformat() + result["limits"] = [{ + "label": "Monthly allowance", + "percent": min(1.0, float(live.get("used") or 0.0) / float(live["entitlement"])), + "resetsAt": resets_at, + }] + result["limitsStale"] = bool(live.get("stale")) + result["limitsFetchedAt"] = number(live.get("fetchedAtMs")) + return result + + if remote is not None and remote.get("reason") not in ("no-token", "disabled"): + result["usageStatusText"] = "Copilot limits unavailable" + + # A budget must not advertise a meter on a machine the CLI has never + # touched: the estimate needs local numbers to mean anything. + budget = config.get("monthlyCredits") + if isinstance(budget, (int, float)) and budget > 0 and has_stats: + result["limits"] = [{ + "label": "Monthly allowance (est.)", + "percent": min(1.0, month_credits / float(budget)), + "resetsAt": resets_at, + }] + return result # ------------------------------------------------------------------ record @@ -934,33 +875,25 @@ def main() -> int: # collector runs. parser.add_argument("--force", action="store_true") parser.add_argument("--limits-only", action="store_true") - parser.add_argument( - "--copilot-home", default=os.environ.get("COPILOT_HOME") or "~/.copilot" - ) + parser.add_argument("--copilot-home", default=os.environ.get("COPILOT_HOME") or "~/.copilot") args = parser.parse_args() home_dir = expand_path(args.copilot_home) - max_age = ( - 0 - if args.force - else (LIMITS_ONLY_REUSE_SECONDS if args.limits_only else SCAN_REUSE_SECONDS) - ) + max_age = 0 if args.force else (LIMITS_ONLY_REUSE_SECONDS if args.limits_only else SCAN_REUSE_SECONDS) stats = cached_scan(home_dir, max_age) config = read_config() now = dt.datetime.now(dt.timezone.utc) + # The quota probe hands a GitHub token to GitHub, so it only runs where the + # Copilot CLI has been used, or where the configuration asks for it. A + # machine with gh signed in but no Copilot stays an empty record. remote = {"available": False, "reason": "disabled"} - if config.get("remote", True) is not False: + if config.get("remote") is True or (config.get("remote", True) is not False and home_dir.is_dir()): remote = fetch_remote_quota() - has_stats = ( - number(stats.get("totalPrompts")) > 0 - or float(stats.get("monthCredits") or 0.0) > 0 - ) - allowance = allowance_limits( - float(stats.get("monthCredits") or 0.0), config, remote, now, has_stats - ) + has_stats = number(stats.get("totalPrompts")) > 0 or float(stats.get("monthCredits") or 0.0) > 0 + allowance = allowance_limits(float(stats.get("monthCredits") or 0.0), config, remote, now, has_stats) record: dict[str, Any] = { "schemaVersion": 1, @@ -973,6 +906,8 @@ def main() -> int: "usageStatusText": allowance["usageStatusText"], "authHelpText": allowance["authHelpText"], "limits": allowance["limits"], + "limitsStale": allowance["limitsStale"], + "limitsFetchedAt": allowance["limitsFetchedAt"], } for key in ( "todayPrompts", @@ -993,4 +928,3 @@ def main() -> int: if __name__ == "__main__": raise SystemExit(main()) - diff --git a/manual/17-ai.md b/manual/17-ai.md index f1fadbf93f7..7d5f5820e80 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, Cursor, OpenCode, 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, Cursor, OpenCode, GitHub Copilot, and Fireworks are covered out of the box. 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 bccf9d6a95f..1e36cf70129 100644 --- a/shell/plugins/agents/README.md +++ b/shell/plugins/agents/README.md @@ -143,46 +143,36 @@ Without a Go plan the panel still shows local token usage. ### Copilot allowance -Copilot reads the CLI's session store under `~/.copilot`, honoring -`COPILOT_HOME`. - GitHub exposes no supported personal-account usage API — its Copilot REST -endpoints are organization- or enterprise-scoped — so the collector reads what -the Copilot CLI keeps locally. Its session store (`~/.copilot/session-store.db`) -records every billed request in an `assistant_usage_events` table: day, model, -session, the token split, and the cost in nano AI units (1e9 nano units = 1 -AI credit ≈ $0.01). The panel's token sections come from that table; CLI -versions whose store predates it are served by the session transcripts under -`~/.copilot/session-state` instead, whose `session.shutdown` events carry the -same per-model token split and cost. Neither source is network-dependent. - -The monthly meter is an allowance of AI credits. The account-wide figure -comes from GitHub's internal quota endpoint (the same one the editor plugins -ask), which counts every machine, IDE, and agent on the entitlement — not just -this machine's CLI. It needs a token from `COPILOT_QUOTA_TOKEN`, `GH_TOKEN`, -`GITHUB_TOKEN`, or `gh auth token`; every failure is soft and falls back to -the local estimate. The included allowance is not discoverable locally, so it -is configured in `~/.config/omarchy/agents/copilot.json`: +endpoints are organization- or enterprise-scoped — so the token sections come +from what the Copilot CLI keeps under `~/.copilot` (or `COPILOT_HOME`). Its +session store (`session-store.db`) records every billed request in an +`assistant_usage_events` table: day, model, session, the token split, and the +cost in nano AI units (1e9 nano units = 1 AI credit ≈ $0.01). CLI versions +whose store predates it are served by the session transcripts under +`session-state` instead, whose `session.shutdown` events carry the same +per-model token split and cost. + +The monthly meter is the account-wide allowance of AI credits from GitHub's +internal quota endpoint (the same one the editor plugins ask), which counts +every machine, IDE, and agent on the entitlement and names the plan and the +reset date. It needs a token from `COPILOT_QUOTA_TOKEN`, `GH_TOKEN`, +`GITHUB_TOKEN`, or `gh auth token`, and only runs where the Copilot CLI has +been used, so a machine with `gh` signed in but no Copilot stays out of the +panel. A failed check keeps the last answer, dimmed. When the endpoint can't +be asked at all, a budget in `~/.config/omarchy/agents/copilot.json` gives a +meter labeled "(est.)" from this machine's CLI spend alone: ```json { - "plan": "pro", "monthlyCredits": 1500, "remote": true } ``` -`plan` is one of `free`, `student`, `pro`, `pro+`, `max`, `business`, -`enterprise` and picks the default allowance (Pro 1500, Pro+ 7000, Max 20000, -Business 1900, Enterprise 3900 credits per month; existing seat-based -customers may see promotional amounts through September 2026). -`monthlyCredits` overrides the plan table, for AI-credit plans without a -published number or for a personal budget. With neither set there is no local -estimate and no plan label, because the plan cannot be read from disk. `remote` -(default true) turns the account-wide quota probe on or off. The allowance -window resets on the 1st of each month at 00:00 UTC, GitHub's own boundary. A -limit labeled "(est.)" is the local estimate — this machine's CLI spend only; a -plain "Monthly allowance" is the live account figure. +`remote` set to `false` turns the quota probe off, and `true` runs it even +without a Copilot CLI home. The estimated window resets on the 1st of each +month at 00:00 UTC, GitHub's own boundary. ### Fireworks balance diff --git a/shell/plugins/agents/manifest.json b/shell/plugins/agents/manifest.json index ea8b873b1ad..9f0ee44718e 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, Cursor, OpenCode, and Fireworks usage, limits, and pace in a native Omarchy bar panel.", + "description": "Claude Code, Codex, Grok, Cursor, OpenCode, GitHub Copilot, and Fireworks usage, limits, and pace in a native Omarchy bar panel.", "kinds": ["bar-widget"], "activation": "on-demand", "entryPoints": { @@ -24,6 +24,7 @@ "grok": { "enabled": true }, "cursor": { "enabled": true }, "opencode": { "enabled": true }, + "copilot": { "enabled": true }, "fireworks": { "enabled": true } }, "refreshIntervalSec": 900, diff --git a/test/shell.d/agent-usage-copilot-test.sh b/test/shell.d/agent-usage-copilot-test.sh index bac95a6838a..4541db2b04e 100755 --- a/test/shell.d/agent-usage-copilot-test.sh +++ b/test/shell.d/agent-usage-copilot-test.sh @@ -18,7 +18,7 @@ trap 'rm -rf "$TEST_HOME" "$EMPTY_HOME" "$TRANSCRIPT_HOME" "$MIDNIGHT_HOME" "$HI for home in "$TEST_HOME" "$EMPTY_HOME" "$TRANSCRIPT_HOME" "$MIDNIGHT_HOME" "$HISTORY_HOME"; do mkdir -p "$home/.copilot/session-state" "$home/.config/omarchy/agents" "$home/.cache" - printf '{"plan": "pro", "remote": false}' >"$home/.config/omarchy/agents/copilot.json" + printf '{"monthlyCredits": 1500, "remote": false}' >"$home/.config/omarchy/agents/copilot.json" done rm -rf "$EMPTY_HOME/.copilot" @@ -99,10 +99,6 @@ pass "Copilot collector splits cached tokens off the prompt" fail "Copilot collector reports today's day row" "$result" pass "Copilot collector reports today's day row" -[[ $(jq -r '.tierLabel' <<<"$result") == "Pro" ]] || - fail "Copilot collector exposes the configured plan label" "$result" -pass "Copilot collector exposes the configured plan label" - [[ $(jq -r '.limits[0].label' <<<"$result") == "Monthly allowance (est.)" ]] || fail "Copilot collector labels the local estimate" "$result" pass "Copilot collector labels the local estimate" @@ -115,33 +111,35 @@ pass "Copilot collector divides month credits by allowance" fail "Copilot collector sets an allowance reset time" "$result" pass "Copilot collector sets an allowance reset time" -# The plan cannot be read from disk, so an unconfigured one is not presented as -# a tier: no label and no estimated meter, while the local stats still show. +# Without the live quota or a budget there is no plan and no meter, while the +# local stats still show. printf '{"remote": false}' >"$TEST_HOME/.config/omarchy/agents/copilot.json" result=$(run_collector_in "$TEST_HOME") [[ $(jq -r '.ready' <<<"$result") == "true" && $(jq -r '.tierLabel' <<<"$result") == "" && $(jq -c '.limits' <<<"$result") == "[]" ]] || - fail "Copilot collector leaves an unconfigured plan unknown" "$result" -pass "Copilot collector leaves an unconfigured plan unknown" + fail "Copilot collector shows no meter without the quota or a budget" "$result" +pass "Copilot collector shows no meter without the quota or a budget" printf '{"monthlyCredits": 3000, "remote": false}' >"$TEST_HOME/.config/omarchy/agents/copilot.json" result=$(run_collector_in "$TEST_HOME") [[ $(jq -r '.tierLabel' <<<"$result") == "" && $(jq -r '.limits[0].percent' <<<"$result") == "0.0005" ]] || - fail "Copilot collector estimates from monthlyCredits without a plan" "$result" -pass "Copilot collector estimates from monthlyCredits without a plan" + fail "Copilot collector estimates from a monthlyCredits budget" "$result" +pass "Copilot collector estimates from a monthlyCredits budget" # The live quota answers in AI credits, account-wide; it wins over the local -# estimate and supplies the plan and reset date. -printf '{"plan": "pro", "remote": true}' >"$TEST_HOME/.config/omarchy/agents/copilot.json" +# estimate and supplies the plan and reset date. A failed check after it +# keeps that answer, marked stale. +printf '{"monthlyCredits": 3000}' >"$TEST_HOME/.config/omarchy/agents/copilot.json" result=$(HOME="$TEST_HOME" COPILOT_HOME="$TEST_HOME/.copilot" XDG_CONFIG_HOME="$TEST_HOME/.config" \ - XDG_CACHE_HOME="$TEST_HOME/.cache" COPILOT_QUOTA_TOKEN="test-token" \ + XDG_CACHE_HOME="$TEST_HOME/.cache" COPILOT_QUOTA_TOKEN="test-token" EMPTY_HOME="$EMPTY_HOME" \ python3 - "$ROOT/bin/omarchy-agent-usage-copilot" <<'PY' import datetime as dt import importlib.machinery import importlib.util import io import json +import os import sys loader = importlib.machinery.SourceFileLoader("collector", sys.argv[1]) @@ -163,17 +161,46 @@ payload = { } }, } -collector.urllib.request.urlopen = lambda request, timeout=None: io.BytesIO(json.dumps(payload).encode()) +def run(*extra): + out = io.StringIO() + sys.argv = ["omarchy-agent-usage-copilot", "--force", *extra] + sys.stdout, real_stdout = out, sys.stdout + collector.main() + sys.stdout = real_stdout + return json.loads(out.getvalue()) + + +def offline(request, timeout=None): + raise collector.urllib.error.URLError("offline") + + +def unasked(request, timeout=None): + raise AssertionError("the quota should not have been asked") -out = io.StringIO() -sys.argv = ["omarchy-agent-usage-copilot", "--force"] -sys.stdout, real_stdout = out, sys.stdout -collector.main() -sys.stdout = real_stdout -print(json.dumps({"record": json.loads(out.getvalue()), "reset": reset.isoformat()})) + +collector.urllib.request.urlopen = lambda request, timeout=None: io.BytesIO(json.dumps(payload).encode()) +record = run() +collector.urllib.request.urlopen = offline +stale = run() +# gh signed in on a machine that never ran the Copilot CLI: nothing is asked. +collector.urllib.request.urlopen = unasked +clean = run("--copilot-home", os.path.join(os.environ["EMPTY_HOME"], ".copilot")) +print(json.dumps({"record": record, "stale": stale, "clean": clean, "reset": reset.isoformat()})) PY ) +[[ $(jq -c '{fresh: .record.limitsStale, stamped: (.record.limitsFetchedAt > 0)}' <<<"$result") == '{"fresh":false,"stamped":true}' ]] || + fail "Copilot collector stamps the live allowance" "$result" +pass "Copilot collector stamps the live allowance" + +[[ $(jq -c '{same: (.stale.limits == .record.limits), stale: .stale.limitsStale, stamp: (.stale.limitsFetchedAt == .record.limitsFetchedAt), tier: .stale.tierLabel}' <<<"$result") == '{"same":true,"stale":true,"stamp":true,"tier":"Business"}' ]] || + fail "Copilot collector keeps the last live allowance, stale, after a failed check" "$result" +pass "Copilot collector keeps the last live allowance, stale, after a failed check" + +[[ $(jq -c '.clean | {ready, tierLabel, limits}' <<<"$result") == '{"ready":false,"tierLabel":"","limits":[]}' ]] || + fail "Copilot collector doesn't ask GitHub where the CLI never ran" "$result" +pass "Copilot collector doesn't ask GitHub where the CLI never ran" + [[ $(jq -r '.record.limits[0].label' <<<"$result") == "Monthly allowance" && $(jq -r '.record.limits[0].percent' <<<"$result") == "0.25" ]] || fail "Copilot collector reads the live allowance in credits" "$result" pass "Copilot collector reads the live allowance in credits" From ac236c2f693c06f3616f3e6563d5a4f6d699759d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Andr=C3=A9s=20Antonio=20Riveros?= Date: Mon, 7 Sep 2026 00:00:47 -0700 Subject: [PATCH 024/546] Add a Muse collector to the agents panel MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Local metering (tokens, prompts, sessions, per-day and per-model breakdowns) from native muse sessions and opencode, following the claude/codex record contract, without counting a session event twice. Probe the client's own key endpoint (POST api.meta.ai/muse-code/key) for subs_usage session + weekly percents with reset times plus the tier name — the same figures the TUI's /usage overlay shows. Limits stay strictly additive and authoritative: pay-as-you-go accounts and active plans that return no subs_usage show local metering only, silently; 401/403 reads as an expired sign-in, a transport failure advises a retry. The OAuth token travels only in the Authorization header; neither it nor the minted key is persisted. Co-Authored-By: Jackson Roussin --- bin/omarchy-agent-usage-muse | 793 ++++++++++++++++++ manual/17-ai.md | 2 +- shell/plugins/agents/README.md | 23 +- shell/plugins/agents/assets/muse-light.svg | 1 + shell/plugins/agents/assets/muse.svg | 1 + shell/plugins/agents/manifest.json | 3 +- test/shell.d/agent-usage-muse-dedup-test.sh | 95 +++ test/shell.d/agent-usage-muse-limits-test.sh | 182 ++++ test/shell.d/agent-usage-muse-scanner-test.sh | 350 ++++++++ 9 files changed, 1446 insertions(+), 4 deletions(-) create mode 100755 bin/omarchy-agent-usage-muse create mode 100644 shell/plugins/agents/assets/muse-light.svg create mode 100644 shell/plugins/agents/assets/muse.svg create mode 100755 test/shell.d/agent-usage-muse-dedup-test.sh create mode 100755 test/shell.d/agent-usage-muse-limits-test.sh create mode 100755 test/shell.d/agent-usage-muse-scanner-test.sh diff --git a/bin/omarchy-agent-usage-muse b/bin/omarchy-agent-usage-muse new file mode 100755 index 00000000000..d5c264484af --- /dev/null +++ b/bin/omarchy-agent-usage-muse @@ -0,0 +1,793 @@ +#!/usr/bin/python3 +# omarchy:summary=Print the Muse (Meta) usage record as JSON +# omarchy:args=[--force] [--limits-only] +# omarchy:hidden=true +"""Collect Muse (Meta Muse Spark) local usage plus subscription windows. + +Local metering (always available, no network): token counts, prompt and +session counts, per-day and per-model breakdowns. Sources (all local): + - native `muse` CLI sessions in $XDG_DATA_HOME/muse/sessions, defaulting to + ~/.local/share/muse/sessions (`model_completed` events carry usage), + - opencode sessions with providerID == "meta" (any Muse Spark model), + - opencode sessions with providerID == "opencode" whose modelID contains + "muse-spark" (e.g. muse-spark-1.3-contributor-free via the opencode proxy). + +Subscription windows (additive, authoritative only): when the CLI is signed +in via `muse login`, the collector probes the client's own key endpoint +(`POST https://api.meta.ai/muse-code/key` with the OAuth token from +`MUSE_AUTH_PATH`, else `~/.config/muse/auth.json`) for `subs_usage` +session (5-hour) + weekly percents with reset times plus `subs_tier_name` — +the same figures the TUI's `/usage` overlay shows. Minting is idempotent so +polling is safe; the token travels only in the Authorization header and +neither it nor the returned API key is persisted. Pay-as-you-go accounts +and active plans that return no `subs_usage` (e.g. Muse Code Power Usage) +have no subscription windows, so limits stay empty — tier plus local +metering still show. No estimated meters from manual caps: without an +authoritative probe the panel shows measured usage only. The agents panel +only ever reads the JSON this prints. +""" + +import argparse +import fcntl +import hashlib +import json +import os +import re +import sqlite3 +import sys +import tempfile +import time +import urllib.error +import urllib.request +from datetime import datetime, timedelta, timezone +from pathlib import Path + +AGENT_ID = "muse" +AGENT_NAME = "Muse" + +# Model ids and session files are read from local data but flow into the +# panel JSON, so bound them: overlong model names fail the row closed, and +# symlinked or gigantic session files are skipped instead of followed. +MAX_MODEL_NAME_LEN = 120 +MAX_SESSION_BYTES = 512 * 1024 * 1024 + +# Cache schema version. Bump when the cached stats shape changes so old +# caches are rescanned instead of trusted. +CACHE_SCHEMA_VERSION = 2 + +# Same cache semantics as the codex collector: a scan this recent only dedups +# concurrent collector runs; --limits-only may reuse a scan for 15 minutes. +SCAN_REUSE_SECONDS = 20 +LIMITS_ONLY_REUSE_SECONDS = 900 + +# Subscription probe reuse: polling mints the same key (idempotent) but the +# answer moves slowly, so repeated panel opens reuse a recent probe. +PROBE_MIN_INTERVAL_SECONDS = 60 + + +def key_endpoint() -> str: + # Overridable for tests, which point it at a local stub instead of Meta's + # key service. + return os.environ.get("MUSE_KEY_ENDPOINT") or "https://api.meta.ai/muse-code/key" + + +def number(value): + try: + return int(float(value or 0)) + except Exception: + return 0 + + +def token_number(value): + return max(0, number(value)) + + +def model_name(raw): + value = str(raw or "muse-spark").rstrip("/").split("/")[-1] + if not value: + value = "muse-spark" + # Model ids flow into the panel JSON: cap the length so a crafted session + # file or DB row cannot bloat the record or the UI. Callers skip the row + # when this returns None (fail closed). + if len(value) > MAX_MODEL_NAME_LEN: + return None + return value + + +def local_day(value): + """Normalize assorted timestamps to a local YYYY-MM-DD date string.""" + # bool is an int subclass: True would otherwise file under 1970. + if value is None or isinstance(value, bool): + return datetime.now().astimezone().date().isoformat() + if isinstance(value, (int, float)): + seconds = float(value) + # Native muse recorded_at is microseconds; opencode created is millis; + # codex-style seconds also possible. + if seconds > 10_000_000_000_000: + seconds = seconds / 1_000_000 + elif seconds > 10_000_000_000: + seconds = seconds / 1000 + try: + return datetime.fromtimestamp(seconds).astimezone().date().isoformat() + except Exception: + return datetime.now().astimezone().date().isoformat() + text = str(value).strip() + try: + # Some writers emit the epoch as a string: route micro/milli/second + # magnitudes through the numeric branch instead of misfiling them as today. + return local_day(float(text)) + except ValueError: + pass + 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().astimezone().date().isoformat() + + +now = datetime.now().astimezone() +today = now.date().isoformat() +recent_dates = [(now.date() - timedelta(days=offset)).isoformat() for offset in range(6, -1, -1)] +recent = {day: 0 for day in recent_dates} +today_tokens_by_model = {} +model_usage = {} +today_sessions = set() +active_days = set() + +today_prompts = 0 +today_total_tokens = 0 +total_prompts = 0 +total_sessions = set() + + +def add_usage(day, session_key, model, input_tokens, output_tokens, cache_read, cache_write): + global today_prompts, today_total_tokens, total_prompts + total = input_tokens + output_tokens + cache_read + cache_write + if total <= 0: + return + total_prompts += 1 + total_sessions.add(session_key) + active_days.add(day) + + bucket = model_usage.setdefault(model, { + "inputTokens": 0, + "outputTokens": 0, + "cacheReadInputTokens": 0, + "cacheCreationInputTokens": 0, + }) + bucket["inputTokens"] += input_tokens + bucket["outputTokens"] += output_tokens + bucket["cacheReadInputTokens"] += cache_read + bucket["cacheCreationInputTokens"] += cache_write + + if day in recent: + recent[day] += total + if day == today: + today_prompts += 1 + today_sessions.add(session_key) + today_total_tokens += total + today_tokens_by_model[model] = today_tokens_by_model.get(model, 0) + total + + +def muse_sessions_root(): + override = (os.environ.get("MUSE_DATA_DIR") or "").strip() + if override: + base = Path(os.path.expandvars(os.path.expanduser(override))) + else: + base = Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) / "muse" + return base / "sessions" + + +def iter_native_events(path): + """Yield (event_dict, recorded_at, event_id) tuples from a native session.jsonl file. + + Handles both plain event lines and retained_frame wrappers whose children + carry record_json payloads. + """ + try: + with path.open(errors="replace") as handle: + for raw in handle: + raw = raw.strip() + if not raw: + continue + try: + entry = json.loads(raw) + except Exception: + continue + if not isinstance(entry, dict): + continue + candidates = [entry] + children = entry.get("children") + if isinstance(children, list): + for child in children: + if not isinstance(child, dict): + continue + nested = child.get("record_json") + if not nested: + continue + try: + parsed = json.loads(nested) + except Exception: + continue + if isinstance(parsed, dict): + candidates.append(parsed) + for event in candidates: + payload = event.get("payload") or {} + if not isinstance(payload, dict): + continue + yield payload, event.get("recorded_at"), event.get("id") + except OSError: + return + + +def scan_native_muse_sessions(): + root = muse_sessions_root() + if not root.exists(): + return + # Full history, no mtime cutoff: all-time totals and active days need every + # session still on disk (the same tradeoff the claude collector makes). + # Session files are small JSONL and the scan cache dedups concurrent runs, + # so each refresh pays one linear pass at most. + try: + resolved_root = root.resolve() + except OSError: + return + files = [] + try: + # Session views contain copies of the same history. Check only paths + # below the sessions root, so a custom data parent is not excluded. + files = [p for p in root.rglob("session.jsonl") + if not any(part.startswith(".msp-view") for part in p.relative_to(root).parts[:-1]) + and p.is_file()] + except OSError: + return + seen_event_ids = set() + for path in files: + try: + # Fail closed on entries that could escape the sessions tree or stall + # the scan: symlinked files are never followed (even when their target + # sits inside the tree, so a link cannot double-count a session), files + # reached via a symlinked directory are rejected by the resolve check, + # and gigantic files are skipped. + if path.is_symlink(): + continue + if not path.resolve().is_relative_to(resolved_root): + continue + st = path.stat() + if st.st_size > MAX_SESSION_BYTES: + continue + mtime_day = local_day(st.st_mtime) + except OSError: + continue + current_model = "" + for payload, recorded_at, event_id in iter_native_events(path): + kind = payload.get("kind") + if kind == "metadata": + record = payload.get("record") or {} + if isinstance(record, dict) and record.get("model_id"): + current_model = model_name(record.get("model_id")) + continue + if kind == "run_model": + record = payload.get("record") or {} + if isinstance(record, dict) and record.get("model_id"): + current_model = model_name(record.get("model_id")) + continue + if kind != "run": + continue + event = payload.get("event") or {} + if not isinstance(event, dict): + continue + if event.get("kind") != "model_completed": + continue + usage = event.get("usage") or {} + if not isinstance(usage, dict): + continue + raw_model = event.get("model") + # An oversized metadata model is an explicit rejection, not a reason to + # silently put the following completion in the generic fallback bucket. + if not raw_model and current_model is None: + continue + model = model_name(raw_model or current_model) + if model is None: + continue + input_tokens = token_number(usage.get("input_tokens")) + # Native Muse reports reasoning as a subset of output_tokens. + output_tokens = token_number(usage.get("output_tokens")) + cache_read = token_number(usage.get("cache_read_tokens", usage.get("cached_tokens"))) + cache_write = token_number(usage.get("cache_write_tokens")) + # Cached tokens are a subset of input: keep the split without double + # counting (same convention as the codex collector). + input_tokens = max(0, input_tokens - cache_read - cache_write) + if not (input_tokens or output_tokens or cache_read or cache_write): + continue + # Native record IDs survive copies and retained-frame wrapping. Claim + # only valid completions: metadata or malformed usage must not suppress + # a later usable record. Missing/malformed IDs cannot identify duplicates. + if isinstance(event_id, str) and event_id: + if event_id in seen_event_ids: + continue + seen_event_ids.add(event_id) + day = local_day(recorded_at) if recorded_at else mtime_day + # Subagent trees roll up into their parent session: normalize both the + # main session.jsonl and any subagent/.../session.jsonl to the session + # directory, so one user session counts once however many subagents + # it spawned. + session_key = "muse:" + re.sub(r"(/subagent/.*)?/session\.jsonl$", "", str(path)) + add_usage(day, session_key, model, + input_tokens, output_tokens, cache_read, cache_write) + + +def scan_opencode_sessions(): + """opencode rows for meta provider Muse models + opencode-proxied muse-spark. + + Returns True when the scan ran to completion (safe to cache). + """ + db = Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) / "opencode" / "opencode.db" + 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") + # LIKE gates are pure acceleration (see the codex collector for the full + # rationale); the json_extract gates match the Python filter below, which + # stays the authority for the rows that reach it. The provider gate keeps + # assistant rows from other providers — including ones that merely mention + # muse-spark in their text — out of the Python parse. + query = ( + "SELECT session_id, data FROM message" + " WHERE data LIKE '%\"role\"%:%\"assistant\"%'" + " AND (data LIKE '%\"providerID\"%:%\"meta\"%' OR data LIKE '%muse-spark%')" + " 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 IN ('meta', 'opencode')" + ) + for session_id, raw in conn.execute(query): + try: + entry = json.loads(raw) + if not isinstance(entry, dict) or entry.get("role") != "assistant": + continue + provider = str(entry.get("providerID") or "") + model = model_name(entry.get("modelID")) + if model is None: + continue + if provider == "meta": + # The meta provider only serves Muse Spark models, so any model on + # it counts (muse-spark 1.1/1.2/1.3, ...). + pass + elif provider == "opencode" and "muse-spark" in model.lower(): + pass # opencode-proxied Muse Spark (e.g. 1.3-contributor-free) + else: + continue + tokens = entry.get("tokens") + if not isinstance(tokens, dict): + tokens = {} + cache = tokens.get("cache") + if not isinstance(cache, dict): + cache = {} + # opencode reports input exclusive of cache: keep the token buckets + # separate, no subtraction. + input_tokens = token_number(tokens.get("input")) + output_tokens = token_number(tokens.get("output")) + token_number(tokens.get("reasoning")) + cache_read = token_number(cache.get("read")) + cache_write = token_number(cache.get("write")) + if not (input_tokens or output_tokens or cache_read or cache_write): + continue + # A malformed time object must not drop the row's valid tokens: + # non-dict shapes fall back to today via local_day(None). + stamped = entry.get("time") + created = stamped.get("created") if isinstance(stamped, dict) else None + day = local_day(created) + add_usage(day, "opencode:" + str(session_id), model, + input_tokens, output_tokens, cache_read, cache_write) + except Exception: + continue + except sqlite3.Error: + return False + finally: + try: + conn.close() + except Exception: + pass + return True + + +# ------------------------------------------------------------------ limits +# +# Additive subscription windows only. Local metering above always runs; this +# section only fills `limits`/`tierLabel`/status when the authoritative key +# probe answers. Pay-as-you-go and active plans without `subs_usage` have no +# subscription windows: no token, a PAYG probe, or an active probe with no +# subs_usage means limits stay empty (never estimated, never zeroed) and +# local metering still shows. The OAuth token goes nowhere but +# the Authorization header; only plan label + usage figures enter the record. + + +def muse_config_dir() -> Path: + base = os.environ.get("XDG_CONFIG_HOME") or (Path.home() / ".config") + return Path(base) / "muse" + + +def read_access_token() -> str: + try: + raw = (os.environ.get("MUSE_AUTH_PATH") or "").strip() + path = Path(os.path.expandvars(os.path.expanduser(raw))) if raw else (muse_config_dir() / "auth.json") + data = json.loads(path.read_text(encoding="utf-8")) + except Exception: + return "" + # Every level checked: a login file that parses to a list, or whose + # `providers` is not an object, is a missing credential, not a crash. + if not isinstance(data, dict): + return "" + providers = data.get("providers") + if not isinstance(providers, dict): + return "" + login = providers.get("meta") + if not isinstance(login, dict): + return "" + return str(login.get("access_token") or "") + + +def normalize_reset_at(value) -> str: + if value is None: + return "" + try: + ts = float(value) + except Exception: + return "" + if ts <= 0: + return "" + # Epoch seconds; accept milliseconds too. + if ts >= 1e12: + ts = ts / 1000 + try: + return datetime.fromtimestamp(ts, timezone.utc).isoformat() + except Exception: + return "" + + +def probe_key(access_token: str) -> dict: + request = urllib.request.Request( + key_endpoint(), + data=json.dumps({}).encode(), + method="POST", + headers={ + "Authorization": "Bearer " + access_token, + "Accept": "application/json", + "Content-Type": "application/json", + }, + ) + try: + with urllib.request.urlopen(request, timeout=10) as response: + raw_body = response.read().decode("utf-8", errors="replace") + except urllib.error.HTTPError as error: + if error.code in (401, 403): + return {"ok": False, "expired": True, + "helpText": "Muse's saved sign-in was rejected. Run `muse login` to refresh it."} + return {"ok": False, + "helpText": f"Muse's usage endpoint returned status {error.code}. Local Muse stats are still shown."} + except Exception: + return {"ok": False, "transport": True, + "helpText": "Couldn't reach Muse's usage endpoint. Retrying shortly. Local Muse stats are still shown."} + + try: + payload = json.loads(raw_body) + except Exception: + # The server answered but not with JSON: a shape change, not a network + # failure. No retryAdvised — re-asking won't fix the parser. + return {"ok": False, "helpText": "Muse's usage endpoint returned an unreadable response. Local Muse stats are still shown."} + + usage = payload.get("subs_usage") if isinstance(payload, dict) else None + if not isinstance(usage, dict): + # PAYG account: key mints fine but there is no subscription window to + # meter (live shape: is_subs_active=false, subs_tier_name=null, no + # subs_usage). Local metering is the whole story — stay silent, not an + # error. + if isinstance(payload, dict) and payload.get("is_subs_active") is False: + return {"ok": False, "payg": True} + # Active plan without windows (live shape on Muse Code Power Usage: + # is_subs_active=true, subs_tier_name set, no subs_usage key at all). + # That is the steady answer, not a snapshot about to appear: keep the + # tier, show local metering only, stay silent like PAYG. + if isinstance(payload, dict) and payload.get("is_subs_active") is True: + tier_hint = str(payload.get("subs_tier_name") or "") + return {"ok": False, "silent": True, "tierLabel": tier_hint} + tier_hint = str(payload.get("subs_tier_name") or "") if isinstance(payload, dict) else "" + return {"ok": False, "tierLabel": tier_hint, + "helpText": "Muse's usage endpoint returned no usage. Local Muse stats are still shown."} + + limits = [] + window = usage.get("window") or {} + if isinstance(window, dict) and window.get("used_percent") is not None: + try: + percent = min(1.0, max(0.0, float(window.get("used_percent")) / 100.0)) + except Exception: + percent = -1.0 + if percent >= 0: + mins = number(window.get("window_duration_mins")) + label = "Session (5-hour)" if mins == 300 else (f"Session ({mins}m)" if mins > 0 else "Session") + limits.append({"label": label, "percent": percent, "resetsAt": normalize_reset_at(window.get("resets_at"))}) + weekly = usage.get("weekly") or {} + if isinstance(weekly, dict) and weekly.get("used_percent") is not None: + try: + percent = min(1.0, max(0.0, float(weekly.get("used_percent")) / 100.0)) + except Exception: + percent = -1.0 + if percent >= 0: + limits.append({"label": "Weekly (7-day)", "percent": percent, "resetsAt": normalize_reset_at(weekly.get("resets_at"))}) + + if not limits: + tier_hint = str(payload.get("subs_tier_name") or "") if isinstance(payload, dict) else "" + return {"ok": False, "tierLabel": tier_hint, + "helpText": "Muse's usage endpoint returned no limits. Local Muse stats are still shown."} + tier = "" + if isinstance(payload, dict): + tier = str(payload.get("subs_tier_name") or "") + return {"ok": True, "limits": limits, "tierLabel": tier} + + +def limit_window_open(entry: dict, now: datetime) -> bool: + # A cached percentage lives only until its window rolls over; past-reset + # figures describe a period that is over. No (or unreadable) reset is kept. + 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) -> list: + 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 collect_limits(force: bool) -> dict: + result = {"limits": [], "tierLabel": "Account", "usageStatusText": "", "authHelpText": ""} + + probe_cache = None + try: + probe_cache = cache_root() / "muse-limits.json" + cached = json.loads(probe_cache.read_text(encoding="utf-8")) if probe_cache.exists() else {} + except Exception as exc: + print(f"omarchy-agent-usage-muse: could not read limits cache ({exc})", file=sys.stderr) + cached = {} + if not isinstance(cached, dict): + cached = {} + fallback = usable_cached_limits(cached) + fallback_tier = str(cached.get("tierLabel") or "") + + access_token = read_access_token() + if access_token == "": + # No login: PAYG or signed-out. Never estimated, never an auth nag — + # local metering is valid on its own. Reuse unexpired cached windows if + # a previous subscription probe left any. + if fallback: + result["limits"] = fallback + if fallback_tier: + result["tierLabel"] = fallback_tier + return result + + fetched_at = number(cached.get("fetchedAtMs")) / 1000 + if (fallback or fallback_tier) and not force and time.time() - fetched_at < PROBE_MIN_INTERVAL_SECONDS: + result["limits"] = fallback + if fallback_tier: + result["tierLabel"] = fallback_tier + return result + + probe = probe_key(access_token) + if probe["ok"]: + result["limits"] = probe["limits"] + if probe.get("tierLabel"): + result["tierLabel"] = probe["tierLabel"] + elif fallback_tier: + result["tierLabel"] = fallback_tier + if probe_cache is not None: + 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-muse: could not write limits cache ({exc})", file=sys.stderr) + return result + + if probe.get("transport"): + result["retryAdvised"] = True + if probe.get("payg"): + # Authenticated PAYG: no subscription windows exist. Keep local metering + # only, silent — not an error, not expired. + return result + if probe.get("silent"): + # Active plan with no windows: same silence as PAYG, but keep the plan + # name. Stale cached windows are dropped — they belong to a different + # plan shape, not this one. + if probe.get("tierLabel"): + result["tierLabel"] = probe["tierLabel"] + return result + if probe.get("expired"): + result["usageStatusText"] = "Sign-in expired" + result["authHelpText"] = probe["helpText"] + result["limits"] = fallback + if fallback_tier: + result["tierLabel"] = fallback_tier + return result + if fallback: + result["limits"] = fallback + if fallback_tier: + result["tierLabel"] = fallback_tier + elif probe.get("tierLabel"): + result["tierLabel"] = probe["tierLabel"] + else: + result["usageStatusText"] = "Muse limits unavailable" + result["authHelpText"] = probe["helpText"] + if probe.get("tierLabel"): + result["tierLabel"] = probe["tierLabel"] + return result + + +def cache_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(): + db = Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) / "opencode" / "opencode.db" + digest = hashlib.sha1( + (str(Path.home()) + "\n" + str(muse_sessions_root()) + "\n" + str(db)).encode("utf-8") + ).hexdigest()[:16] + root = cache_root() + return root / f"muse-scan-{digest}.json", root / f"muse-scan-{digest}.lock" + + +def read_fresh_json(path, max_age_seconds): + if max_age_seconds <= 0 or not path.exists(): + return None + try: + 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, payload): + 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 read_cached_stats(cache_file, max_age_seconds): + cached = read_fresh_json(cache_file, max_age_seconds) + if not isinstance(cached, dict) or cached.get("schemaVersion") != CACHE_SCHEMA_VERSION: + return None + if cached.get("scanDate") != today: + return None + stats = cached.get("stats") + if not isinstance(stats, dict): + return None + required = ( + "todayPrompts", "todaySessions", "todayTotalTokens", "todayTokensByModel", + "recentDays", "totalPrompts", "totalSessions", "activeDays", "activeDates", + "modelUsage", + ) + if not all(key in stats for key in required): + return None + return stats + + +def write_cached_stats(cache_file, stats): + try: + write_json(cache_file, {"schemaVersion": CACHE_SCHEMA_VERSION, "scanDate": today, "stats": stats}) + except Exception as exc: + print(f"omarchy-agent-usage-muse: could not write usage cache ({exc})", file=sys.stderr) + + +def local_stats(): + return { + "todayPrompts": today_prompts, + "todaySessions": len(today_sessions), + "todayTotalTokens": today_total_tokens, + "todayTokensByModel": today_tokens_by_model, + "recentDays": [{"date": day, "messageCount": recent[day]} for day in recent_dates], + "totalPrompts": total_prompts, + "totalSessions": len(total_sessions), + "activeDays": len(active_days), + "activeDates": sorted(active_days), + "modelUsage": model_usage, + } + + +def run_local_scans(): + scan_native_muse_sessions() + complete = scan_opencode_sessions() + return local_stats(), complete + + +def cached_local_stats(max_age): + try: + return _cached_local_stats(max_age) + except Exception as exc: + print(f"omarchy-agent-usage-muse: cache unavailable ({exc}); scanning directly", file=sys.stderr) + stats, _ = run_local_scans() + return stats + + +def _cached_local_stats(max_age): + cache_file, lock_file = scan_cache_paths() + cached = read_cached_stats(cache_file, max_age) + if cached is not None: + return cached + with lock_file.open("w") as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + cached = read_cached_stats(cache_file, max_age) + if cached is not None: + return cached + stats, complete = run_local_scans() + if complete: + write_cached_stats(cache_file, stats) + return stats + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--force", action="store_true", help="rescan sessions, ignoring caches") + parser.add_argument("--limits-only", action="store_true", help="reuse any recent scan; kept for CLI compatibility with the panel refresh") + args = parser.parse_args() + + max_age = 0 if args.force else (LIMITS_ONLY_REUSE_SECONDS if args.limits_only else SCAN_REUSE_SECONDS) + stats = cached_local_stats(max_age) + limits = collect_limits(args.force) + + record = { + "schemaVersion": 1, + "id": AGENT_ID, + "name": AGENT_NAME, + "updatedAt": datetime.now(timezone.utc).isoformat(), + "ready": number(stats.get("totalPrompts")) > 0 or len(limits["limits"]) > 0, + "hasLocalStats": True, + # Prompts and sessions are counted, so the panel shows them in tooltips. + "hasPromptStats": True, + "todayPrompts": number(stats.get("todayPrompts")), + "todaySessions": number(stats.get("todaySessions", len(today_sessions))), + "todayTotalTokens": number(stats.get("todayTotalTokens")), + "todayTokensByModel": stats.get("todayTokensByModel") or {}, + "recentDays": stats.get("recentDays") or [], + "totalPrompts": number(stats.get("totalPrompts")), + "totalSessions": number(stats.get("totalSessions", len(total_sessions))), + "activeDays": number(stats.get("activeDays")), + "activeDates": stats.get("activeDates") or [], + "modelUsage": stats.get("modelUsage") or {}, + "tierLabel": limits["tierLabel"], + "limits": limits["limits"], + "usageStatusText": limits["usageStatusText"], + "authHelpText": limits["authHelpText"], + } + if limits.get("retryAdvised"): + record["retryAdvised"] = True + print(json.dumps(record, separators=(",", ":"), sort_keys=True)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/manual/17-ai.md b/manual/17-ai.md index 7d5f5820e80..c99186c3896 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, Cursor, OpenCode, GitHub Copilot, 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, Cursor, OpenCode, GitHub Copilot, Fireworks, and Muse are covered out of the box. 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 1e36cf70129..0df0f8bab18 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. | `opencode` | OpenCode Go's usage endpoint (rolling + weekly + monthly), with an API key or the Console session | pi and omp sessions on the Zen and Go providers, plus opencode's own `message` and `session_message` stores | | `fireworks` | Estimated prepaid balance: configured funding minus rated account costs | Fireworks billing API, grouped by day and model for the last 30 days | | `copilot` | Account-wide AI credit allowance (GitHub quota endpoint, or estimated locally) | the Copilot CLI session store's `assistant_usage_events`, with `~/.copilot/session-state` transcripts as fallback | +| `muse` | The client's key endpoint (session + weekly percents), else local metering only — never estimated | native `muse` CLI session files plus opencode sessions on the `meta` provider or `muse-spark` models | When `~/.local/state/omarchy/agents/accounts/.json` registers more than one account, the `claude`, `codex`, and `grok` records @@ -202,6 +203,24 @@ 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. +### Muse limits + +Muse reads native CLI sessions from `$XDG_DATA_HOME/muse/sessions` +(defaulting to `~/.local/share/muse/sessions`, and overridable via +`MUSE_DATA_DIR`) plus opencode sessions on the `meta` provider or +`muse-spark` models, and the OAuth token from `MUSE_AUTH_PATH` when set, +otherwise from `~/.config/muse/auth.json` (which `muse login` creates). + +The collector asks the client's own key endpoint for the subscription's +session and weekly percents with reset times — the same figures the TUI's +`/usage` overlay shows. Minting is idempotent (the same Model API key comes +back every call), so polling is safe; the token travels only in the +`Authorization` header and neither it nor the returned key is persisted. +Pay-as-you-go accounts have no subscription windows, so the tab shows +measured local usage with no meters — never estimated, never zeroed. An +active plan that returns no `subs_usage` (Muse Code Power Usage answers this +way on every probe) behaves the same way, except the tab keeps the plan name. + ## Interactions - Bar icon: left = panel, right = launch agent, middle = refresh. It turns @@ -270,5 +289,5 @@ the same account synced from two machines is not counted twice. One caveat on "all-time": the Codex collector only reads native session files touched in the last 30 days, and Fireworks requests the last 30 days from its -billing API, so their totals and day counts cover that window. Claude's cover -every transcript still on disk. +billing API, so their totals and day counts cover that window. Claude's and +Muse's cover every session still on disk. diff --git a/shell/plugins/agents/assets/muse-light.svg b/shell/plugins/agents/assets/muse-light.svg new file mode 100644 index 00000000000..0850837cf9c --- /dev/null +++ b/shell/plugins/agents/assets/muse-light.svg @@ -0,0 +1 @@ +Muse diff --git a/shell/plugins/agents/assets/muse.svg b/shell/plugins/agents/assets/muse.svg new file mode 100644 index 00000000000..daf5cc78f53 --- /dev/null +++ b/shell/plugins/agents/assets/muse.svg @@ -0,0 +1 @@ +Muse diff --git a/shell/plugins/agents/manifest.json b/shell/plugins/agents/manifest.json index 9f0ee44718e..af914b0e586 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, Cursor, OpenCode, GitHub Copilot, and Fireworks usage, limits, and pace in a native Omarchy bar panel.", + "description": "Claude Code, Codex, Grok, Cursor, OpenCode, GitHub Copilot, Fireworks, and Muse usage, limits, and pace in a native Omarchy bar panel.", "kinds": ["bar-widget"], "activation": "on-demand", "entryPoints": { @@ -25,6 +25,7 @@ "cursor": { "enabled": true }, "opencode": { "enabled": true }, "copilot": { "enabled": true }, + "muse": { "enabled": true }, "fireworks": { "enabled": true } }, "refreshIntervalSec": 900, diff --git a/test/shell.d/agent-usage-muse-dedup-test.sh b/test/shell.d/agent-usage-muse-dedup-test.sh new file mode 100755 index 00000000000..a08a176bb92 --- /dev/null +++ b/test/shell.d/agent-usage-muse-dedup-test.sh @@ -0,0 +1,95 @@ +#!/bin/bash + +set -euo pipefail + +source "$(dirname "$0")/base-test.sh" + +require_command python3 +require_command jq + +TEST_HOME=$(mktemp -d) +trap 'rm -rf "$TEST_HOME"' EXIT + +export HOME="$TEST_HOME" +export XDG_CONFIG_HOME="$TEST_HOME/config" +export XDG_CACHE_HOME="$TEST_HOME/cache" +export XDG_DATA_HOME="$TEST_HOME/data" +# Only view directories within the sessions tree should be excluded. +export MUSE_DATA_DIR="$TEST_HOME/.msp-view-parent/muse" +export MUSE_AUTH_PATH="$TEST_HOME/no-auth.json" + +python3 - <<'PY' +import json +import os +import time +from pathlib import Path + +root = Path(os.environ["MUSE_DATA_DIR"]) / "sessions" + +def event(event_id, tokens): + return { + "id": event_id, + "recorded_at": int(time.time() * 1_000_000), + "payload": {"kind": "run", "event": { + "kind": "model_completed", "model": "muse-test", + "usage": {"input_tokens": tokens}, + }}, + } + +def write(relative, entries): + path = root / relative / "session.jsonl" + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("".join(json.dumps(entry) + "\n" for entry in entries)) + +first = event("first", 110) +child = event("child", 55) +rejected = event("later-valid", 0) +rejected["payload"]["event"]["usage"] = [] +write("2026/09/08/session", [ + {"id": "first", "payload": {"kind": "metadata", "record": {"model_id": "muse-test"}}}, + first, first, + {"children": [{"record_json": json.dumps(first)}, {"record_json": json.dumps(child)}]}, + event(None, 7), event(None, 7), event([], 3), + rejected, event("later-valid", 5), +]) +write("2026/09/08/session/subagent/child", [child]) +write("2026/09/08/copied-session", [first]) +for view in (".msp-view", ".msp-view.saved", "2026/09/08/session/.msp-view-backup"): + write(view, [event(None, 9000)]) +PY + +result=$("$ROOT/bin/omarchy-agent-usage-muse" --force) +[[ $(jq -c '[.todayTotalTokens,.totalPrompts,.modelUsage["muse-test"].inputTokens]' <<<"$result") == '[187,6,187]' ]] || + fail "Muse counts unique completions across plain records, retained frames, and copied sessions" "$result" +pass "Muse counts unique completions across plain records, retained frames, and copied sessions" + +# The same IDs without view copies must produce exactly the same record +# totals; records without a usable ID are still counted individually. +python3 - <<'PY' +import os +import shutil +from pathlib import Path +root = Path(os.environ["MUSE_DATA_DIR"]) / "sessions" +for relative in (".msp-view", ".msp-view.saved", "2026/09/08/session/.msp-view-backup"): + shutil.rmtree(root / relative) +PY +without_views=$("$ROOT/bin/omarchy-agent-usage-muse" --force) +[[ $(jq -c '[.todayTotalTokens,.totalPrompts]' <<<"$without_views") == '[187,6]' ]] || + fail "Muse excludes view copies without dropping missing or malformed IDs" "$without_views" +pass "Muse excludes view copies without dropping missing or malformed IDs" + +# Discard a pre-deduplication cache even during --limits-only refreshes. +python3 - <<'PY' +import json +import os +from pathlib import Path +for path in Path(os.environ["XDG_CACHE_HOME"]).rglob("muse-scan-*.json"): + cached = json.loads(path.read_text()) + cached["schemaVersion"] = 1 + cached["stats"]["todayTotalTokens"] = 999999 + path.write_text(json.dumps(cached)) +PY +refreshed=$("$ROOT/bin/omarchy-agent-usage-muse" --limits-only) +[[ $(jq '.todayTotalTokens' <<<"$refreshed") == 187 ]] || + fail "Muse rescans caches produced before event deduplication" "$refreshed" +pass "Muse rescans caches produced before event deduplication" diff --git a/test/shell.d/agent-usage-muse-limits-test.sh b/test/shell.d/agent-usage-muse-limits-test.sh new file mode 100755 index 00000000000..dadf3f73d1c --- /dev/null +++ b/test/shell.d/agent-usage-muse-limits-test.sh @@ -0,0 +1,182 @@ +#!/bin/bash + +set -euo pipefail + +source "$(dirname "$0")/base-test.sh" + +require_command jq +require_command python3 + +unset MUSE_DATA_DIR MUSE_AUTH_PATH MUSE_KEY_ENDPOINT + +TEST_HOME=$(mktemp -d) +trap 'rm -rf "$TEST_HOME"' EXIT + +mkdir -p "$TEST_HOME/.config/muse" "$TEST_HOME/.cache" +cat >"$TEST_HOME/.config/muse/auth.json" <<'EOF' +{"schema_version":1,"providers":{"meta":{"access_token":"test-token","api_base_url":"https://api.meta.ai/v1"}}} +EOF + +# A stub key service: records the Authorization header it was called with +# and answers from a canned payload file. The collector must send the OAuth +# token as a Bearer token and never persist it anywhere. +cat >"$TEST_HOME/stub.py" <<'EOF' +import http.server +import json +import os +import sys + +seen_file = os.environ["STUB_SEEN_FILE"] +payload_file = os.environ["STUB_PAYLOAD_FILE"] + +class Handler(http.server.BaseHTTPRequestHandler): + def do_POST(self): + length = int(self.headers.get("Content-Length") or 0) + self.rfile.read(length) + with open(seen_file, "a") as handle: + handle.write((self.headers.get("Authorization") or "") + "\n") + try: + with open(payload_file) as handle: + spec = json.load(handle) + except OSError: + spec = {"status": 500} + status = spec.get("status", 200) + self.send_response(status) + if status != 200: + self.end_headers() + return + raw = spec.get("raw") + body = raw.encode() if isinstance(raw, str) else json.dumps(spec.get("body", {})).encode() + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, *args): + pass + +http.server.HTTPServer(("127.0.0.1", int(sys.argv[1])), Handler).serve_forever() +EOF + +port=$(python3 -c 'import socket; s=socket.socket(); s.bind(("127.0.0.1",0)); print(s.getsockname()[1])') +export STUB_SEEN_FILE="$TEST_HOME/seen-headers" +export STUB_PAYLOAD_FILE="$TEST_HOME/stub-payload.json" +: >"$STUB_SEEN_FILE" +printf '{"status":500}' >"$STUB_PAYLOAD_FILE" +python3 "$TEST_HOME/stub.py" "$port" & stub_pid=$! +trap 'kill "$stub_pid" 2>/dev/null; rm -rf "$TEST_HOME"' EXIT +for _ in $(seq 1 50); do + (echo >/dev/tcp/127.0.0.1/"$port") 2>/dev/null && break + sleep 0.1 +done + +run_collector() { + HOME="$TEST_HOME" XDG_DATA_HOME="$TEST_HOME/.local/share" XDG_CACHE_HOME="$TEST_HOME/.cache" XDG_CONFIG_HOME="$TEST_HOME/.config" \ + MUSE_KEY_ENDPOINT="http://127.0.0.1:$port/key" \ + "$ROOT/bin/omarchy-agent-usage-muse" "$@" +} + +set_payload() { + printf '%s' "$1" >"$STUB_PAYLOAD_FILE" + : >"$STUB_SEEN_FILE" + # Each stub case starts without a previous probe's cache, so fallback + # behavior never leaks across cases. + rm -f "$TEST_HOME/.cache/omarchy/agent-usage/muse-limits.json" +} + +# Subscriber with usage: window + weekly percents, resets, tier. +set_payload '{"body":{"subs_tier_name":"Power Usage","subs_usage":{"window":{"used_percent":8,"window_duration_mins":300,"resets_at":1788900190},"weekly":{"used_percent":15,"resets_at":1789344000}}}}' +result=$(run_collector --force) + +[[ $(jq -c '[.limits[]|{label,percent}]' <<<"$result") == '[{"label":"Session (5-hour)","percent":0.08},{"label":"Weekly (7-day)","percent":0.15}]' ]] || + fail "Muse collector reports window and weekly limits from the key endpoint" "$result" +pass "Muse collector reports window and weekly limits from the key endpoint" + +[[ $(jq -r '.limits[0].resetsAt' <<<"$result") == "2026-09-08T20:43:10+00:00" ]] || + fail "Muse collector converts window reset times to ISO" "$result" +pass "Muse collector converts window reset times to ISO" + +[[ $(jq -r '.tierLabel' <<<"$result") == "Power Usage" ]] || + fail "Muse collector reports the subscription tier" "$result" +pass "Muse collector reports the subscription tier" + +[[ $(cat "$STUB_SEEN_FILE") == "Bearer test-token" ]] || + fail "Muse collector authenticates the probe with the saved token" "$(cat "$STUB_SEEN_FILE")" +pass "Muse collector authenticates the probe with the saved token" + +# Token and minted key must not leak into the record or the caches. +grep -rq "test-token" "$TEST_HOME/.cache" "$TEST_HOME/.local" 2>/dev/null && + fail "Muse collector keeps the token out of caches" "$(grep -r "test-token" "$TEST_HOME/.cache" 2>/dev/null)" +pass "Muse collector keeps the token out of caches" + +# A fresh probe result is reused briefly, so repeated panel opens do not +# turn into a request per flick. +: >"$STUB_SEEN_FILE" +second=$(run_collector) +[[ -s $STUB_SEEN_FILE ]] && + fail "Muse collector reuses a recent probe result" "$(cat "$STUB_SEEN_FILE")" +pass "Muse collector reuses a recent probe result" + +[[ $(jq -r '.tierLabel' <<<"$second") == "Power Usage" ]] || + fail "Muse collector reuses the cached tier" "$second" +pass "Muse collector reuses the cached tier" + +# Pay-as-you-go: key mints fine but there is no subscription window to +# meter. Silent local metering — not an error, not expired. +set_payload '{"body":{"api_key":"minted","is_subs_active":false,"subs_tier_name":null}}' +payg=$(run_collector --force) +[[ $(jq -c '[.limits,.tierLabel,.usageStatusText,.authHelpText]' <<<"$payg") == '[[],"Account","",""]' ]] || + fail "Muse collector stays silent for pay-as-you-go accounts" "$payg" +pass "Muse collector stays silent for pay-as-you-go accounts" + +[[ $(jq 'has("retryAdvised")' <<<"$payg") == "false" ]] || + fail "Muse collector does not retry a pay-as-you-go answer" "$payg" +pass "Muse collector does not retry a pay-as-you-go answer" + +# Active plan without usage windows (live Muse Code Power Usage shape): +# keep the tier, show local metering only — silent, not an error. +set_payload '{"body":{"api_key":"minted","is_subs_active":true,"subs_tier_name":"Muse Code Everyday Usage"}}' +nosnap=$(run_collector --force) +[[ $(jq -c '[.limits,.tierLabel,.usageStatusText,.authHelpText]' <<<"$nosnap") == '[[],"Muse Code Everyday Usage","",""]' ]] || + fail "Muse collector stays silent for an active plan without usage windows" "$nosnap" +pass "Muse collector stays silent for an active plan without usage windows" + +[[ $(jq 'has("retryAdvised")' <<<"$nosnap") == "false" ]] || + fail "Muse collector does not retry an active plan without usage windows" "$nosnap" +pass "Muse collector does not retry an active plan without usage windows" + +# Neither active nor PAYG and no windows: still a warning, never invented meters. +set_payload '{"body":{"api_key":"minted","subs_tier_name":"Mystery"}}' +unknown=$(run_collector --force) +[[ $(jq -r '.usageStatusText' <<<"$unknown") == "Muse limits unavailable" ]] || + fail "Muse collector reports unavailable limits for an unknown shape" "$unknown" +pass "Muse collector reports unavailable limits for an unknown shape" + +# A rejected sign-in is an auth problem, not missing data: say so. +set_payload '{"status":401}' +expired=$(run_collector --force) +[[ $(jq -r '.usageStatusText' <<<"$expired") == "Sign-in expired" ]] || + fail "Muse collector reports an expired sign-in" "$expired" +pass "Muse collector reports an expired sign-in" + +# A malformed login file is a missing credential, not a crash. +for malformed in '[1,2]' '"nope"' '{"providers":[]}' '{"providers":{"meta":7}}'; do + printf '%s' "$malformed" >"$TEST_HOME/.config/muse/auth.json" + garbled=$(run_collector --force) || + fail "Muse collector survives a malformed login file" "$malformed" + [[ $(jq -c '[.limits,.usageStatusText]' <<<"$garbled") == '[[],""]' ]] || + fail "Muse collector survives a malformed login file" "$malformed -> $garbled" +done +pass "Muse collector survives a malformed login file" +printf '{"providers":{"meta":{"access_token":"test-token"}}}' >"$TEST_HOME/.config/muse/auth.json" + +# The override must win over a different login at the default location. +export MUSE_AUTH_PATH="$TEST_HOME/custom auth.json" +printf '{"providers":{"meta":{"access_token":"test-token"}}}' >"$MUSE_AUTH_PATH" +printf '{"providers":{"meta":{"access_token":"default-token"}}}' >"$TEST_HOME/.config/muse/auth.json" +set_payload '{"body":{"subs_tier_name":"T","subs_usage":{"window":{"used_percent":5,"window_duration_mins":300}}}}' +custom=$(run_collector --force) +[[ $(cat "$STUB_SEEN_FILE") == "Bearer test-token" ]] || + fail "Muse collector honors MUSE_AUTH_PATH over the default login" "$(cat "$STUB_SEEN_FILE")" +pass "Muse collector honors MUSE_AUTH_PATH over the default login" +unset MUSE_AUTH_PATH diff --git a/test/shell.d/agent-usage-muse-scanner-test.sh b/test/shell.d/agent-usage-muse-scanner-test.sh new file mode 100755 index 00000000000..98de59d40e6 --- /dev/null +++ b/test/shell.d/agent-usage-muse-scanner-test.sh @@ -0,0 +1,350 @@ +#!/bin/bash + +set -euo pipefail + +source "$(dirname "$0")/base-test.sh" + +require_command jq +require_command python3 +require_command sqlite3 +PYTHON=$(command -p -v python3) + +TEST_HOME=$(mktemp -d) +trap 'rm -rf "$TEST_HOME"' EXIT + +export HOME="$TEST_HOME" +export XDG_CACHE_HOME="$TEST_HOME/.cache" +export XDG_CONFIG_HOME="$TEST_HOME/.config" +export XDG_DATA_HOME="$TEST_HOME/.local/share" +export MUSE_DATA_DIR="$TEST_HOME/muse-data" +unset META_API_KEY META_API_BASE_URL MUSE_AUTH_PATH + +now_s=$(date +%s) +now_us=$((now_s * 1000000)) +now_ms=$((now_s * 1000)) +old_s=$((now_s - 8 * 86400)) +old_us=$((old_s * 1000000)) +old_ms=$((old_s * 1000)) +old_date=$(date -d "@$old_s" +%F) + +# Native muse sessions: one main session with a model_completed event, one +# subagent event under the same session, and one session from 8 days ago. +mkdir -p "$MUSE_DATA_DIR/sessions/sess-A/subagent/sub-1" "$MUSE_DATA_DIR/sessions/sess-B" +cat >"$MUSE_DATA_DIR/sessions/sess-A/session.jsonl" <"$MUSE_DATA_DIR/sessions/sess-A/subagent/sub-1/session.jsonl" <"$MUSE_DATA_DIR/sessions/sess-B/session.jsonl" <"$TEST_HOME/outside/evil/session.jsonl" <"$MUSE_DATA_DIR/sessions/sess-huge/session.jsonl" <"$MUSE_DATA_DIR/sessions/sess-huge-metadata/session.jsonl" <"$MUSE_DATA_DIR/sessions/sess-negative/session.jsonl" <>"$MUSE_DATA_DIR/sessions/sess-negative/session.jsonl" <"$xdg_home/data/muse/sessions/xdg-session/session.jsonl" <"$ev_home/muse-data/sessions/sess-events/session.jsonl" <"$corrupt_home/muse-data/sessions/sess-native/session.jsonl" <"$corrupt_home/data/opencode/opencode.db" +corrupt_env=(env HOME="$corrupt_home" XDG_DATA_HOME="$corrupt_home/data" XDG_CACHE_HOME="$corrupt_home/cache" MUSE_DATA_DIR="$corrupt_home/muse-data") +result_corrupt=$("${corrupt_env[@]}" "$ROOT/bin/omarchy-agent-usage-muse" --force) +[[ $(jq -r '.todayTotalTokens' <<<"$result_corrupt") == "33" ]] || + fail "Muse collector serves native stats past a corrupt database" "$result_corrupt" +[[ $(jq -r '.ready' <<<"$result_corrupt") == "true" ]] || + fail "Muse collector reports ready past a corrupt database" "$result_corrupt" +shopt -s nullglob +corrupt_cached=("$corrupt_home"/cache/omarchy/agent-usage/muse-scan-*.json) +shopt -u nullglob +(( ${#corrupt_cached[@]} == 0 )) || + fail "Muse collector refuses to cache an incomplete scan" "$result_corrupt" +result_corrupt_again=$("${corrupt_env[@]}" "$ROOT/bin/omarchy-agent-usage-muse") +[[ $(jq -r '.todayTotalTokens' <<<"$result_corrupt_again") == "33" ]] || + fail "Muse collector repeats the scan past a corrupt database" "$result_corrupt_again" +pass "Muse collector survives a corrupt opencode database" + +# Test 16: --limits-only reuses a scan for far longer than a normal run +lo_home="$TEST_HOME/limits-only" +mkdir -p "$lo_home/muse-data/sessions/sess-lo" +cat >"$lo_home/muse-data/sessions/sess-lo/session.jsonl" <>"$lo_home/muse-data/sessions/sess-lo/session.jsonl" <"$odd_home/muse-data/sessions/sess-odd/session.jsonl" < Date: Fri, 2 Oct 2026 06:22:14 -0400 Subject: [PATCH 025/546] Dim stale Muse limits and leave Zen's Muse Spark to OpenCode Cache each sign-in's last good windows and, when a check fails or the sign-in is refused, show them marked limitsStale with their limitsFetchedAt so the panel dims them. Without a sign-in the record names no plan, so a machine nobody uses Muse on gives a record the panel skips. Muse Spark through the OpenCode Zen gateway (providerID "opencode") is Zen usage the opencode record already counts, so only the meta provider counts here. The login file is only ever read; a test now holds that. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-muse | 118 +++++++++--------- shell/plugins/agents/README.md | 13 +- test/shell.d/agent-usage-muse-limits-test.sh | 35 +++++- test/shell.d/agent-usage-muse-scanner-test.sh | 42 ++++--- 4 files changed, 124 insertions(+), 84 deletions(-) diff --git a/bin/omarchy-agent-usage-muse b/bin/omarchy-agent-usage-muse index d5c264484af..0a90cb44eb1 100755 --- a/bin/omarchy-agent-usage-muse +++ b/bin/omarchy-agent-usage-muse @@ -8,9 +8,9 @@ Local metering (always available, no network): token counts, prompt and session counts, per-day and per-model breakdowns. Sources (all local): - native `muse` CLI sessions in $XDG_DATA_HOME/muse/sessions, defaulting to ~/.local/share/muse/sessions (`model_completed` events carry usage), - - opencode sessions with providerID == "meta" (any Muse Spark model), - - opencode sessions with providerID == "opencode" whose modelID contains - "muse-spark" (e.g. muse-spark-1.3-contributor-free via the opencode proxy). + - opencode sessions with providerID == "meta" (any Muse Spark model). + Muse Spark served through OpenCode Zen (providerID == "opencode") is + Zen usage, and the opencode record counts it. Subscription windows (additive, authoritative only): when the CLI is signed in via `muse login`, the collector probes the client's own key endpoint @@ -19,7 +19,9 @@ in via `muse login`, the collector probes the client's own key endpoint session (5-hour) + weekly percents with reset times plus `subs_tier_name` — the same figures the TUI's `/usage` overlay shows. Minting is idempotent so polling is safe; the token travels only in the Authorization header and -neither it nor the returned API key is persisted. Pay-as-you-go accounts +neither it nor the returned API key is persisted, and the login file is only +ever read, never refreshed or rewritten. A failed check keeps the last good +windows, marked stale. Pay-as-you-go accounts and active plans that return no `subs_usage` (e.g. Muse Code Power Usage) have no subscription windows, so limits stay empty — tier plus local metering still show. No estimated meters from manual caps: without an @@ -28,6 +30,7 @@ only ever reads the JSON this prints. """ import argparse +import base64 import fcntl import hashlib import json @@ -323,7 +326,7 @@ def scan_native_muse_sessions(): def scan_opencode_sessions(): - """opencode rows for meta provider Muse models + opencode-proxied muse-spark. + """opencode rows for meta provider Muse models. Returns True when the scan ran to completion (safe to cache). """ @@ -339,14 +342,13 @@ def scan_opencode_sessions(): # LIKE gates are pure acceleration (see the codex collector for the full # rationale); the json_extract gates match the Python filter below, which # stays the authority for the rows that reach it. The provider gate keeps - # assistant rows from other providers — including ones that merely mention - # muse-spark in their text — out of the Python parse. + # assistant rows from other providers out of the Python parse. query = ( "SELECT session_id, data FROM message" " WHERE data LIKE '%\"role\"%:%\"assistant\"%'" - " AND (data LIKE '%\"providerID\"%:%\"meta\"%' OR data LIKE '%muse-spark%')" + " AND data LIKE '%\"providerID\"%:%\"meta\"%'" " 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 IN ('meta', 'opencode')" + " AND CASE WHEN json_valid(data) THEN json_extract(data, '$.providerID') END = 'meta'" ) for session_id, raw in conn.execute(query): try: @@ -357,13 +359,10 @@ def scan_opencode_sessions(): model = model_name(entry.get("modelID")) if model is None: continue - if provider == "meta": - # The meta provider only serves Muse Spark models, so any model on - # it counts (muse-spark 1.1/1.2/1.3, ...). - pass - elif provider == "opencode" and "muse-spark" in model.lower(): - pass # opencode-proxied Muse Spark (e.g. 1.3-contributor-free) - else: + # The meta provider only serves Muse Spark models, so any model on it + # counts (muse-spark 1.1/1.2/1.3, ...). Muse Spark through the Zen + # gateway is OpenCode usage, so the opencode record keeps it. + if provider != "meta": continue tokens = entry.get("tokens") if not isinstance(tokens, dict): @@ -555,50 +554,63 @@ def usable_cached_limits(cached: dict) -> list: return [entry for entry in entries if isinstance(entry, dict) and limit_window_open(entry, now)] +# Keyed by who the token signs in as (the JWT's subject when it is one), so a +# machine signed in to someone else never shows the last one's windows, and a +# refreshed token keeps the same cache. +def limits_cache_path(access_token: str) -> Path: + subject = "" + parts = access_token.split(".") + if len(parts) == 3: + try: + claims = json.loads(base64.urlsafe_b64decode(parts[1] + "=" * (-len(parts[1]) % 4))) + subject = str(claims.get("sub") or "") if isinstance(claims, dict) else "" + except Exception: + subject = "" + digest = hashlib.sha256((subject or access_token).encode("utf-8")).hexdigest()[:16] + return cache_root() / f"muse-limits-{digest}.json" + + def collect_limits(force: bool) -> dict: - result = {"limits": [], "tierLabel": "Account", "usageStatusText": "", "authHelpText": ""} + result = {"limits": [], "tierLabel": "", "usageStatusText": "", "authHelpText": "", "stale": False, "fetchedAtMs": 0} + + access_token = read_access_token() + if access_token == "": + # No login: PAYG or signed-out. Never estimated, never an auth nag — + # local metering is valid on its own. + return result probe_cache = None + cached = {} try: - probe_cache = cache_root() / "muse-limits.json" + probe_cache = limits_cache_path(access_token) cached = json.loads(probe_cache.read_text(encoding="utf-8")) if probe_cache.exists() else {} except Exception as exc: print(f"omarchy-agent-usage-muse: could not read limits cache ({exc})", file=sys.stderr) - cached = {} if not isinstance(cached, dict): cached = {} fallback = usable_cached_limits(cached) - fallback_tier = str(cached.get("tierLabel") or "") + fetched_at = number(cached.get("fetchedAtMs")) + result["tierLabel"] = str(cached.get("tierLabel") or "") or "Account" - access_token = read_access_token() - if access_token == "": - # No login: PAYG or signed-out. Never estimated, never an auth nag — - # local metering is valid on its own. Reuse unexpired cached windows if - # a previous subscription probe left any. - if fallback: - result["limits"] = fallback - if fallback_tier: - result["tierLabel"] = fallback_tier - return result + # The last good windows, dimmed as stale, with when they were measured. + def keep_cached(): + result["limits"] = fallback + result["stale"] = True + result["fetchedAtMs"] = fetched_at if fallback else 0 - fetched_at = number(cached.get("fetchedAtMs")) / 1000 - if (fallback or fallback_tier) and not force and time.time() - fetched_at < PROBE_MIN_INTERVAL_SECONDS: + if (fallback or cached.get("tierLabel")) and not force and 0 <= time.time() - fetched_at / 1000 < PROBE_MIN_INTERVAL_SECONDS: result["limits"] = fallback - if fallback_tier: - result["tierLabel"] = fallback_tier + result["fetchedAtMs"] = fetched_at if fallback else 0 return result probe = probe_key(access_token) if probe["ok"]: result["limits"] = probe["limits"] - if probe.get("tierLabel"): - result["tierLabel"] = probe["tierLabel"] - elif fallback_tier: - result["tierLabel"] = fallback_tier + result["tierLabel"] = probe.get("tierLabel") or result["tierLabel"] + result["fetchedAtMs"] = round(time.time() * 1000) if probe_cache is not None: try: - write_json(probe_cache, {"fetchedAtMs": round(time.time() * 1000), - "limits": probe["limits"], "tierLabel": result["tierLabel"]}) + write_json(probe_cache, {"fetchedAtMs": result["fetchedAtMs"], "limits": probe["limits"], "tierLabel": result["tierLabel"]}) except Exception as exc: print(f"omarchy-agent-usage-muse: could not write limits cache ({exc})", file=sys.stderr) return result @@ -608,32 +620,22 @@ def collect_limits(force: bool) -> dict: if probe.get("payg"): # Authenticated PAYG: no subscription windows exist. Keep local metering # only, silent — not an error, not expired. + result["tierLabel"] = "Account" return result if probe.get("silent"): # Active plan with no windows: same silence as PAYG, but keep the plan - # name. Stale cached windows are dropped — they belong to a different - # plan shape, not this one. - if probe.get("tierLabel"): - result["tierLabel"] = probe["tierLabel"] + # name. Cached windows are dropped — they belong to a different plan + # shape, not this one. + result["tierLabel"] = probe.get("tierLabel") or result["tierLabel"] return result + keep_cached() if probe.get("expired"): result["usageStatusText"] = "Sign-in expired" result["authHelpText"] = probe["helpText"] - result["limits"] = fallback - if fallback_tier: - result["tierLabel"] = fallback_tier - return result - if fallback: - result["limits"] = fallback - if fallback_tier: - result["tierLabel"] = fallback_tier - elif probe.get("tierLabel"): - result["tierLabel"] = probe["tierLabel"] - else: + elif not fallback: result["usageStatusText"] = "Muse limits unavailable" result["authHelpText"] = probe["helpText"] - if probe.get("tierLabel"): - result["tierLabel"] = probe["tierLabel"] + result["tierLabel"] = probe.get("tierLabel") or result["tierLabel"] return result @@ -780,6 +782,8 @@ def main(): "modelUsage": stats.get("modelUsage") or {}, "tierLabel": limits["tierLabel"], "limits": limits["limits"], + "limitsStale": limits["stale"], + "limitsFetchedAt": limits["fetchedAtMs"], "usageStatusText": limits["usageStatusText"], "authHelpText": limits["authHelpText"], } diff --git a/shell/plugins/agents/README.md b/shell/plugins/agents/README.md index 0df0f8bab18..8b691186c52 100644 --- a/shell/plugins/agents/README.md +++ b/shell/plugins/agents/README.md @@ -77,7 +77,7 @@ light surfaces — and the bar glyph stands in when there is none. | `opencode` | OpenCode Go's usage endpoint (rolling + weekly + monthly), with an API key or the Console session | pi and omp sessions on the Zen and Go providers, plus opencode's own `message` and `session_message` stores | | `fireworks` | Estimated prepaid balance: configured funding minus rated account costs | Fireworks billing API, grouped by day and model for the last 30 days | | `copilot` | Account-wide AI credit allowance (GitHub quota endpoint, or estimated locally) | the Copilot CLI session store's `assistant_usage_events`, with `~/.copilot/session-state` transcripts as fallback | -| `muse` | The client's key endpoint (session + weekly percents), else local metering only — never estimated | native `muse` CLI session files plus opencode sessions on the `meta` provider or `muse-spark` models | +| `muse` | The client's key endpoint (session + weekly percents), else local metering only — never estimated | native `muse` CLI session files plus opencode sessions on the `meta` provider | When `~/.local/state/omarchy/agents/accounts/.json` registers more than one account, the `claude`, `codex`, and `grok` records @@ -207,15 +207,18 @@ only adds the meter and the spent-of-funded line under the real figure. Muse reads native CLI sessions from `$XDG_DATA_HOME/muse/sessions` (defaulting to `~/.local/share/muse/sessions`, and overridable via -`MUSE_DATA_DIR`) plus opencode sessions on the `meta` provider or -`muse-spark` models, and the OAuth token from `MUSE_AUTH_PATH` when set, -otherwise from `~/.config/muse/auth.json` (which `muse login` creates). +`MUSE_DATA_DIR`) plus opencode sessions on the `meta` provider. Muse Spark +through the OpenCode Zen gateway counts in the `opencode` record instead. The +OAuth token comes from `MUSE_AUTH_PATH` when set, otherwise from +`~/.config/muse/auth.json` (which `muse login` creates); the file is only +ever read, never refreshed or rewritten. The collector asks the client's own key endpoint for the subscription's session and weekly percents with reset times — the same figures the TUI's `/usage` overlay shows. Minting is idempotent (the same Model API key comes back every call), so polling is safe; the token travels only in the -`Authorization` header and neither it nor the returned key is persisted. +`Authorization` header and neither it nor the returned key is persisted. A +failed check keeps the last good windows, dimmed, until they reset. Pay-as-you-go accounts have no subscription windows, so the tab shows measured local usage with no meters — never estimated, never zeroed. An active plan that returns no `subs_usage` (Muse Code Power Usage answers this diff --git a/test/shell.d/agent-usage-muse-limits-test.sh b/test/shell.d/agent-usage-muse-limits-test.sh index dadf3f73d1c..cf2e821049d 100755 --- a/test/shell.d/agent-usage-muse-limits-test.sh +++ b/test/shell.d/agent-usage-muse-limits-test.sh @@ -16,6 +16,12 @@ mkdir -p "$TEST_HOME/.config/muse" "$TEST_HOME/.cache" cat >"$TEST_HOME/.config/muse/auth.json" <<'EOF' {"schema_version":1,"providers":{"meta":{"access_token":"test-token","api_base_url":"https://api.meta.ai/v1"}}} EOF +touch -d '2026-01-01 00:00:00' "$TEST_HOME/.config/muse/auth.json" +auth_state() { + stat -c '%Y %s' "$TEST_HOME/.config/muse/auth.json" + sha256sum <"$TEST_HOME/.config/muse/auth.json" +} +auth_before=$(auth_state) # A stub key service: records the Authorization header it was called with # and answers from a canned payload file. The collector must send the OAuth @@ -81,18 +87,18 @@ set_payload() { : >"$STUB_SEEN_FILE" # Each stub case starts without a previous probe's cache, so fallback # behavior never leaks across cases. - rm -f "$TEST_HOME/.cache/omarchy/agent-usage/muse-limits.json" + rm -f "$TEST_HOME"/.cache/omarchy/agent-usage/muse-limits-*.json } # Subscriber with usage: window + weekly percents, resets, tier. -set_payload '{"body":{"subs_tier_name":"Power Usage","subs_usage":{"window":{"used_percent":8,"window_duration_mins":300,"resets_at":1788900190},"weekly":{"used_percent":15,"resets_at":1789344000}}}}' +set_payload '{"body":{"subs_tier_name":"Power Usage","subs_usage":{"window":{"used_percent":8,"window_duration_mins":300,"resets_at":4102444800},"weekly":{"used_percent":15,"resets_at":4103049600}}}}' result=$(run_collector --force) [[ $(jq -c '[.limits[]|{label,percent}]' <<<"$result") == '[{"label":"Session (5-hour)","percent":0.08},{"label":"Weekly (7-day)","percent":0.15}]' ]] || fail "Muse collector reports window and weekly limits from the key endpoint" "$result" pass "Muse collector reports window and weekly limits from the key endpoint" -[[ $(jq -r '.limits[0].resetsAt' <<<"$result") == "2026-09-08T20:43:10+00:00" ]] || +[[ $(jq -r '.limits[0].resetsAt' <<<"$result") == "2100-01-01T00:00:00+00:00" ]] || fail "Muse collector converts window reset times to ISO" "$result" pass "Muse collector converts window reset times to ISO" @@ -121,6 +127,29 @@ pass "Muse collector reuses a recent probe result" fail "Muse collector reuses the cached tier" "$second" pass "Muse collector reuses the cached tier" +[[ $(jq -c '{stale: .limitsStale, stamped: (.limitsFetchedAt > 0)}' <<<"$result") == '{"stale":false,"stamped":true}' ]] || + fail "Muse collector stamps fresh limits" "$result" +pass "Muse collector stamps fresh limits" + +# A failed check keeps the last good windows, marked stale with when they +# were measured; a rejected sign-in does too, and says so. +printf '{"status":500}' >"$STUB_PAYLOAD_FILE" +down=$(run_collector --force) +[[ $(jq -c --argjson fresh "$result" '{same: (.limits == $fresh.limits), stale: .limitsStale, stamp: (.limitsFetchedAt == $fresh.limitsFetchedAt), tier: .tierLabel, status: .usageStatusText}' <<<"$down") == '{"same":true,"stale":true,"stamp":true,"tier":"Power Usage","status":""}' ]] || + fail "Muse collector keeps the last good limits, stale, after a failed check" "$down" +pass "Muse collector keeps the last good limits, stale, after a failed check" + +printf '{"status":401}' >"$STUB_PAYLOAD_FILE" +refused=$(run_collector --force) +[[ $(jq -c '{kept: (.limits | length), stale: .limitsStale, status: .usageStatusText}' <<<"$refused") == '{"kept":2,"stale":true,"status":"Sign-in expired"}' ]] || + fail "Muse collector keeps the last limits through a rejected sign-in" "$refused" +pass "Muse collector keeps the last limits through a rejected sign-in" + +# The login file belongs to the Muse CLI: every run so far only read it. +[[ $(auth_state) == "$auth_before" && $(ls -A "$TEST_HOME/.config/muse") == "auth.json" ]] || + fail "Muse collector never writes the login file" "$(ls -lA --time-style=+%s "$TEST_HOME/.config/muse")" +pass "Muse collector never writes the login file" + # Pay-as-you-go: key mints fine but there is no subscription window to # meter. Silent local metering — not an error, not expired. set_payload '{"body":{"api_key":"minted","is_subs_active":false,"subs_tier_name":null}}' diff --git a/test/shell.d/agent-usage-muse-scanner-test.sh b/test/shell.d/agent-usage-muse-scanner-test.sh index 98de59d40e6..7e226067fe0 100755 --- a/test/shell.d/agent-usage-muse-scanner-test.sh +++ b/test/shell.d/agent-usage-muse-scanner-test.sh @@ -41,8 +41,9 @@ cat >"$MUSE_DATA_DIR/sessions/sess-B/session.jsonl" <"$MUSE_DATA_DIR/sessions/sess-huge-metadata/session.jsonl" < Date: Fri, 2 Oct 2026 06:29:41 -0400 Subject: [PATCH 026/546] 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 027/546] 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 028/546] 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 029/546] 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" From 33f8eaa265c02bbb643eeac93f04989ff9889665 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Fri, 2 Oct 2026 06:41:47 -0400 Subject: [PATCH 030/546] Keep Cursor's carried numbers to their own account and day Runs that fetch no fresh tokens carried the token sections of whatever cursor.json held, even after signing in as someone else, and kept yesterday's today figures past midnight. Carry this account's own dated scan instead: a scan from an earlier day moves onto this week's dates with its today figures cleared. Kept limits drop once their billing cycle ends, and each day's window now ends at the next local midnight, so days that change for DST are 23 or 25 hours long. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-cursor | 94 ++++++++++++++++--------- shell/plugins/agents/README.md | 7 +- test/shell.d/agent-usage-cursor-test.sh | 74 +++++++++++++++++-- 3 files changed, 136 insertions(+), 39 deletions(-) diff --git a/bin/omarchy-agent-usage-cursor b/bin/omarchy-agent-usage-cursor index 3611619bd7b..cd697dff0ad 100755 --- a/bin/omarchy-agent-usage-cursor +++ b/bin/omarchy-agent-usage-cursor @@ -59,14 +59,6 @@ def fraction(value: Any) -> float: return 0.0 -def state_home() -> Path: - return Path(os.environ.get("XDG_STATE_HOME") or (Path.home() / ".local" / "state")) - - -def record_path() -> Path: - return state_home() / "omarchy" / "agents" / "usage" / f"{AGENT_ID}.json" - - def default_auth_path() -> Path: config_home = Path(os.environ.get("XDG_CONFIG_HOME") or (Path.home() / ".config")) return config_home / "cursor" / "auth.json" @@ -226,8 +218,10 @@ def day_windows(today: date, count: int) -> list[tuple[str, int, int]]: windows = [] for offset in range(count - 1, -1, -1): day = today - timedelta(days=offset) + # Each midnight converts on its own: a day that springs forward or falls + # back is 23 or 25 hours long, never a fixed 24. start = datetime.combine(day, time.min).astimezone() - end = start + timedelta(days=1) + end = datetime.combine(day + timedelta(days=1), time.min).astimezone() windows.append((day.isoformat(), int(start.timestamp() * 1000), int(end.timestamp() * 1000))) return windows @@ -277,26 +271,41 @@ def summarize_usage(days: list[tuple[str, dict[str, dict[str, int]]]]) -> dict[s } -def carried_stats() -> dict[str, Any]: - """--limits-only promises fresh limits only, so the token sections stand.""" +def stats_cache(key: str) -> Path: + return cache_root() / f"cursor-stats-{key}.json" + + +def carried_stats(key: str, today: date) -> dict[str, Any]: + """The last scan of this account, for runs that fetch no fresh tokens. + + Another account's scan never stands in for this one's. A scan from an + earlier day is moved onto this week's dates, and its today figures are + cleared, since nothing in it happened today. + """ stats = empty_stats() - try: - previous = json.loads(record_path().read_text()) - except (OSError, ValueError): - return stats + cached = read_json(stats_cache(key)) + previous = cached.get("stats") if not isinstance(previous, dict): return stats - for key in stats: - if key in previous: - stats[key] = previous[key] + for name in stats: + if name in previous: + stats[name] = previous[name] + if cached.get("scanDate") != today.isoformat(): + counts = {str(row.get("date")): row.get("messageCount", 0) for row in stats["recentDays"] if isinstance(row, dict)} + week = [window[0] for window in day_windows(today, RECENT_DAYS)] + stats["recentDays"] = [{"date": day, "messageCount": counts.get(day, 0)} for day in week] + stats["todayTotalTokens"] = 0 + stats["todayTokensByModel"] = {} return stats -def collect_stats(client: CursorClient, today: date) -> dict[str, Any]: +def collect_stats(client: CursorClient, today: date, key: str) -> dict[str, Any]: windows = day_windows(today, RECENT_DAYS) with ThreadPoolExecutor(max_workers=4) as pool: rows = list(pool.map(lambda window: client.aggregations(window[1], window[2]), windows)) - return summarize_usage([(window[0], summarize_day(day_rows)) for window, day_rows in zip(windows, rows)]) + stats = summarize_usage([(window[0], summarize_day(day_rows)) for window, day_rows in zip(windows, rows)]) + write_json(stats_cache(key), {"scanDate": today.isoformat(), "stats": stats}) + return stats def summarize_limits(period: dict[str, Any]) -> list[dict[str, Any]]: @@ -358,10 +367,10 @@ def write_json(path: Path, payload: dict[str, Any]) -> None: raise -# Keyed by who the token signs in as (the JWT's subject), so a machine signed -# in to someone else never shows the last one's limits, and a refreshed token -# keeps the same cache. An opaque API key keys by its own hash. -def limits_cache(token: str) -> Path: +# Who the token signs in as (the JWT's subject), so a machine signed in to +# someone else never shows the last one's limits or tokens, and a refreshed +# token keeps the same caches. An opaque API key keys by its own hash. +def account_key(token: str) -> str: subject = "" parts = token.split(".") if len(parts) == 3: @@ -370,10 +379,29 @@ def limits_cache(token: str) -> Path: subject = str(claims.get("sub") or "") if isinstance(claims, dict) else "" except (ValueError, UnicodeDecodeError): subject = "" - key = hashlib.sha256((subject or token).encode("utf-8")).hexdigest()[:16] + return hashlib.sha256((subject or token).encode("utf-8")).hexdigest()[:16] + + +def limits_cache(key: str) -> Path: return cache_root() / f"cursor-limits-{key}.json" +# A kept limit is only good until its billing cycle ends. +def open_limits(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_at = datetime.fromisoformat(str(limit.get("resetsAt") or "")) + except ValueError: + resets_at = None + if resets_at is None or resets_at.tzinfo is None or resets_at > now: + kept.append(limit) + return kept + + def now_ms() -> int: return round(datetime.now(timezone.utc).timestamp() * 1000) @@ -384,14 +412,16 @@ def scan(api_base_url: str, auth_path: Path, limits_only: bool, state_db: Path | return base_record(usageStatusText="Cursor unavailable", authHelpText=AUTH_HELP) client = CursorClient(token, api_base_url) - cache = limits_cache(token) + key = account_key(token) + cache = limits_cache(key) + today = datetime.now().astimezone().date() try: period = client.period() except CursorError as error: # A failed check keeps the last good limits, marked stale so the panel # dims them, and the week of tokens already on disk. cached = read_json(cache) - limits = cached.get("limits") or [] + limits = open_limits(cached.get("limits")) status = "Waiting for auth" if error.auth else ("" if limits else "Cursor limits unavailable") record = base_record( ready=True, @@ -400,9 +430,9 @@ def scan(api_base_url: str, auth_path: Path, limits_only: bool, state_db: Path | authHelpText=str(error), limits=limits, limitsStale=True, - limitsFetchedAt=number(cached.get("fetchedAtMs")), + limitsFetchedAt=number(cached.get("fetchedAtMs")) if limits else 0, ) - record.update(carried_stats()) + record.update(carried_stats(key, today)) return record record = base_record( @@ -414,12 +444,12 @@ def scan(api_base_url: str, auth_path: Path, limits_only: bool, state_db: Path | ) write_json(cache, {"fetchedAtMs": record["limitsFetchedAt"], "tierLabel": record["tierLabel"], "limits": record["limits"]}) if limits_only: - record.update(carried_stats()) + record.update(carried_stats(key, today)) else: try: - record.update(collect_stats(client, datetime.now().astimezone().date())) + record.update(collect_stats(client, today, key)) except CursorError: - record.update(carried_stats()) + record.update(carried_stats(key, today)) return record diff --git a/shell/plugins/agents/README.md b/shell/plugins/agents/README.md index 8b691186c52..14ebd063531 100644 --- a/shell/plugins/agents/README.md +++ b/shell/plugins/agents/README.md @@ -113,9 +113,10 @@ editor's own sign-in in `~/.config/Cursor/User/globalStorage/state.vscdb` Cursor's session files on disk carry no token counts, so this is the one collector with nothing to fall back on: without a credential the record stays empty and the panel skips it rather than showing a week of zeros. A failed -check keeps the last good meters, dimmed, and the week already on disk; a -refused sign-in also says so under the name. Tokens by day and by -model cover the last seven days, one aggregation call per day. +check keeps the account's last good meters, dimmed, until its billing cycle +ends, and its last week of tokens (today cleared once the day turns over); a +refused sign-in also says so under the name. Tokens by day and by model cover +the last seven local days, one aggregation call per day. The meters are Cursor's own percentages — included total, auto-bucket models, and named models — so they always agree with the Cursor dashboard. Spend is diff --git a/test/shell.d/agent-usage-cursor-test.sh b/test/shell.d/agent-usage-cursor-test.sh index ff82d4a1b4c..d17166c1803 100755 --- a/test/shell.d/agent-usage-cursor-test.sh +++ b/test/shell.d/agent-usage-cursor-test.sh @@ -102,12 +102,42 @@ record = collector.scan("https://api.example", Path(os.environ["XDG_CONFIG_HOME" windows = collector.day_windows(today, collector.RECENT_DAYS) first_start = datetime.fromtimestamp(windows[0][1] / 1000).astimezone() -# --limits-only keeps the token sections of the record already on disk. -record_path = collector.record_path() -record_path.parent.mkdir(parents=True, exist_ok=True) -record_path.write_text(json.dumps(record)) +# --limits-only keeps the token sections of this account's last scan. carried = collector.scan("https://api.example", Path(os.environ["XDG_CONFIG_HOME"]) / "cursor" / "auth.json", True) +# Another account's scan never stands in: signed in as someone else with the +# dashboard down, there are no tokens to show. +class DownForOther(StubClient): + def call(self, method, payload): + raise collector.CursorError("Cursor API unreachable") + + +collector.CursorClient = DownForOther +os.environ["CURSOR_API_KEY"] = "someone_else" +other = collector.scan("https://api.example", Path(os.environ["XDG_CONFIG_HOME"]) / "cursor" / "auth.json", False) +os.environ.pop("CURSOR_API_KEY") +collector.CursorClient = StubClient + +# A scan from yesterday moves onto this week's dates, and its today figures +# are cleared, since none of it happened today. +key = collector.account_key("cursor_test") +cache = collector.stats_cache(key) +scan = json.loads(cache.read_text()) +yesterday_scan = dict(scan, scanDate=yesterday.isoformat()) +yesterday_scan["stats"] = dict(scan["stats"], recentDays=[{"date": (today - timedelta(days=offset)).isoformat(), "messageCount": offset} for offset in range(7, 0, -1)]) +cache.write_text(json.dumps(yesterday_scan)) +overnight = collector.carried_stats(key, today) +cache.write_text(json.dumps(scan)) + +# A kept limit whose billing cycle has ended is dropped, not shown stale. +limits_cache = collector.limits_cache(key) +good = json.loads(limits_cache.read_text()) +limits_cache.write_text(json.dumps(dict(good, limits=[dict(good["limits"][0], resetsAt="2020-01-01T00:00:00+00:00")] + good["limits"][1:]))) +collector.CursorClient = DownForOther +lapsed = collector.scan("https://api.example", Path(os.environ["XDG_CONFIG_HOME"]) / "cursor" / "auth.json", False) +limits_cache.write_text(json.dumps(good)) +collector.CursorClient = StubClient + # A failed check keeps the last good limits and the week on disk, marked stale. class DownClient(StubClient): @@ -126,6 +156,9 @@ collector.CursorClient = RefusedClient refused = collector.scan("https://api.example", Path(os.environ["XDG_CONFIG_HOME"]) / "cursor" / "auth.json", False) print(json.dumps({ + "other": other, + "overnight": overnight, + "lapsed": [limit["label"] for limit in lapsed["limits"]], "down": down, "refused": refused, "record": record, @@ -185,6 +218,39 @@ pass "Cursor collector keeps the last good limits, stale, after a failed check" fail "Cursor collector says a refused sign-in needs attention" "$result" pass "Cursor collector says a refused sign-in needs attention" +[[ $(jq -c '.other | {limits, tierLabel, todayTotalTokens, recentDays, modelUsage}' <<<"$result") == '{"limits":[],"tierLabel":"","todayTotalTokens":0,"recentDays":[],"modelUsage":{}}' ]] || + fail "Cursor collector never carries another account's numbers" "$result" +pass "Cursor collector never carries another account's numbers" + +[[ $(jq -c '.overnight | {todayTotalTokens, todayTokensByModel, counts: [.recentDays[].messageCount], last: .recentDays[6].date}' <<<"$result") == "{\"todayTotalTokens\":0,\"todayTokensByModel\":{},\"counts\":[6,5,4,3,2,1,0],\"last\":\"$(TZ=UTC date +%F)\"}" ]] || + fail "Cursor collector moves a scan from yesterday onto today's week" "$result" +pass "Cursor collector moves a scan from yesterday onto today's week" + +[[ $(jq -c '.lapsed' <<<"$result") == '["Auto models","Named models","On-demand $9.51 / $10.00"]' ]] || + fail "Cursor collector drops a kept limit whose cycle ended" "$result" +pass "Cursor collector drops a kept limit whose cycle ended" + +# Days that spring forward or fall back are 23 or 25 hours long. +dst=$(TZ=America/New_York python3 - "$ROOT/bin/omarchy-agent-usage-cursor" <<'PY' +import importlib.machinery +import importlib.util +import sys +from datetime import date + +loader = importlib.machinery.SourceFileLoader("cursor_collector", sys.argv[1]) +spec = importlib.util.spec_from_loader(loader.name, loader) +collector = importlib.util.module_from_spec(spec) +loader.exec_module(collector) + +spring = collector.day_windows(date(2026, 3, 9), 2) +fall = collector.day_windows(date(2026, 11, 1), 1) +hours = [(end - start) / 3_600_000 for _, start, end in spring + fall] +print(" ".join(str(int(h)) for h in hours), spring[0][2] == spring[1][1]) +PY +) +[[ $dst == "23 24 25 True" ]] || fail "Cursor collector sizes daily windows across DST" "$dst" +pass "Cursor collector sizes daily windows across DST" + # Someone who only uses the Cursor editor has no CLI auth file; the editor's # own sign-in in state.vscdb stands in, read without touching the database. ide=$(python3 - "$ROOT/bin/omarchy-agent-usage-cursor" "$TEST_HOME/ide" "$TEST_HOME/.config/cursor/auth.json" <<'PY' From 24136aebcaa9b52ba5162787cd0103c38c6af008 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Fri, 2 Oct 2026 06:42:30 -0400 Subject: [PATCH 031/546] Drop a kept Copilot allowance once its window resets A failed quota check after the allowance reset kept last month's usage and gave it the new month's reset date. A kept answer now only stands until its own reset (or, without one, through the UTC month it was fetched in). Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-copilot | 14 +++++++++++++- test/shell.d/agent-usage-copilot-test.sh | 11 ++++++++++- 2 files changed, 23 insertions(+), 2 deletions(-) diff --git a/bin/omarchy-agent-usage-copilot b/bin/omarchy-agent-usage-copilot index de8467c1e94..d306921c096 100755 --- a/bin/omarchy-agent-usage-copilot +++ b/bin/omarchy-agent-usage-copilot @@ -798,12 +798,24 @@ def fetch_remote_quota() -> dict[str, Any]: return quota cached = read_fresh_json(cache, float("inf")) or {} - if cached.get("available"): + if cached.get("available") and quota_window_open(cached): cached.update(stale=True, failure=quota.get("reason", "")) return cached return quota +# A kept answer describes its own allowance window and nothing after it: once +# that window resets, its usage is gone. Without a reset date, the window is +# the calendar month (UTC) it was fetched in. +def quota_window_open(cached: dict[str, Any]) -> bool: + now = dt.datetime.now(dt.timezone.utc) + resets_at = parse_event_time(str(cached.get("resetsAt") or "")) + if resets_at is not None: + return resets_at > now + fetched = number(cached.get("fetchedAtMs")) / 1000 + return fetched > 0 and utc_month_key(fetched) == now.strftime("%Y-%m") + + # ---------------------------------------------------------------- allowance diff --git a/test/shell.d/agent-usage-copilot-test.sh b/test/shell.d/agent-usage-copilot-test.sh index 4541db2b04e..ca8edb86c51 100755 --- a/test/shell.d/agent-usage-copilot-test.sh +++ b/test/shell.d/agent-usage-copilot-test.sh @@ -182,10 +182,15 @@ collector.urllib.request.urlopen = lambda request, timeout=None: io.BytesIO(json record = run() collector.urllib.request.urlopen = offline stale = run() +# Once the kept answer's window has reset, its usage no longer applies. +for cache in collector.cache_root().glob("copilot-quota-*.json"): + kept = json.loads(cache.read_text()) + cache.write_text(json.dumps(dict(kept, resetsAt="2020-01-01T00:00:00.000Z"))) +lapsed = run() # gh signed in on a machine that never ran the Copilot CLI: nothing is asked. collector.urllib.request.urlopen = unasked clean = run("--copilot-home", os.path.join(os.environ["EMPTY_HOME"], ".copilot")) -print(json.dumps({"record": record, "stale": stale, "clean": clean, "reset": reset.isoformat()})) +print(json.dumps({"record": record, "stale": stale, "lapsed": lapsed, "clean": clean, "reset": reset.isoformat()})) PY ) @@ -197,6 +202,10 @@ pass "Copilot collector stamps the live allowance" fail "Copilot collector keeps the last live allowance, stale, after a failed check" "$result" pass "Copilot collector keeps the last live allowance, stale, after a failed check" +[[ $(jq -c '.lapsed | {label: .limits[0].label, stale: .limitsStale, tierLabel}' <<<"$result") == '{"label":"Monthly allowance (est.)","stale":false,"tierLabel":""}' ]] || + fail "Copilot collector drops a kept allowance once its window resets" "$result" +pass "Copilot collector drops a kept allowance once its window resets" + [[ $(jq -c '.clean | {ready, tierLabel, limits}' <<<"$result") == '{"ready":false,"tierLabel":"","limits":[]}' ]] || fail "Copilot collector doesn't ask GitHub where the CLI never ran" "$result" pass "Copilot collector doesn't ask GitHub where the CLI never ran" From 8dd74d49bd722a711815deb54f636625b1a309c7 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Fri, 2 Oct 2026 06:44:19 -0400 Subject: [PATCH 032/546] Read Muse's OpenCode V2 sessions and show a plan before its first session Muse runs in OpenCode V2 live in session_message, with the provider nested under `model`, and were counted nowhere: the opencode record leaves the meta provider to Muse. Read both stores, deduplicated by message id and leaving out fork copies, from the database OPENCODE_DB names when set. A signed-in plan with no windows and no sessions yet came out not ready, so the panel hid it; a named plan now counts as ready. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-muse | 149 ++++++++++++------ shell/plugins/agents/README.md | 5 +- test/shell.d/agent-usage-muse-limits-test.sh | 6 + test/shell.d/agent-usage-muse-scanner-test.sh | 17 ++ 4 files changed, 126 insertions(+), 51 deletions(-) diff --git a/bin/omarchy-agent-usage-muse b/bin/omarchy-agent-usage-muse index 0a90cb44eb1..34da423a019 100755 --- a/bin/omarchy-agent-usage-muse +++ b/bin/omarchy-agent-usage-muse @@ -325,12 +325,74 @@ def scan_native_muse_sessions(): input_tokens, output_tokens, cache_read, cache_write) +def opencode_db(): + # The database opencode itself reads, as the opencode collector finds it: + # OPENCODE_DB when set (a relative name lives beside the default one), and + # opencode-v2.db only for the builds that wrote nothing else. + root = Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) / "opencode" + configured = os.environ.get("OPENCODE_DB", "").strip() + if configured: + path = Path(configured).expanduser() + return path if path.is_absolute() else root / path + normal = root / "opencode.db" + legacy = root / "opencode-v2.db" + return legacy if not normal.is_file() and legacy.is_file() else normal + + +def table_columns(conn, table): + try: + return {row[1] for row in conn.execute(f"PRAGMA table_info({table})")} + except sqlite3.Error: + return set() + + +def add_opencode_row(message_id, session_id, entry, row_type, created, seen): + # The original `message` table keeps the provider and model at the top + # level; V2's `session_message` nests them under `model` and keeps the role + # in its `type` column, with a compaction's tokens spent like a reply's. + if not isinstance(entry, dict) or (entry.get("role") or row_type) not in ("assistant", "compaction"): + return + nested = entry.get("model") if isinstance(entry.get("model"), dict) else {} + # The meta provider only serves Muse Spark models, so any model on it + # counts (muse-spark 1.1/1.2/1.3, ...). Muse Spark through the Zen gateway + # is OpenCode usage, so the opencode record keeps it. + if str(entry.get("providerID") or nested.get("providerID") or "") != "meta": + return + model = model_name(entry.get("modelID") or nested.get("id") or nested.get("modelID")) + if model is None: + return + tokens = entry.get("tokens") + if not isinstance(tokens, dict): + tokens = {} + cache = tokens.get("cache") + if not isinstance(cache, dict): + cache = {} + # opencode reports input exclusive of cache: keep the token buckets + # separate, no subtraction. + input_tokens = token_number(tokens.get("input")) + output_tokens = token_number(tokens.get("output")) + token_number(tokens.get("reasoning")) + cache_read = token_number(cache.get("read")) + cache_write = token_number(cache.get("write")) + if not (input_tokens or output_tokens or cache_read or cache_write): + return + # A message migrated into both stores counts once. + if message_id in seen: + return + seen.add(message_id) + # A malformed time object must not drop the row's valid tokens: non-dict + # shapes fall back to the row's own time, then today. + stamped = entry.get("time") + day = local_day((stamped.get("created") if isinstance(stamped, dict) else None) or created) + add_usage(day, "opencode:" + str(session_id), model, input_tokens, output_tokens, cache_read, cache_write) + + def scan_opencode_sessions(): - """opencode rows for meta provider Muse models. + """opencode rows for meta provider Muse models, from both of opencode's + message stores, deduplicated by message id. Returns True when the scan ran to completion (safe to cache). """ - db = Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) / "opencode" / "opencode.db" + db = opencode_db() if not db.is_file(): return True try: @@ -339,54 +401,42 @@ def scan_opencode_sessions(): return False try: conn.execute("PRAGMA query_only = ON") + tables = {name for (name,) in conn.execute("SELECT name FROM sqlite_master WHERE type = 'table'")} + seen = set() # LIKE gates are pure acceleration (see the codex collector for the full - # rationale); the json_extract gates match the Python filter below, which - # stays the authority for the rows that reach it. The provider gate keeps - # assistant rows from other providers out of the Python parse. - query = ( - "SELECT session_id, data FROM message" - " WHERE data LIKE '%\"role\"%:%\"assistant\"%'" - " AND data LIKE '%\"providerID\"%:%\"meta\"%'" - " 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 = 'meta'" - ) - for session_id, raw in conn.execute(query): - try: - entry = json.loads(raw) - if not isinstance(entry, dict) or entry.get("role") != "assistant": - continue - provider = str(entry.get("providerID") or "") - model = model_name(entry.get("modelID")) - if model is None: - continue - # The meta provider only serves Muse Spark models, so any model on it - # counts (muse-spark 1.1/1.2/1.3, ...). Muse Spark through the Zen - # gateway is OpenCode usage, so the opencode record keeps it. - if provider != "meta": + # rationale); the Python filter stays the authority for the rows that + # reach it. The provider gate keeps other providers' rows out of the parse. + columns = table_columns(conn, "session_message") if "session_message" in tables else set() + if {"id", "session_id", "type", "data"} <= columns: + created = "m.time_created" if "time_created" in columns else "NULL" + source = "session_message m" + fork_filter = "" + if "time_created" in columns and "session_v2" in tables and {"id", "fork_session_id", "time_created"} <= table_columns(conn, "session_v2"): + # A fork copies its source session's messages under new ids; only + # what came after the fork was spent in it. + source += " LEFT JOIN session_v2 s ON s.id = m.session_id" + fork_filter = " AND (s.fork_session_id IS NULL OR m.time_created >= s.time_created)" + for message_id, session_id, row_type, time_created, raw in conn.execute( + f"SELECT m.id, m.session_id, m.type, {created}, m.data FROM {source}" + " WHERE m.type IN ('assistant', 'compaction')" + " AND m.data LIKE '%\"providerID\"%:%\"meta\"%'" + fork_filter + ): + try: + add_opencode_row(str(message_id), session_id, json.loads(raw), str(row_type or ""), time_created, seen) + except Exception: continue - tokens = entry.get("tokens") - if not isinstance(tokens, dict): - tokens = {} - cache = tokens.get("cache") - if not isinstance(cache, dict): - cache = {} - # opencode reports input exclusive of cache: keep the token buckets - # separate, no subtraction. - input_tokens = token_number(tokens.get("input")) - output_tokens = token_number(tokens.get("output")) + token_number(tokens.get("reasoning")) - cache_read = token_number(cache.get("read")) - cache_write = token_number(cache.get("write")) - if not (input_tokens or output_tokens or cache_read or cache_write): + columns = table_columns(conn, "message") if "message" in tables else set() + if {"session_id", "data"} <= columns: + row_id = "id" if "id" in columns else "rowid" + for message_id, session_id, raw in conn.execute( + f"SELECT {row_id}, session_id, data FROM message" + " WHERE data LIKE '%\"role\"%:%\"assistant\"%'" + " AND data LIKE '%\"providerID\"%:%\"meta\"%'" + ): + try: + add_opencode_row(str(message_id), session_id, json.loads(raw), "", None, seen) + except Exception: continue - # A malformed time object must not drop the row's valid tokens: - # non-dict shapes fall back to today via local_day(None). - stamped = entry.get("time") - created = stamped.get("created") if isinstance(stamped, dict) else None - day = local_day(created) - add_usage(day, "opencode:" + str(session_id), model, - input_tokens, output_tokens, cache_read, cache_write) - except Exception: - continue except sqlite3.Error: return False finally: @@ -646,7 +696,7 @@ def cache_root(): def scan_cache_paths(): - db = Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) / "opencode" / "opencode.db" + db = opencode_db() digest = hashlib.sha1( (str(Path.home()) + "\n" + str(muse_sessions_root()) + "\n" + str(db)).encode("utf-8") ).hexdigest()[:16] @@ -766,7 +816,8 @@ def main(): "id": AGENT_ID, "name": AGENT_NAME, "updatedAt": datetime.now(timezone.utc).isoformat(), - "ready": number(stats.get("totalPrompts")) > 0 or len(limits["limits"]) > 0, + # A signed-in plan shows even before its first session or window. + "ready": number(stats.get("totalPrompts")) > 0 or len(limits["limits"]) > 0 or limits["tierLabel"] not in ("", "Account"), "hasLocalStats": True, # Prompts and sessions are counted, so the panel shows them in tooltips. "hasPromptStats": True, diff --git a/shell/plugins/agents/README.md b/shell/plugins/agents/README.md index 14ebd063531..670aba47ef4 100644 --- a/shell/plugins/agents/README.md +++ b/shell/plugins/agents/README.md @@ -208,8 +208,9 @@ only adds the meter and the spent-of-funded line under the real figure. Muse reads native CLI sessions from `$XDG_DATA_HOME/muse/sessions` (defaulting to `~/.local/share/muse/sessions`, and overridable via -`MUSE_DATA_DIR`) plus opencode sessions on the `meta` provider. Muse Spark -through the OpenCode Zen gateway counts in the `opencode` record instead. The +`MUSE_DATA_DIR`) plus opencode sessions on the `meta` provider, from both of +opencode's message stores (deduplicated by message id). Muse Spark through +the OpenCode Zen gateway counts in the `opencode` record instead. The OAuth token comes from `MUSE_AUTH_PATH` when set, otherwise from `~/.config/muse/auth.json` (which `muse login` creates); the file is only ever read, never refreshed or rewritten. diff --git a/test/shell.d/agent-usage-muse-limits-test.sh b/test/shell.d/agent-usage-muse-limits-test.sh index cf2e821049d..2afb1aea6db 100755 --- a/test/shell.d/agent-usage-muse-limits-test.sh +++ b/test/shell.d/agent-usage-muse-limits-test.sh @@ -170,6 +170,12 @@ nosnap=$(run_collector --force) fail "Muse collector stays silent for an active plan without usage windows" "$nosnap" pass "Muse collector stays silent for an active plan without usage windows" +# A signed-in plan shows before its first session; pay-as-you-go with no +# usage has nothing to show. +[[ $(jq -r '.ready' <<<"$nosnap") == "true" && $(jq -r '.ready' <<<"$payg") == "false" ]] || + fail "Muse collector shows a signed-in plan without sessions" "$nosnap" +pass "Muse collector shows a signed-in plan without sessions" + [[ $(jq 'has("retryAdvised")' <<<"$nosnap") == "false" ]] || fail "Muse collector does not retry an active plan without usage windows" "$nosnap" pass "Muse collector does not retry an active plan without usage windows" diff --git a/test/shell.d/agent-usage-muse-scanner-test.sh b/test/shell.d/agent-usage-muse-scanner-test.sh index 7e226067fe0..887e9418a43 100755 --- a/test/shell.d/agent-usage-muse-scanner-test.sh +++ b/test/shell.d/agent-usage-muse-scanner-test.sh @@ -352,3 +352,20 @@ result_odd=$(HOME="$odd_home" XDG_DATA_HOME="$odd_home/data" XDG_CACHE_HOME="$od [[ $(jq -r '.modelUsage.["odd-cache-model"] | [.inputTokens, .outputTokens] == [11, 0]' <<<"$result_odd") == "true" ]] || fail "Muse collector keeps input past a malformed cache object" "$result_odd" pass "Muse collector survives odd timestamp and token shapes" + +# OpenCode V2 nests the provider under `model` in session_message. A message +# migrated into both stores counts once, and Zen's Muse Spark stays out. +v2_home="$TEST_HOME/v2" +mkdir -p "$v2_home/data/opencode" +sqlite3 "$v2_home/data/opencode/opencode.db" < Date: Fri, 2 Oct 2026 06:45:11 -0400 Subject: [PATCH 033/546] Skip malformed OpenCode rows instead of losing the record A row with valid JSON but a nested field of the wrong shape (a numeric time, a list of tokens) raised past the decode guard, and the cache fallback repeated the same failing scan, so no record came out. Read such fields as empty, and skip any row that still fails, in both message stores and in pi sessions. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-opencode | 29 ++++++++------ .../agent-usage-opencode-scanner-test.sh | 39 ++++++++++++++++++- 2 files changed, 55 insertions(+), 13 deletions(-) diff --git a/bin/omarchy-agent-usage-opencode b/bin/omarchy-agent-usage-opencode index 9b1ceb49c74..b3762bbd300 100755 --- a/bin/omarchy-agent-usage-opencode +++ b/bin/omarchy-agent-usage-opencode @@ -60,6 +60,12 @@ SCAN_REUSE_SECONDS = 20 LIMITS_ONLY_REUSE_SECONDS = 900 +# A nested field of the wrong shape (a list of tokens, a numeric time) reads +# as empty rather than failing the row, or the whole scan with it. +def obj(value): + return value if isinstance(value, dict) else {} + + def number(value): try: return int(value or 0) @@ -210,19 +216,19 @@ def scan_pi_sessions(): except Exception: continue - if entry.get("type") != "message": + if not isinstance(entry, dict) or entry.get("type") != "message": continue message_key = path + ":" + str(entry.get("id") or "") if message_key in seen_pi_messages: continue seen_pi_messages.add(message_key) - message = entry.get("message") or {} + message = obj(entry.get("message")) if message.get("role") != "assistant": continue if str(message.get("provider") or "") not in PROVIDERS: continue - usage = message.get("usage") or {} + usage = obj(message.get("usage")) if not usage: continue input_tokens = number(usage.get("input")) @@ -286,8 +292,8 @@ def store_row(message_id, session_id, entry, row_type, created, seen): # runs never match, so the claude and codex collectors keep those. if str(provider or "") not in PROVIDERS: return - tokens = entry.get("tokens") or {} - cache = tokens.get("cache") or {} + tokens = obj(entry.get("tokens")) + cache = obj(tokens.get("cache")) 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")) @@ -299,7 +305,7 @@ def store_row(message_id, session_id, entry, row_type, created, seen): if message_id in seen: return seen.add(message_id) - day = local_day((entry.get("time") or {}).get("created") or created) + day = local_day(obj(entry.get("time")).get("created") or created) add_usage(day, "opencode:" + str(session_id), model_name(entry.get("modelID") or model), input_tokens, output_tokens, cache_read, cache_write) @@ -322,11 +328,11 @@ def scan_session_message_rows(conn, tables, seen): " WHERE m.type IN ('assistant', 'compaction')" " AND m.data LIKE '%\"providerID\"%:%\"opencode%'" + fork_filter ): + # One row that still fails must not take the scan down with it. try: - entry = json.loads(raw) - except (TypeError, ValueError): + store_row(str(message_id), session_id, json.loads(raw), str(row_type or ""), time_created, seen) + except Exception: continue - store_row(str(message_id), session_id, entry, str(row_type or ""), time_created, seen) def scan_message_rows(conn, seen): @@ -346,10 +352,9 @@ def scan_message_rows(conn, seen): " AND (CASE WHEN json_valid(data) THEN json_extract(data, '$.providerID') END IN ('opencode', 'opencode-go'))" ): try: - entry = json.loads(raw) - except (TypeError, ValueError): + store_row(str(row_id), session_id, json.loads(raw), "", None, seen) + except Exception: continue - store_row(str(row_id), session_id, entry, "", None, seen) def scan_opencode_db(): diff --git a/test/shell.d/agent-usage-opencode-scanner-test.sh b/test/shell.d/agent-usage-opencode-scanner-test.sh index 80942b3c957..49161f5cc2e 100644 --- a/test/shell.d/agent-usage-opencode-scanner-test.sh +++ b/test/shell.d/agent-usage-opencode-scanner-test.sh @@ -212,10 +212,47 @@ result=$(HOME="$V2_HOME" XDG_CACHE_HOME="$V2_HOME/.cache" XDG_DATA_HOME="$V2_HOM fail "OpenCode collector reads V2 records once each, without fork copies" "$result" pass "OpenCode collector reads V2 records once each, without fork copies" +# Valid JSON with a nested field of the wrong shape (a numeric time, a list of +# tokens or of cache counts, a list for a pi message) costs that field or that +# row, never the record. +ODD_HOME=$(mktemp -d) +trap 'rm -rf "$PI_HOME" "$OPENCODE_HOME" "$CACHE_HOME" "$SESSION_HOME" "$V2_HOME" "$ODD_HOME"' EXIT +mkdir -p "$ODD_HOME/.pi/agent/sessions/project" +cat >"$ODD_HOME/.pi/agent/sessions/project/pi.jsonl" < Date: Mon, 12 Oct 2026 10:05:24 -0400 Subject: [PATCH 034/546] Reset kept Cursor, OpenCode, Copilot, and Muse limits to zero once they end Match quattro's rule for kept limits: a window whose reset has passed comes back at 0% with no reset time instead of disappearing, and stale limits are no longer described as dimmed, since the panel now gives their age on hover. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-copilot | 8 +++++--- bin/omarchy-agent-usage-cursor | 7 +++++-- bin/omarchy-agent-usage-muse | 10 +++++++--- bin/omarchy-agent-usage-opencode | 9 ++++++--- shell/plugins/agents/README.md | 11 ++++++----- test/shell.d/agent-usage-copilot-test.sh | 8 ++++---- test/shell.d/agent-usage-cursor-test.sh | 10 +++++----- test/shell.d/agent-usage-muse-limits-test.sh | 10 ++++++++++ test/shell.d/agent-usage-opencode-limits-test.sh | 7 ++++--- 9 files changed, 52 insertions(+), 28 deletions(-) diff --git a/bin/omarchy-agent-usage-copilot b/bin/omarchy-agent-usage-copilot index d306921c096..85e24d9ea02 100755 --- a/bin/omarchy-agent-usage-copilot +++ b/bin/omarchy-agent-usage-copilot @@ -778,7 +778,7 @@ def fetch_remote_quota() -> dict[str, Any]: """The live quota, or the last good one marked stale. Every failure is soft: a missing token, an offline machine, or a changed - payload keeps the last answer this token got, dimmed, or leaves the record + payload keeps the last answer this token got, marked stale, or leaves the record on its local stats rather than failing the whole collection. The cache is keyed by the token, so a machine signed in to another account never shows the last one's allowance. @@ -798,14 +798,16 @@ def fetch_remote_quota() -> dict[str, Any]: return quota cached = read_fresh_json(cache, float("inf")) or {} - if cached.get("available") and quota_window_open(cached): + if cached.get("available"): + if not quota_window_open(cached): + cached.update(used=0.0, resetsAt="") cached.update(stale=True, failure=quota.get("reason", "")) return cached return quota # A kept answer describes its own allowance window and nothing after it: once -# that window resets, its usage is gone. Without a reset date, the window is +# that window resets, its usage is back at zero. Without a reset date, the window is # the calendar month (UTC) it was fetched in. def quota_window_open(cached: dict[str, Any]) -> bool: now = dt.datetime.now(dt.timezone.utc) diff --git a/bin/omarchy-agent-usage-cursor b/bin/omarchy-agent-usage-cursor index cd697dff0ad..d5e64ebc60e 100755 --- a/bin/omarchy-agent-usage-cursor +++ b/bin/omarchy-agent-usage-cursor @@ -386,7 +386,8 @@ def limits_cache(key: str) -> Path: return cache_root() / f"cursor-limits-{key}.json" -# A kept limit is only good until its billing cycle ends. +# A kept limit resets to zero once its billing cycle ends; the next cycle's +# reset time is unknown until a check succeeds. def open_limits(limits: Any) -> list[dict[str, Any]]: now = datetime.now(timezone.utc) kept = [] @@ -399,6 +400,8 @@ def open_limits(limits: Any) -> list[dict[str, Any]]: resets_at = None if resets_at is None or resets_at.tzinfo is None or resets_at > now: kept.append(limit) + else: + kept.append(dict(limit, percent=0.0, resetsAt="")) return kept @@ -419,7 +422,7 @@ def scan(api_base_url: str, auth_path: Path, limits_only: bool, state_db: Path | period = client.period() except CursorError as error: # A failed check keeps the last good limits, marked stale so the panel - # dims them, and the week of tokens already on disk. + # gives their age on hover, and the week of tokens already on disk. cached = read_json(cache) limits = open_limits(cached.get("limits")) status = "Waiting for auth" if error.auth else ("" if limits else "Cursor limits unavailable") diff --git a/bin/omarchy-agent-usage-muse b/bin/omarchy-agent-usage-muse index 34da423a019..0639df6759a 100755 --- a/bin/omarchy-agent-usage-muse +++ b/bin/omarchy-agent-usage-muse @@ -583,7 +583,8 @@ def probe_key(access_token: str) -> dict: def limit_window_open(entry: dict, now: datetime) -> bool: # A cached percentage lives only until its window rolls over; past-reset - # figures describe a period that is over. No (or unreadable) reset is kept. + # figures describe a period that is over and read as zero. No (or + # unreadable) reset keeps its number. raw = str(entry.get("resetsAt") or "") if raw == "": return True @@ -601,7 +602,10 @@ def usable_cached_limits(cached: dict) -> list: 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)] + return [ + entry if limit_window_open(entry, now) else dict(entry, percent=0.0, resetsAt="") + for entry in entries if isinstance(entry, dict) + ] # Keyed by who the token signs in as (the JWT's subject when it is one), so a @@ -642,7 +646,7 @@ def collect_limits(force: bool) -> dict: fetched_at = number(cached.get("fetchedAtMs")) result["tierLabel"] = str(cached.get("tierLabel") or "") or "Account" - # The last good windows, dimmed as stale, with when they were measured. + # The last good windows, marked stale, with when they were measured. def keep_cached(): result["limits"] = fallback result["stale"] = True diff --git a/bin/omarchy-agent-usage-opencode b/bin/omarchy-agent-usage-opencode index b3762bbd300..7211f497620 100755 --- a/bin/omarchy-agent-usage-opencode +++ b/bin/omarchy-agent-usage-opencode @@ -683,7 +683,7 @@ def probe_limits(credential): def limit_window_open(entry, now): # A cached percentage stops meaning anything once its window resets: a # stale 78% would misreport an allowance that is now untouched. An - # unreadable timestamp is no reason to drop a real number. + # unreadable timestamp is no reason to doubt a real number. raw = str(entry.get("resetsAt") or "") if raw == "": return True @@ -701,7 +701,10 @@ def open_cached_limits(cached): 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)] + return [ + entry if limit_window_open(entry, now) else dict(entry, percent=0.0, resetsAt="") + for entry in entries if isinstance(entry, dict) + ] # One cache per credential, so a machine signed in to another account never @@ -754,7 +757,7 @@ def collect_limits(force): refused = refused[:3] + (expired,) continue - # Last good numbers, dimmed as stale, beat an empty limits section while + # Last good numbers, marked stale, beat an empty limits section while # the probe is down. result["retryAdvised"] = bool(probe.get("transport")) answer(fallback, fetched_at, True) diff --git a/shell/plugins/agents/README.md b/shell/plugins/agents/README.md index 5113a305c0a..dec9af4c4ab 100644 --- a/shell/plugins/agents/README.md +++ b/shell/plugins/agents/README.md @@ -118,8 +118,8 @@ editor's own sign-in in `~/.config/Cursor/User/globalStorage/state.vscdb` Cursor's session files on disk carry no token counts, so this is the one collector with nothing to fall back on: without a credential the record stays empty and the panel skips it rather than showing a week of zeros. A failed -check keeps the account's last good meters, dimmed, until its billing cycle -ends, and its last week of tokens (today cleared once the day turns over); a +check keeps the account's last good meters, their age on hover, back at zero +once its billing cycle ends, and its last week of tokens (today cleared once the day turns over); a refused sign-in also says so under the name. Tokens by day and by model cover the last seven local days, one aggregation call per day. @@ -145,7 +145,7 @@ its database and `auth.json`, then pi's, and last the OpenCode Console session, which asks Console's own endpoint with its workspace id. A refused credential hands over to the next one. A recent answer is reused for 15 seconds, each credential caches its own, a failed check keeps the last good -windows, dimmed, until they reset, and a rate-limited window reads as full. +windows, back at zero once they reset, and a rate-limited window reads as full. Without a Go plan the panel still shows local token usage. ### Copilot allowance @@ -166,7 +166,8 @@ every machine, IDE, and agent on the entitlement and names the plan and the reset date. It needs a token from `COPILOT_QUOTA_TOKEN`, `GH_TOKEN`, `GITHUB_TOKEN`, or `gh auth token`, and only runs where the Copilot CLI has been used, so a machine with `gh` signed in but no Copilot stays out of the -panel. A failed check keeps the last answer, dimmed. When the endpoint can't +panel. A failed check keeps the last answer, back at zero once the month +resets. When the endpoint can't be asked at all, a budget in `~/.config/omarchy/agents/copilot.json` gives a meter labeled "(est.)" from this machine's CLI spend alone: @@ -225,7 +226,7 @@ session and weekly percents with reset times — the same figures the TUI's `/usage` overlay shows. Minting is idempotent (the same Model API key comes back every call), so polling is safe; the token travels only in the `Authorization` header and neither it nor the returned key is persisted. A -failed check keeps the last good windows, dimmed, until they reset. +failed check keeps the last good windows, back at zero once they reset. Pay-as-you-go accounts have no subscription windows, so the tab shows measured local usage with no meters — never estimated, never zeroed. An active plan that returns no `subs_usage` (Muse Code Power Usage answers this diff --git a/test/shell.d/agent-usage-copilot-test.sh b/test/shell.d/agent-usage-copilot-test.sh index ca8edb86c51..989c5aff4bd 100755 --- a/test/shell.d/agent-usage-copilot-test.sh +++ b/test/shell.d/agent-usage-copilot-test.sh @@ -182,7 +182,7 @@ collector.urllib.request.urlopen = lambda request, timeout=None: io.BytesIO(json record = run() collector.urllib.request.urlopen = offline stale = run() -# Once the kept answer's window has reset, its usage no longer applies. +# Once the kept answer's window has reset, its usage is back at zero. for cache in collector.cache_root().glob("copilot-quota-*.json"): kept = json.loads(cache.read_text()) cache.write_text(json.dumps(dict(kept, resetsAt="2020-01-01T00:00:00.000Z"))) @@ -202,9 +202,9 @@ pass "Copilot collector stamps the live allowance" fail "Copilot collector keeps the last live allowance, stale, after a failed check" "$result" pass "Copilot collector keeps the last live allowance, stale, after a failed check" -[[ $(jq -c '.lapsed | {label: .limits[0].label, stale: .limitsStale, tierLabel}' <<<"$result") == '{"label":"Monthly allowance (est.)","stale":false,"tierLabel":""}' ]] || - fail "Copilot collector drops a kept allowance once its window resets" "$result" -pass "Copilot collector drops a kept allowance once its window resets" +[[ $(jq -c '.lapsed | {label: .limits[0].label, zero: (.limits[0].percent == 0), stale: .limitsStale, tierLabel}' <<<"$result") == '{"label":"Monthly allowance","zero":true,"stale":true,"tierLabel":"Business"}' ]] || + fail "Copilot collector resets a kept allowance once its window ends" "$result" +pass "Copilot collector resets a kept allowance once its window ends" [[ $(jq -c '.clean | {ready, tierLabel, limits}' <<<"$result") == '{"ready":false,"tierLabel":"","limits":[]}' ]] || fail "Copilot collector doesn't ask GitHub where the CLI never ran" "$result" diff --git a/test/shell.d/agent-usage-cursor-test.sh b/test/shell.d/agent-usage-cursor-test.sh index d17166c1803..b4a748e5297 100755 --- a/test/shell.d/agent-usage-cursor-test.sh +++ b/test/shell.d/agent-usage-cursor-test.sh @@ -129,7 +129,7 @@ cache.write_text(json.dumps(yesterday_scan)) overnight = collector.carried_stats(key, today) cache.write_text(json.dumps(scan)) -# A kept limit whose billing cycle has ended is dropped, not shown stale. +# A kept limit whose billing cycle has ended comes back at zero. limits_cache = collector.limits_cache(key) good = json.loads(limits_cache.read_text()) limits_cache.write_text(json.dumps(dict(good, limits=[dict(good["limits"][0], resetsAt="2020-01-01T00:00:00+00:00")] + good["limits"][1:]))) @@ -158,7 +158,7 @@ refused = collector.scan("https://api.example", Path(os.environ["XDG_CONFIG_HOME print(json.dumps({ "other": other, "overnight": overnight, - "lapsed": [limit["label"] for limit in lapsed["limits"]], + "lapsed": [[limit["label"], round(limit["percent"] * 100), limit["resetsAt"]] for limit in lapsed["limits"][:2]], "down": down, "refused": refused, "record": record, @@ -226,9 +226,9 @@ pass "Cursor collector never carries another account's numbers" fail "Cursor collector moves a scan from yesterday onto today's week" "$result" pass "Cursor collector moves a scan from yesterday onto today's week" -[[ $(jq -c '.lapsed' <<<"$result") == '["Auto models","Named models","On-demand $9.51 / $10.00"]' ]] || - fail "Cursor collector drops a kept limit whose cycle ended" "$result" -pass "Cursor collector drops a kept limit whose cycle ended" +[[ $(jq -c '.lapsed' <<<"$result") == '[["Included total",0,""],["Auto models",12,"2026-10-15T21:32:06+00:00"]]' ]] || + fail "Cursor collector resets a kept limit whose cycle ended" "$result" +pass "Cursor collector resets a kept limit whose cycle ended" # Days that spring forward or fall back are 23 or 25 hours long. dst=$(TZ=America/New_York python3 - "$ROOT/bin/omarchy-agent-usage-cursor" <<'PY' diff --git a/test/shell.d/agent-usage-muse-limits-test.sh b/test/shell.d/agent-usage-muse-limits-test.sh index 2afb1aea6db..959807dc773 100755 --- a/test/shell.d/agent-usage-muse-limits-test.sh +++ b/test/shell.d/agent-usage-muse-limits-test.sh @@ -145,6 +145,16 @@ refused=$(run_collector --force) fail "Muse collector keeps the last limits through a rejected sign-in" "$refused" pass "Muse collector keeps the last limits through a rejected sign-in" +# A kept window that has since reset comes back at zero, with no reset time +# until a check succeeds again. +limits_cache=$(echo "$TEST_HOME"/.cache/omarchy/agent-usage/muse-limits-*.json) +jq '.limits[0].resetsAt = "2020-01-01T00:00:00+00:00"' "$limits_cache" >"$limits_cache.tmp" && mv "$limits_cache.tmp" "$limits_cache" +printf '{"status":500}' >"$STUB_PAYLOAD_FILE" +lapsed=$(run_collector --force) +[[ $(jq -c '[.limits[] | {zero: (.percent == 0), reset: (.resetsAt != "")}]' <<<"$lapsed") == '[{"zero":true,"reset":false},{"zero":false,"reset":true}]' ]] || + fail "Muse collector resets a kept window once it has ended" "$lapsed" +pass "Muse collector resets a kept window once it has ended" + # The login file belongs to the Muse CLI: every run so far only read it. [[ $(auth_state) == "$auth_before" && $(ls -A "$TEST_HOME/.config/muse") == "auth.json" ]] || fail "Muse collector never writes the login file" "$(ls -lA --time-style=+%s "$TEST_HOME/.config/muse")" diff --git a/test/shell.d/agent-usage-opencode-limits-test.sh b/test/shell.d/agent-usage-opencode-limits-test.sh index 860f581ea64..a9b06de0061 100644 --- a/test/shell.d/agent-usage-opencode-limits-test.sh +++ b/test/shell.d/agent-usage-opencode-limits-test.sh @@ -153,7 +153,7 @@ check( ) # A failed probe keeps the last good numbers, marked stale with when they -# were measured, but only until each window resets. +# were measured, with each window that has since reset back at zero. os.environ["OPENCODE_API_KEY"] = "zen_test" collector.cache_root().mkdir(parents=True, exist_ok=True) collector.limits_cache_path("zen_test").write_text(json.dumps({ @@ -166,8 +166,9 @@ collector.limits_cache_path("zen_test").write_text(json.dumps({ collector.urllib.request.urlopen = answer(error=urllib.error.HTTPError("https://x", 500, "Server Error", {}, None)) result = collector.collect_limits(True) check( - "OpenCode collector keeps open cached windows, stale, after a failed probe", - [w["label"] for w in result["limits"]] == ["Monthly"] and result["stale"] is True + "OpenCode collector keeps cached windows, stale, after a failed probe", + [(w["label"], w["percent"], w["resetsAt"]) for w in result["limits"]] + == [("Rolling (5h)", 0.0, ""), ("Monthly", 0.4, "2999-01-01T00:00:00.000Z")] and result["stale"] is True and result["fetchedAtMs"] == 1700000000000 and result["tierLabel"] == "Go" and result["usageStatusText"] == "", json.dumps(result), ) From 7dc296d81ed88ad04d802279e7d1e6cbfd964f34 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Mon, 12 Oct 2026 10:09:13 -0400 Subject: [PATCH 035/546] Send OpenCode only its own key and count id-less pi messages A custom provider merely named like opencode (an "opencode-local" gateway) had its key sent to opencode.ai, and pi lines without an id all shared one dedup key, so only the first was counted. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-opencode | 17 +++++++----- .../agent-usage-opencode-limits-test.sh | 5 ++++ .../agent-usage-opencode-scanner-test.sh | 26 ++++++++++++++----- 3 files changed, 35 insertions(+), 13 deletions(-) diff --git a/bin/omarchy-agent-usage-opencode b/bin/omarchy-agent-usage-opencode index 7211f497620..e668da83e3a 100755 --- a/bin/omarchy-agent-usage-opencode +++ b/bin/omarchy-agent-usage-opencode @@ -218,10 +218,14 @@ def scan_pi_sessions(): if not isinstance(entry, dict) or entry.get("type") != "message": continue - message_key = path + ":" + str(entry.get("id") or "") - if message_key in seen_pi_messages: - continue - seen_pi_messages.add(message_key) + # Only an id can say two lines are the same message; without one, each + # line is its own. + message_id = str(entry.get("id") or "") + if message_id: + message_key = path + ":" + message_id + if message_key in seen_pi_messages: + continue + seen_pi_messages.add(message_key) message = obj(entry.get("message")) if message.get("role") != "assistant": continue @@ -541,10 +545,9 @@ def read_auth_key(path): return "" if not isinstance(parsed, dict): return "" - # Prefer the exact provider ids, then any opencode* entry: the login flow - # sometimes files the key under a different id than the runtime uses. + # Only the subscription's own provider ids: a custom gateway named like + # "opencode-local" holds a key for somewhere else, never to be sent here. entries = [parsed[provider] for provider in PROVIDERS if provider in parsed] - entries += [entry for name, entry in parsed.items() if "opencode" in str(name).lower()] for entry in entries: if isinstance(entry, dict) and entry.get("key"): return str(entry["key"]).strip() diff --git a/test/shell.d/agent-usage-opencode-limits-test.sh b/test/shell.d/agent-usage-opencode-limits-test.sh index a9b06de0061..b861b20ae5f 100644 --- a/test/shell.d/agent-usage-opencode-limits-test.sh +++ b/test/shell.d/agent-usage-opencode-limits-test.sh @@ -212,6 +212,11 @@ with open(os.path.join(opencode_home, "opencode", "auth.json"), "w") as handle: os.environ["XDG_DATA_HOME"] = opencode_home check("OpenCode collector reads opencode's own credential", [c.token for c in collector.credential_candidates()] == ["from_opencode"]) +# A custom provider merely named like opencode holds a key for somewhere else. +with open(os.path.join(opencode_home, "opencode", "auth.json"), "w") as handle: + json.dump({"opencode-local": {"type": "api_key", "key": "gateway_secret"}}, handle) +check("OpenCode collector never sends a look-alike provider's key", "gateway_secret" not in [c.token for c in collector.credential_candidates()]) + clear_key() pi_home = os.path.join(home, "pi") os.makedirs(os.path.join(pi_home, ".pi", "agent"), exist_ok=True) diff --git a/test/shell.d/agent-usage-opencode-scanner-test.sh b/test/shell.d/agent-usage-opencode-scanner-test.sh index 49161f5cc2e..abcf48b38d4 100644 --- a/test/shell.d/agent-usage-opencode-scanner-test.sh +++ b/test/shell.d/agent-usage-opencode-scanner-test.sh @@ -34,6 +34,20 @@ result=$(HOME="$PI_HOME" XDG_CACHE_HOME="$PI_HOME/.cache" XDG_DATA_HOME="$PI_HOM fail "OpenCode collector folds reasoning into output and strips the thinking suffix" "$result" pass "OpenCode collector counts pi and omp subscription usage" +# pi lines without an id can't be told apart, so each one counts. +NOID_HOME=$(mktemp -d) +trap 'rm -rf "$PI_HOME" "$NOID_HOME"' EXIT +mkdir -p "$NOID_HOME/.pi/agent/sessions/project" +cat >"$NOID_HOME/.pi/agent/sessions/project/pi.jsonl" <"$CACHE_HOME/.pi/agent/sessions/project/pi.jsonl" <"$ODD_HOME/.pi/agent/sessions/project/pi.jsonl" < Date: Mon, 12 Oct 2026 10:10:00 -0400 Subject: [PATCH 036/546] Don't cache a Copilot scan whose session store couldn't be read A busy or oversized store left only the transcripts, which newer CLIs may not write, and that short scan was served for the next 15 minutes. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-copilot | 22 ++++++++++++++-------- test/shell.d/agent-usage-copilot-test.sh | 12 ++++++++++++ 2 files changed, 26 insertions(+), 8 deletions(-) diff --git a/bin/omarchy-agent-usage-copilot b/bin/omarchy-agent-usage-copilot index 85e24d9ea02..637ea674429 100755 --- a/bin/omarchy-agent-usage-copilot +++ b/bin/omarchy-agent-usage-copilot @@ -677,29 +677,35 @@ def cached_scan(home_dir: Path, max_age_seconds: float) -> dict[str, Any]: cached = read_cached_stats(cache_file, max_age_seconds, today) if cached is not None: return cached - stats = direct_scan(home_dir) - write_cached_stats(cache_file, stats) + stats, complete = direct_scan(home_dir) + # A scan cut short isn't served to the next runs; they try again. + if complete: + write_cached_stats(cache_file, stats) return stats except Exception as exc: print( f"omarchy-agent-usage-copilot: cache unavailable ({exc}); scanning directly", file=sys.stderr, ) - return direct_scan(home_dir) + return direct_scan(home_dir)[0] -def direct_scan(home_dir: Path) -> dict[str, Any]: - conn, tmpdir = open_store(home_dir / "session-store.db") +# The stats, and whether they are complete: a store that is there but could +# not be snapshotted (busy, or too large) leaves only the transcripts, which +# newer CLIs may not write. +def direct_scan(home_dir: Path) -> tuple[dict[str, Any], bool]: + store = home_dir / "session-store.db" + conn, tmpdir = open_store(store) if conn is None: - return scan_transcripts(home_dir / "session-state") + return scan_transcripts(home_dir / "session-state"), not store.exists() try: if has_table(conn, "assistant_usage_events"): stats = scan_store(conn) if number(stats.get("totalPrompts")) > 0: - return stats + return stats, True # A store without the usage table, or one that is still empty, falls back # to the transcripts exactly like a store-less machine. - return scan_transcripts(home_dir / "session-state") + return scan_transcripts(home_dir / "session-state"), True finally: conn.close() if tmpdir is not None: diff --git a/test/shell.d/agent-usage-copilot-test.sh b/test/shell.d/agent-usage-copilot-test.sh index 989c5aff4bd..6e5afc5ed08 100755 --- a/test/shell.d/agent-usage-copilot-test.sh +++ b/test/shell.d/agent-usage-copilot-test.sh @@ -293,6 +293,18 @@ pass "Copilot transcript fallback sums today's tokens" fail "Copilot transcript fallback estimates spend from shutdown cost" "$result" pass "Copilot transcript fallback estimates spend from shutdown cost" +# A store that is there but can't be snapshotted leaves only the transcripts, +# and that short scan isn't cached for the runs after it. +[[ -n $(find "$TRANSCRIPT_HOME/.cache" -name 'copilot-scan-*.json') ]] || + fail "Copilot collector caches a complete scan" "$(find "$TRANSCRIPT_HOME/.cache")" +find "$TRANSCRIPT_HOME/.cache" -name 'copilot-scan-*.json' -delete +printf 'not a database' >"$TRANSCRIPT_HOME/.copilot/session-store.db" +result=$(run_collector_in "$TRANSCRIPT_HOME") +[[ $(jq -r '.todayPrompts' <<<"$result") == "3" && -z $(find "$TRANSCRIPT_HOME/.cache" -name 'copilot-scan-*.json') ]] || + fail "Copilot collector doesn't cache a scan its store couldn't join" "$result" +pass "Copilot collector doesn't cache a scan its store couldn't join" +rm "$TRANSCRIPT_HOME/.copilot/session-store.db" + # A segment that spans midnight books its shutdown totals on the days its # replies landed on: output by each day's reply output (40 / 60), everything # else by reply count (2 / 1). From b01407ab92deb446a96e4a35c44887fb001265c7 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Mon, 12 Oct 2026 10:10:30 -0400 Subject: [PATCH 037/546] Keep Copilot's billed requests out of the panel's prompt totals The session store holds a row per billed model request, and one prompt can make many, so counting rows as prompts inflated the totals. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-copilot | 5 ++++- test/shell.d/agent-usage-copilot-test.sh | 8 ++++++++ 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/bin/omarchy-agent-usage-copilot b/bin/omarchy-agent-usage-copilot index 637ea674429..81ad101e374 100755 --- a/bin/omarchy-agent-usage-copilot +++ b/bin/omarchy-agent-usage-copilot @@ -407,7 +407,9 @@ def scan_store(conn: sqlite3.Connection) -> dict[str, Any]: ) except sqlite3.Error: return blank_stats() - return tally.stats() + # A row is a billed model request, and one prompt can make many, so its + # count stays out of the panel's prompt totals. + return dict(tally.stats(), hasPromptStats=False) # ------------------------------------------------------ transcript fallback @@ -922,6 +924,7 @@ def main() -> int: "updatedAt": now.isoformat(), "ready": number(stats.get("totalPrompts")) > 0 or len(allowance["limits"]) > 0, "hasLocalStats": True, + "hasPromptStats": stats.get("hasPromptStats") is not False, "tierLabel": allowance["tierLabel"], "usageStatusText": allowance["usageStatusText"], "authHelpText": allowance["authHelpText"], diff --git a/test/shell.d/agent-usage-copilot-test.sh b/test/shell.d/agent-usage-copilot-test.sh index 6e5afc5ed08..746969d950d 100755 --- a/test/shell.d/agent-usage-copilot-test.sh +++ b/test/shell.d/agent-usage-copilot-test.sh @@ -79,6 +79,10 @@ pass "Copilot collector is ready with usage" fail "Copilot collector counts billed requests" "$result" pass "Copilot collector counts billed requests" +[[ $(jq -r '.hasPromptStats' <<<"$result") == "false" ]] || + fail "Copilot collector keeps billed requests out of the prompt totals" "$result" +pass "Copilot collector keeps billed requests out of the prompt totals" + [[ $(jq -r '.todaySessions' <<<"$result") == "2" ]] || fail "Copilot collector counts distinct sessions" "$result" pass "Copilot collector counts distinct sessions" @@ -281,6 +285,10 @@ result=$(run_collector_in "$TRANSCRIPT_HOME") fail "Copilot transcript fallback counts the user's prompts" "$result" pass "Copilot transcript fallback counts the user's prompts" +[[ $(jq -r '.hasPromptStats' <<<"$result") == "true" ]] || + fail "Copilot transcript fallback offers its prompts to the panel" "$result" +pass "Copilot transcript fallback offers its prompts to the panel" + [[ $(jq -c '.modelUsage["claude-sonnet-4.6"]' <<<"$result") == '{"inputTokens":500,"outputTokens":340,"cacheReadInputTokens":900,"cacheCreationInputTokens":100}' ]] || fail "Copilot transcript fallback reads shutdown usage once" "$result" pass "Copilot transcript fallback reads shutdown usage once" From 229e4c35b79bdb4eb21c310fb9c09e4889ad45b0 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Mon, 12 Oct 2026 10:11:09 -0400 Subject: [PATCH 038/546] Fetch Cursor tokens again once a limits-only run's scan is 15 minutes old Limits-only runs keep going while any meter sits near its limit, which can last hours, and Cursor served its last token scan through all of them. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-cursor | 14 ++++++++++++-- test/shell.d/agent-usage-cursor-test.sh | 15 +++++++++++++++ 2 files changed, 27 insertions(+), 2 deletions(-) diff --git a/bin/omarchy-agent-usage-cursor b/bin/omarchy-agent-usage-cursor index d5e64ebc60e..3f72f39b092 100755 --- a/bin/omarchy-agent-usage-cursor +++ b/bin/omarchy-agent-usage-cursor @@ -36,6 +36,8 @@ API_BASE_URL = "https://api2.cursor.sh" SERVICE = "aiserver.v1.DashboardService" REQUEST_TIMEOUT = 20 RECENT_DAYS = 7 +# How long --limits-only serves the last token scan before asking again. +LIMITS_ONLY_REUSE_SECONDS = 900 class CursorError(Exception): @@ -299,12 +301,18 @@ def carried_stats(key: str, today: date) -> dict[str, Any]: return stats +def carried_stats_fresh(key: str, today: date) -> bool: + cached = read_json(stats_cache(key)) + age = now_ms() - number(cached.get("scannedAtMs")) + return cached.get("scanDate") == today.isoformat() and 0 <= age < LIMITS_ONLY_REUSE_SECONDS * 1000 + + def collect_stats(client: CursorClient, today: date, key: str) -> dict[str, Any]: windows = day_windows(today, RECENT_DAYS) with ThreadPoolExecutor(max_workers=4) as pool: rows = list(pool.map(lambda window: client.aggregations(window[1], window[2]), windows)) stats = summarize_usage([(window[0], summarize_day(day_rows)) for window, day_rows in zip(windows, rows)]) - write_json(stats_cache(key), {"scanDate": today.isoformat(), "stats": stats}) + write_json(stats_cache(key), {"scanDate": today.isoformat(), "scannedAtMs": now_ms(), "stats": stats}) return stats @@ -446,7 +454,9 @@ def scan(api_base_url: str, auth_path: Path, limits_only: bool, state_db: Path | limitsFetchedAt=now_ms(), ) write_json(cache, {"fetchedAtMs": record["limitsFetchedAt"], "tierLabel": record["tierLabel"], "limits": record["limits"]}) - if limits_only: + # --limits-only runs while a meter is near its limit, which can last hours, + # so the tokens it carries are only so old before they are fetched again. + if limits_only and carried_stats_fresh(key, today): record.update(carried_stats(key, today)) else: try: diff --git a/test/shell.d/agent-usage-cursor-test.sh b/test/shell.d/agent-usage-cursor-test.sh index b4a748e5297..acca6685c1d 100755 --- a/test/shell.d/agent-usage-cursor-test.sh +++ b/test/shell.d/agent-usage-cursor-test.sh @@ -105,6 +105,16 @@ first_start = datetime.fromtimestamp(windows[0][1] / 1000).astimezone() # --limits-only keeps the token sections of this account's last scan. carried = collector.scan("https://api.example", Path(os.environ["XDG_CONFIG_HOME"]) / "cursor" / "auth.json", True) +# ...but only for so long: an old scan is fetched again. +stats_file = collector.stats_cache(collector.account_key("cursor_test")) +fresh_scan = json.loads(stats_file.read_text()) +stats_file.write_text(json.dumps(dict(fresh_scan, scannedAtMs=fresh_scan["scannedAtMs"] - collector.LIMITS_ONLY_REUSE_SECONDS * 1000))) +calls_before = len(calls) +collector.scan("https://api.example", Path(os.environ["XDG_CONFIG_HOME"]) / "cursor" / "auth.json", True) +refetched = sum(1 for method, _ in calls[calls_before:] if method == "GetAggregatedUsageEvents") +del calls[calls_before:] +stats_file.write_text(json.dumps(fresh_scan)) + # Another account's scan never stands in: signed in as someone else with the # dashboard down, there are no tokens to show. class DownForOther(StubClient): @@ -165,6 +175,7 @@ print(json.dumps({ "aggregationCalls": sum(1 for method, _ in calls if method == "GetAggregatedUsageEvents"), "windowIsLocalMidnight": (first_start.hour, first_start.minute, first_start.second) == (0, 0, 0), "carriedRecentDays": carried["recentDays"] == record["recentDays"], + "refetched": refetched, "carriedSkipsAggregation": sum(1 for method, _ in calls if method == "GetAggregatedUsageEvents") == collector.RECENT_DAYS, })) PY @@ -206,6 +217,10 @@ pass "Cursor collector separates today from the rest of the week" fail "Cursor collector keeps the token sections on --limits-only" "$result" pass "Cursor collector keeps the token sections on --limits-only" +[[ $(jq -r '.refetched' <<<"$result") == "7" ]] || + fail "Cursor collector fetches tokens again once the carried scan is old" "$result" +pass "Cursor collector fetches tokens again once the carried scan is old" + [[ $(jq -c '.record | {limitsStale, fetched: (.limitsFetchedAt > 0)}' <<<"$result") == '{"limitsStale":false,"fetched":true}' ]] || fail "Cursor collector stamps fresh limits" "$result" pass "Cursor collector stamps fresh limits" From d377e69e7dfe23b053ebb2d16f37c5dd8f46c728 Mon Sep 17 00:00:00 2001 From: Rene Limberger <4841242+rlimberger@users.noreply.github.com> Date: Thu, 8 Oct 2026 11:28:14 +0200 Subject: [PATCH 039/546] Add Grok Bot's weekly limit to the agents panel Grok Bot runs on a Cursor account, but its weekly allowance is metered apart from Cursor's plan, by the dashboard's GetSandUsageStatus RPC. A grokbot collector asks it with the Cursor CLI's sign-in, or the editor's, and writes a limits-only record; an account without a Grok Bot allowance gives an empty one, so the panel only shows Grok Bot to people who have it. A failed check keeps the last limit, dimmed, until its week ends. No panel changes: the collector and its mark are picked up by convention. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-grokbot | 218 ++++++++++++++++++ shell/plugins/agents/assets/grokbot-light.svg | 1 + shell/plugins/agents/assets/grokbot.svg | 1 + test/shell.d/agent-usage-grokbot-test.sh | 152 ++++++++++++ 4 files changed, 372 insertions(+) create mode 100755 bin/omarchy-agent-usage-grokbot create mode 100644 shell/plugins/agents/assets/grokbot-light.svg create mode 100644 shell/plugins/agents/assets/grokbot.svg create mode 100755 test/shell.d/agent-usage-grokbot-test.sh diff --git a/bin/omarchy-agent-usage-grokbot b/bin/omarchy-agent-usage-grokbot new file mode 100755 index 00000000000..e263d3726b0 --- /dev/null +++ b/bin/omarchy-agent-usage-grokbot @@ -0,0 +1,218 @@ +#!/usr/bin/python3 +# omarchy:summary=Print the Grok Bot usage record as JSON +# omarchy:args=[--force] [--limits-only] +# omarchy:hidden=true +"""Collect Grok Bot usage into one display-ready JSON record. + +Grok Bot runs on a Cursor account, but its weekly allowance is metered apart +from Cursor's own plan: the dashboard's GetSandUsageStatus RPC (Sand is Grok +Bot's internal name) says how much of this week's included usage is gone and +when it resets. It is asked as the same account the Cursor collector reads, +so the two never describe different people. Grok Bot keeps +nothing countable on disk, so the record carries the limit alone, and an +account without a Grok Bot allowance gives an empty record, which the panel +skips. The agents panel only ever reads the JSON this prints. +""" + +import argparse +import base64 +import hashlib +import http.client +import json +import os +import sqlite3 +import tempfile +import urllib.error +import urllib.request +from datetime import datetime, timezone +from pathlib import Path + +AGENT_ID = "grokbot" +AGENT_NAME = "Grok Bot" +AUTH_HELP = "Sign in to the Cursor editor, or run `cursor-agent login`, with your Grok Bot account." +USAGE_URL = "https://api2.cursor.sh/aiserver.v1.DashboardService/GetSandUsageStatus" + + +def config_home(): + return Path(os.environ.get("XDG_CONFIG_HOME") or (Path.home() / ".config")) + + +def cache_root(): + return Path(os.environ.get("XDG_CACHE_HOME") or (Path.home() / ".cache")) / "omarchy" / "agent-usage" + + +def read_json(path): + try: + data = json.loads(Path(path).read_text(encoding="utf-8")) + except (OSError, ValueError): + return {} + return data if isinstance(data, dict) else {} + + +# A temporary file of its own per write, so overlapping runs can't rename +# each other's away. +def write_json(path, payload): + 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): + try: + parsed = datetime.fromisoformat(str(value or "").replace("Z", "+00:00")) + except ValueError: + return None + return parsed if parsed.tzinfo else parsed.replace(tzinfo=timezone.utc) + + +# The editor's sign-in, read-only: the editor may be writing its database +# right now. No immutable=1, so a live database with an active WAL still reads. +def read_ide_token(state_db): + if not state_db.is_file(): + return "" + try: + conn = sqlite3.connect(state_db.resolve().as_uri() + "?mode=ro", uri=True, timeout=2) + except sqlite3.Error: + return "" + try: + row = conn.execute("SELECT value FROM ItemTable WHERE key = 'cursorAuth/accessToken' LIMIT 1").fetchone() + except sqlite3.Error: + return "" + finally: + conn.close() + value = row[0] if row else "" + if isinstance(value, bytes): + value = value.decode("utf-8", errors="replace") + return str(value or "").strip() + + +# The same account the Cursor collector reads, in the same order: the key the +# user exported, the credential the Cursor CLI stores, or the editor's sign-in. +def read_token(): + token = str(os.environ.get("CURSOR_API_KEY") or "").strip() + token = token or str(read_json(config_home() / "cursor" / "auth.json").get("accessToken") or "").strip() + return token or read_ide_token(config_home() / "Cursor" / "User" / "globalStorage" / "state.vscdb") + + +# Who the token signs in as (the JWT's subject), so a machine signed in to +# someone else never shows the last one's limit, and a refreshed token keeps +# the same cache. +def account_key(token): + parts = token.split(".") + subject = "" + if len(parts) == 3: + try: + subject = str(json.loads(base64.urlsafe_b64decode(parts[1] + "=" * (-len(parts[1]) % 4))).get("sub") or "") + except (ValueError, AttributeError): + subject = "" + return hashlib.sha256((subject or token).encode("utf-8")).hexdigest()[:16] + + +def fetch_status(token): + request = urllib.request.Request( + USAGE_URL, + data=b"{}", + headers={ + "Content-Type": "application/json", + "Connect-Protocol-Version": "1", + "Authorization": f"Bearer {token}", + "x-cursor-client-type": "cli", + }, + method="POST", + ) + with urllib.request.urlopen(request, timeout=15) as response: + return json.loads(response.read() or b"{}") + + +# The answer is protobuf JSON, which leaves out a field holding zero or +# false: no included limit means no Grok Bot allowance on this account, and +# a week with nothing used yet has no percentage at all. +def weekly_limits(status): + if status.get("hasNonZeroIncludedLimit") is not True: + return [] + resets = parse_time(status.get("nextResetTimestampUtc")) + return [{ + "label": "Weekly", + "percent": max(0.0, min(1.0, float(status.get("usagePercent") or 0) / 100)), + "resetsAt": resets.isoformat() if resets else "", + }] + + +def collect(): + record = { + "schemaVersion": 1, + "id": AGENT_ID, + "name": AGENT_NAME, + "updatedAt": datetime.now(timezone.utc).isoformat(), + "ready": False, + "hasLocalStats": False, + "hasPromptStats": False, + "tierLabel": "", + "usageStatusText": "", + "authHelpText": AUTH_HELP, + "limits": [], + "limitsStale": False, + "limitsFetchedAt": 0, + } + token = read_token() + if not token: + return record + + cache = cache_root() / f"grokbot-limits-{account_key(token)}.json" + try: + status = fetch_status(token) + limits = weekly_limits(status) + except urllib.error.HTTPError as error: + auth = error.code in (401, 403) + # Unreachable, timed out, or reset mid-read are all OSError; a body cut + # short is an HTTPException. + except (OSError, http.client.HTTPException, ValueError, TypeError, AttributeError): + auth = False + else: + fetched_at = round(datetime.now(timezone.utc).timestamp() * 1000) + tier = str(status.get("grokPlanLabel") or "") + record.update(ready=bool(limits), tierLabel=tier, authHelpText="", limits=limits, limitsFetchedAt=fetched_at) + # The cache only backs a later failed check; failing to keep it never + # costs this run's live answer. + try: + write_json(cache, {"fetchedAtMs": fetched_at, "tierLabel": tier, "limits": limits}) + except OSError: + pass + return record + + # A failed check keeps the last good limit until its week ends, marked + # stale so the panel dims it; one with no known end isn't kept at all. + # With nothing kept, the record stays empty, so a Cursor account that + # never had Grok Bot doesn't appear on a blip. + cached = read_json(cache) + now = datetime.now(timezone.utc) + limits = [limit for limit in cached.get("limits") or [] if (resets := parse_time(limit.get("resetsAt"))) and resets > now] + record.update( + ready=bool(limits), + tierLabel=str(cached.get("tierLabel") or ""), + usageStatusText="Waiting for auth" if auth else ("" if limits else "Grok Bot limits unavailable"), + limits=limits, + limitsStale=True, + limitsFetchedAt=int(cached.get("fetchedAtMs") or 0) if limits else 0, + ) + return record + + +def main(): + parser = argparse.ArgumentParser(description="Print the Grok Bot usage record as JSON") + # Every run is one live call, so there is no cache to force past and no + # scan to skip; the flags exist so every collector takes the same invocation. + 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__": + main() diff --git a/shell/plugins/agents/assets/grokbot-light.svg b/shell/plugins/agents/assets/grokbot-light.svg new file mode 100644 index 00000000000..92d54f32957 --- /dev/null +++ b/shell/plugins/agents/assets/grokbot-light.svg @@ -0,0 +1 @@ +Grok Bot diff --git a/shell/plugins/agents/assets/grokbot.svg b/shell/plugins/agents/assets/grokbot.svg new file mode 100644 index 00000000000..80eb13cad8d --- /dev/null +++ b/shell/plugins/agents/assets/grokbot.svg @@ -0,0 +1 @@ +Grok Bot diff --git a/test/shell.d/agent-usage-grokbot-test.sh b/test/shell.d/agent-usage-grokbot-test.sh new file mode 100755 index 00000000000..38798eab184 --- /dev/null +++ b/test/shell.d/agent-usage-grokbot-test.sh @@ -0,0 +1,152 @@ +#!/bin/bash + +set -euo pipefail + +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" + +require_command jq +require_command python3 + +test_tmp=$(mktemp -d) +trap 'rm -rf "$test_tmp"' EXIT + +export HOME="$test_tmp/home" +export XDG_CONFIG_HOME="$test_tmp/config" +export XDG_CACHE_HOME="$test_tmp/cache" +unset CURSOR_API_KEY + +reset_at=$(python3 -c 'import datetime as dt; print((dt.datetime.now(dt.timezone.utc) + dt.timedelta(days=3)).strftime("%Y-%m-%dT%H:%M:%S.%fZ"))') + +# A JWT whose subject names the account, as Cursor's sign-in tokens do. +jwt() { + printf 'h.%s.s' "$(printf '{"sub":"%s"}' "$1" | base64 -w0 | tr '+/' '-_' | tr -d '=')" +} + +cli_signed_in() { + mkdir -p "$XDG_CONFIG_HOME/cursor" + jq -n --arg token "$(jwt "$1")" '{accessToken: $token}' >"$XDG_CONFIG_HOME/cursor/auth.json" +} + +# ANSWER is the dashboard's JSON reply, an HTTP status code to fail with, or +# cut/reset for a connection that breaks mid-reply. With EXPECT set, any other +# token is refused, as Cursor would refuse it. +collect() { + COLLECTOR="$ROOT/bin/omarchy-agent-usage-grokbot" python3 - <<'PY' +import http.client, importlib.machinery, importlib.util, io, os, sys + +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) + +def urlopen(request, timeout=None): + assert request.full_url == collector.USAGE_URL + answer = os.environ["ANSWER"] + expected = os.environ.get("EXPECT") + if expected and request.get_header("Authorization") != f"Bearer {expected}": + answer = "401" + if answer == "cut": + raise http.client.IncompleteRead(b"{") + if answer == "reset": + raise ConnectionResetError() + if answer.isdigit(): + raise collector.urllib.error.HTTPError(request.full_url, int(answer), "", {}, None) + return io.BytesIO(answer.encode()) + +collector.urllib.request.urlopen = urlopen +sys.argv = ["omarchy-agent-usage-grokbot", "--limits-only"] +collector.main() +PY +} + +record=$(ANSWER='{}' collect) +[[ $(jq -c '{id, ready, limits}' <<<"$record") == '{"id":"grokbot","ready":false,"limits":[]}' ]] || + fail "a machine signed in to Cursor nowhere gives an empty record" "$record" +pass "a machine signed in to Cursor nowhere gives an empty record" + +cli_signed_in u-main +answer=$(jq -nc --arg at "$reset_at" '{usagePercent: 42.5, nextResetTimestampUtc: $at, hasNonZeroIncludedLimit: true, grokPlanLabel: "Pro"}') +record=$(ANSWER="$answer" collect) +[[ $(jq -c '{ready, tierLabel, stale: .limitsStale, label: .limits[0].label, percent: .limits[0].percent, help: .authHelpText}' <<<"$record") == '{"ready":true,"tierLabel":"Pro","stale":false,"label":"Weekly","percent":0.425,"help":""}' ]] || + fail "Grok Bot's weekly allowance comes from the Cursor sign-in" "$record" +python3 -c 'import datetime as dt, sys; d = dt.datetime.fromisoformat(sys.argv[1]) - dt.datetime.now(dt.timezone.utc); sys.exit(0 if dt.timedelta(days=2) < d < dt.timedelta(days=3) else 1)' "$(jq -r '.limits[0].resetsAt' <<<"$record")" || + fail "the weekly window says when it resets" "$record" +pass "Grok Bot's weekly allowance comes from the Cursor sign-in" + +# Failed checks keep the last good limit, dimmed, and say why when it's auth. +record=$(ANSWER=500 collect) +[[ $(jq -c '{ready, tierLabel, stale: .limitsStale, percent: .limits[0].percent, status: .usageStatusText}' <<<"$record") == '{"ready":true,"tierLabel":"Pro","stale":true,"percent":0.425,"status":""}' ]] || + fail "a failed check keeps the last limit, marked stale" "$record" +record=$(ANSWER=401 collect) +[[ $(jq -c '{stale: .limitsStale, percent: .limits[0].percent, status: .usageStatusText}' <<<"$record") == '{"stale":true,"percent":0.425,"status":"Waiting for auth"}' ]] || + fail "a refused sign-in keeps the last limit and says so" "$record" +for broken in cut reset; do + record=$(ANSWER=$broken collect) + [[ $(jq -c '{stale: .limitsStale, percent: .limits[0].percent}' <<<"$record") == '{"stale":true,"percent":0.425}' ]] || + fail "a connection that breaks mid-reply ($broken) keeps the last limit" "$record" +done +pass "failed checks keep the last limit, marked stale" + +# A cache that can't be written never costs the live answer. +chmod a-w "$XDG_CACHE_HOME/omarchy/agent-usage" +record=$(ANSWER="$answer" collect) +chmod u+w "$XDG_CACHE_HOME/omarchy/agent-usage" +[[ $(jq -c '{ready, stale: .limitsStale, percent: .limits[0].percent}' <<<"$record") == '{"ready":true,"stale":false,"percent":0.425}' ]] || + fail "an unwritable cache still gives the live limit" "$record" +pass "an unwritable cache still gives the live limit" + +# Another account's limit never stands in for this one's. +cli_signed_in u-other +record=$(ANSWER=500 collect) +[[ $(jq -c '{ready, limits}' <<<"$record") == '{"ready":false,"limits":[]}' ]] || + fail "another account's cached limit is never shown" "$record" +pass "another account's cached limit is never shown" + +# Protobuf JSON leaves out zero and false: an untouched week reads 0%, and an +# account with no included limit has no Grok Bot, so the panel skips it. +answer=$(jq -nc --arg at "$reset_at" '{nextResetTimestampUtc: $at, hasNonZeroIncludedLimit: true}') +record=$(ANSWER="$answer" collect) +[[ $(jq -c '.limits[0] | {label, percent}' <<<"$record") == '{"label":"Weekly","percent":0.0}' ]] || + fail "an untouched week reads as 0%" "$record" +record=$(ANSWER='{"usagePercent": 10}' collect) +[[ $(jq -c '{ready, limits}' <<<"$record") == '{"ready":false,"limits":[]}' ]] || + fail "a Cursor account without Grok Bot gives an empty record" "$record" +pass "an untouched week reads 0%, and no allowance means no Grok Bot" + +# A limit with no known end shows while live, but is never kept through a +# failed check, where it could stand for a week long gone. +cli_signed_in u-no-reset +record=$(ANSWER='{"usagePercent": 5, "hasNonZeroIncludedLimit": true}' collect) +[[ $(jq -c '{ready, percent: .limits[0].percent, resetsAt: .limits[0].resetsAt}' <<<"$record") == '{"ready":true,"percent":0.05,"resetsAt":""}' ]] || + fail "a limit with no known reset still shows while live" "$record" +record=$(ANSWER=500 collect) +[[ $(jq -c '{ready, limits}' <<<"$record") == '{"ready":false,"limits":[]}' ]] || + fail "a limit with no known reset is not kept through a failed check" "$record" +pass "a limit with no known reset is never kept through a failed check" + +# The same account the Cursor collector reads: an exported key, then the CLI's +# sign-in, then the editor's, with all three present and each asked with its +# own token. +cli_signed_in p-cli +state_db="$XDG_CONFIG_HOME/Cursor/User/globalStorage/state.vscdb" +mkdir -p "$(dirname "$state_db")" +python3 - "$state_db" "$(jwt p-editor)" <<'PY' +import sqlite3, sys +conn = sqlite3.connect(sys.argv[1]) +conn.execute("CREATE TABLE ItemTable (key TEXT UNIQUE ON CONFLICT REPLACE, value BLOB)") +conn.execute("INSERT INTO ItemTable VALUES ('cursorAuth/accessToken', ?)", (sys.argv[2],)) +conn.commit() +PY +answer=$(jq -nc --arg at "$reset_at" '{usagePercent: 7, nextResetTimestampUtc: $at, hasNonZeroIncludedLimit: true}') +live='{"ready":true,"stale":false,"percent":0.07}' +record=$(CURSOR_API_KEY="$(jwt p-key)" EXPECT="$(jwt p-key)" ANSWER="$answer" collect) +[[ $(jq -c '{ready, stale: .limitsStale, percent: .limits[0].percent}' <<<"$record") == "$live" ]] || + fail "an exported CURSOR_API_KEY wins over both sign-ins" "$record" +record=$(EXPECT="$(jwt p-cli)" ANSWER="$answer" collect) +[[ $(jq -c '{ready, stale: .limitsStale, percent: .limits[0].percent}' <<<"$record") == "$live" ]] || + fail "the Cursor CLI's sign-in wins over the editor's" "$record" +rm -f "$XDG_CONFIG_HOME/cursor/auth.json" +record=$(EXPECT="$(jwt p-editor)" ANSWER="$answer" collect) +[[ $(jq -c '{ready, stale: .limitsStale, percent: .limits[0].percent}' <<<"$record") == "$live" ]] || + fail "the Cursor editor's sign-in stands in for the CLI's" "$record" +pass "the sign-in is the one the Cursor collector reads, in the same order" From 5c162d67fd9d23e3ac64788e6991d5b7c15ba484 Mon Sep 17 00:00:00 2001 From: Oleksandr Anosov Date: Sat, 3 Oct 2026 21:48:22 +0300 Subject: [PATCH 040/546] Add Kimchi to the agents panel Kimchi's plan, prepaid credits, and spend budgets come from the gateway endpoints its own status line reads; tokens come from its pi session files. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-kimchi | 376 ++++++++++++++++++++++++ manual/17-ai.md | 2 +- shell/plugins/agents/README.md | 16 +- shell/plugins/agents/assets/kimchi.svg | 4 + shell/plugins/agents/manifest.json | 3 +- test/shell.d/agent-usage-kimchi-test.sh | 182 ++++++++++++ 6 files changed, 580 insertions(+), 3 deletions(-) create mode 100755 bin/omarchy-agent-usage-kimchi create mode 100644 shell/plugins/agents/assets/kimchi.svg create mode 100755 test/shell.d/agent-usage-kimchi-test.sh diff --git a/bin/omarchy-agent-usage-kimchi b/bin/omarchy-agent-usage-kimchi new file mode 100755 index 00000000000..1f9d268efe9 --- /dev/null +++ b/bin/omarchy-agent-usage-kimchi @@ -0,0 +1,376 @@ +#!/usr/bin/python3 +# omarchy:summary=Print the Kimchi usage record as JSON +# omarchy:args=[--force] [--limits-only] +# omarchy:hidden=true +"""Collect Kimchi usage into one display-ready JSON record. + +The plan, prepaid credits, and spend budgets come from the gateway endpoints +Kimchi's own status line reads (/v1/credits and /v1/budget), asked +with the API key Kimchi signs in with. Tokens and prompts come from its pi +session files: one line per message, each assistant reply carrying its model +and token usage. The agents panel only ever reads the JSON this prints. +""" + +import argparse +import hashlib +import json +import os +import tempfile +import time +import urllib.error +import urllib.request +from datetime import datetime, timedelta, timezone +from pathlib import Path + +AGENT_ID = "kimchi" +AGENT_NAME = "Kimchi" +AUTH_HELP = "Start Kimchi and log in, or set KIMCHI_API_KEY." +REGION_LLM_BASE = {"us": "https://llm.kimchi.dev", "eu": "https://llm.eu.kimchi.dev"} +PLAN_LABELS = {"community": "Community", "coder": "Coder", "teams": "Teams", "enterprise": "Enterprise"} +BUDGET_LABELS = { + "API_KEY": "API key budget", + "USER": "Budget", + "TEAM_PER_USER": "Team budget", + "TEAM_POOLED": "Team pool", + "ORGANIZATION_SOFT": "Org budget", + "ORGANIZATION_HARD": "Org budget", +} +# A limits-only refresh needs fresh credits, not a fresh history scan. +LIMITS_ONLY_SCAN_REUSE_SECONDS = 900 + + +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 cache_root(): + return Path(os.environ.get("XDG_CACHE_HOME") or (Path.home() / ".cache")) / "omarchy" / "agent-usage" + + +def number(value): + try: + return max(0, int(value or 0)) + except Exception: + return 0 + + +# Kimchi sends dollar amounts as strings, sometimes with a sign or unit. +def usd(value): + if isinstance(value, bool): + return None + if isinstance(value, (int, float)): + return float(value) + text = str(value or "").lower().replace("$", "").replace("usd", "").replace(",", ".").strip() + try: + return float(text) if text else None + except ValueError: + return None + + +# pi stamps entries with ISO strings and messages with epoch milliseconds. +def local_day(value): + if isinstance(value, (int, float)) and not isinstance(value, bool): + return datetime.fromtimestamp(value / 1000 if value > 10_000_000_000 else value).astimezone().date().isoformat() + try: + parsed = datetime.fromisoformat(str(value or "").replace("Z", "+00:00")) + except ValueError: + return None + return (parsed if parsed.tzinfo else parsed.replace(tzinfo=timezone.utc)).astimezone().date().isoformat() + + +# The harness directory Kimchi keeps its sessions under: KIMCHI_CODING_AGENT_DIR +# when set, as the CLI itself reads it, else ~/.config/kimchi/harness. +def harness_dir(): + return Path(os.environ.get("KIMCHI_CODING_AGENT_DIR") or (Path.home() / ".config" / "kimchi" / "harness")) + + +# Every session file under /sessions//. Only assistant +# messages carry usage, and a resumed session can repeat an entry, so each +# entry id counts once per file. +def local_stats(sessions_root): + now = datetime.now().astimezone() + today = now.date().isoformat() + week = [(now.date() - timedelta(days=offset)).isoformat() for offset in range(6, -1, -1)] + prompts = today_prompts = today_tokens = 0 + sessions, today_sessions, days = set(), set(), set() + by_day = {day: 0 for day in week} + today_by_model = {} + models = {} + files = sorted(Path(sessions_root).glob("*/*.jsonl")) + for path in files: + seen = set() + try: + lines = path.read_text(encoding="utf-8", errors="replace").splitlines() + except OSError: + continue + for line in lines: + if '"usage"' not in line: + continue + try: + entry = json.loads(line) + except ValueError: + continue + message = entry.get("message") if isinstance(entry, dict) and entry.get("type") == "message" else None + if not isinstance(message, dict) or message.get("role") != "assistant": + continue + entry_id = str(entry.get("id") or "") + if entry_id and entry_id in seen: + continue + seen.add(entry_id) + + usage = message.get("usage") or {} + uncached = number(usage.get("input")) + output = number(usage.get("output")) + cached = number(usage.get("cacheRead")) + created = number(usage.get("cacheWrite")) + if not (uncached or output or cached or created): + uncached = number(usage.get("totalTokens")) + tokens = uncached + output + cached + created + if not tokens: + continue + + day = local_day(entry.get("timestamp") or message.get("timestamp")) or today + model = str(message.get("model") or AGENT_ID) + prompts += 1 + sessions.add(path.name) + days.add(day) + totals = models.setdefault(model, {"inputTokens": 0, "outputTokens": 0, "cacheReadInputTokens": 0, "cacheCreationInputTokens": 0}) + totals["inputTokens"] += uncached + totals["outputTokens"] += output + totals["cacheReadInputTokens"] += cached + totals["cacheCreationInputTokens"] += created + if day in by_day: + by_day[day] += tokens + if day == today: + today_prompts += 1 + today_sessions.add(path.name) + today_tokens += tokens + today_by_model[model] = today_by_model.get(model, 0) + tokens + + return { + "hasLocalStats": bool(sessions), + "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 sessions as today's. +def cached_local_stats(sessions_root, max_age): + cache = cache_root() / "kimchi-stats.json" + cached = read_json(cache) or {} + today = datetime.now().astimezone().date().isoformat() + if cached.get("root") == str(sessions_root) and cached.get("day") == today and time.time() - float(cached.get("at") or 0) < max_age: + return cached["stats"] + stats = local_stats(sessions_root) + write_json(cache, {"root": str(sessions_root), "day": today, "at": time.time(), "stats": stats}) + return stats + + +# The key and gateway Kimchi itself would use: KIMCHI_API_KEY, KIMCHI_BASE_URL, +# and KIMCHI_REGION over ~/.config/kimchi/config.json, the region's gateway +# when nothing names one. +def gateway(): + config = read_json(Path.home() / ".config" / "kimchi" / "config.json") or {} + key = os.environ.get("KIMCHI_API_KEY") or config.get("apiKey") or config.get("api_key") or "" + base = os.environ.get("KIMCHI_BASE_URL") or config.get("llmEndpoint") or "" + if not base: + region = os.environ.get("KIMCHI_REGION") or config.get("region") + base = REGION_LLM_BASE.get(str(region or ""), REGION_LLM_BASE["us"]) + base = str(base).strip().rstrip("/") + if base.lower().endswith("/openai/v1"): + base = base[: -len("/openai/v1")] + return str(key).strip(), base + + +def get_json(url, key): + request = urllib.request.Request(url, headers={"Authorization": f"Bearer {key}", "Accept": "application/json", "User-Agent": "omarchy-agent-usage-kimchi"}) + with urllib.request.urlopen(request, timeout=8) as response: + return json.loads(response.read()) + + +def plan_label(raw): + text = str(raw or "").strip().lower().replace("_", "-").replace(" ", "-") + if text in ("community", "free", "free-tier", "free-slow"): + return PLAN_LABELS["community"] + if text in ("starter", "individual"): + return PLAN_LABELS["coder"] + if text == "team": + return PLAN_LABELS["teams"] + for plan in ("coder", "teams", "enterprise"): + if text == plan or text.startswith(plan + "-"): + return PLAN_LABELS[plan] + return str(raw or "") + + +# The credits answer names the plan and what's left of the prepaid balance. +# It says how much remains but not how much was bought, so the gauge has +# no funded figure unless a budget supplies one. +def apply_credits(result, data): + if data.get("serverless") is False: + return + tier = next((data[k] for k in ("tier", "tier_name", "tierName") if k in data), "") + result["tierLabel"] = plan_label(tier) + remaining = usd(data.get("remaining")) + if remaining is not None and remaining >= 0: + result["balance"] = {"remaining": remaining, "funded": 0, "spent": 0, "currency": "USD", "estimated": False} + has_credits = next((data[k] for k in ("has_credits", "hasCredits") if k in data), None) + if has_credits is False or (remaining is not None and remaining <= 0): + result["usageStatusText"] = "Out of credits, rate limited" + + +def parse_time(value): + try: + parsed = datetime.fromisoformat(str(value or "").replace("Z", "+00:00")) + except ValueError: + return None + return parsed if parsed.tzinfo else parsed.replace(tzinfo=timezone.utc) + + +# The panel titles a window by its clock, so a budget period that spans a +# month or a week says so rather than leaving the panel to guess from a label. +def period_title(period): + start, end = parse_time(period.get("startTime")), parse_time(period.get("endTime")) + if not (start and end): + return "" + days = (end - start).days + if 27 <= days <= 31: + return "Monthly" + if days == 7: + return "Weekly" + return "" + + +# Each spend budget is a limit window: spend against its cap until the +# period ends. A provider capped inside it is its own allowance on the same +# clock, which the panel draws as a tick on that budget's meter. A personal +# budget also gives the balance its funded figure. +def apply_budget(result, data): + period = data.get("period") if isinstance(data.get("period"), dict) else {} + resets = str(period.get("endTime") or "") + title = period_title(period) + for entry in data.get("budgets") or []: + if not isinstance(entry, dict): + continue + cap = usd(entry.get("budgetLimitUsd")) + spent = usd(entry.get("totalSpendUsd")) + if not cap or spent is None: + continue + limit = {"label": BUDGET_LABELS.get(entry.get("scope"), "Budget"), "percent": max(0.0, min(1.0, spent / cap))} + if title: + limit["title"] = title + if resets: + limit["resetsAt"] = resets + result["limits"].append(limit) + for provider in (entry.get("providerBudgets") or []) if title else []: + if not isinstance(provider, dict) or not str(provider.get("limitType") or "").endswith("CAPPED"): + continue + provider_cap = usd(provider.get("budgetLimitUsd")) + used = usd(provider.get("usageUsd")) + name = str(provider.get("provider") or "").strip() + if not (provider_cap and used is not None and name): + continue + scoped = {"label": f"{name} budget", "title": f"{name[:1].upper()}{name[1:]} {title}", "percent": max(0.0, min(1.0, used / provider_cap))} + if resets: + scoped["resetsAt"] = resets + result["limits"].append(scoped) + balance = result.get("balance") + if balance and not balance["funded"] and entry.get("scope") in ("USER", "API_KEY"): + balance.update(funded=cap, spent=spent) + + +# Keyed by the API key, so a machine signed in to someone else never shows +# the last one's numbers. +def billing_cache(key): + return cache_root() / f"kimchi-billing-{hashlib.sha256(key.encode()).hexdigest()[:16]}.json" + + +def collect_billing(): + result = {"tierLabel": "", "limits": [], "usageStatusText": "", "authHelpText": AUTH_HELP, "signedIn": False, "limitsStale": False, "limitsFetchedAt": 0} + key, base = gateway() + if not key: + return result + result["signedIn"] = True + reached = False + for resource, apply in (("credits", apply_credits), ("budget", apply_budget)): + try: + data = get_json(f"{base}/v1/{resource}", key) + except urllib.error.HTTPError as error: + if error.code in (401, 403): + result["usageStatusText"] = "Waiting for auth" + continue + except Exception: + continue + reached = True + if isinstance(data, dict): + apply(result, data) + + cache = billing_cache(key) + if reached: + result["authHelpText"] = "" + result["limitsFetchedAt"] = round(time.time() * 1000) + write_json(cache, {"fetchedAtMs": result["limitsFetchedAt"], "tierLabel": result["tierLabel"], "limits": result["limits"], "balance": result.get("balance")}) + return result + + # Nothing reached the gateway: the last answer stands, dimmed and dated, + # rather than a blank section. + cached = read_json(cache) or {} + if cached.get("fetchedAtMs"): + result.update(tierLabel=cached.get("tierLabel") or "", limits=cached.get("limits") or [], limitsStale=True, limitsFetchedAt=cached["fetchedAtMs"]) + if cached.get("balance"): + result["balance"] = cached["balance"] + if not result["usageStatusText"]: + # Usually the seconds after login before the network is up; the panel + # tries once more soon instead of waiting out its refresh interval. + if not cached.get("fetchedAtMs"): + result["usageStatusText"] = "Kimchi credits unavailable" + result["retryAdvised"] = True + return result + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--force", action="store_true") + parser.add_argument("--limits-only", action="store_true") + args = parser.parse_args() + + billing = collect_billing() + record = { + "schemaVersion": 1, + "id": AGENT_ID, + "name": AGENT_NAME, + "updatedAt": datetime.now(timezone.utc).isoformat(), + "ready": billing.pop("signedIn"), + } + record.update(billing) + record.update(cached_local_stats(harness_dir() / "sessions", 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/manual/17-ai.md b/manual/17-ai.md index 0af299655c9..61de87494f8 100644 --- a/manual/17-ai.md +++ b/manual/17-ai.md @@ -39,7 +39,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, a meter for each of its 5-hour session and weekly limits with the time until it resets (hover for the exact percentage), 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, Kimi, Cursor, GitHub Copilot, Muse, and Antigravity are covered out of the box with their 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: OpenCode Zen and Go, the Z.ai GLM Coding Plan, Ollama Cloud, and the MiniMax Token Plan with their limits, and Fireworks, OpenRouter, and DeepSeek with their prepaid balances. +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, a meter for each of its 5-hour session and weekly limits with the time until it resets (hover for the exact percentage), 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, Kimi, Cursor, GitHub Copilot, Muse, Antigravity, and Kimchi are covered out of the box with their 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: OpenCode Zen and Go, the Z.ai GLM Coding Plan, Ollama Cloud, and the MiniMax Token Plan with their limits, and Fireworks, OpenRouter, and DeepSeek with their prepaid balances. 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 26e69fb5e4f..23568287753 100644 --- a/shell/plugins/agents/README.md +++ b/shell/plugins/agents/README.md @@ -72,6 +72,7 @@ Adding an agent needs no QML changes: ship a collector that prints the record co | `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 | | `cursor` | Cursor's dashboard RPCs: included, auto-model, named-model, and on-demand meters | the same dashboard RPCs, one call per day for the last week | | `opencode` | OpenCode Go's usage endpoint (rolling + weekly + monthly), with an API key or the Console session | pi and omp sessions on the Zen and Go providers, plus opencode's own `message` and `session_message` stores | +| `kimchi` | The plan, prepaid credits, and any spend budgets from the gateway's `/v1/credits` and `/v1/budget`, the endpoints Kimchi's own status line reads | Kimchi's pi session files under `~/.config/kimchi/harness/sessions`: tokens by model per assistant reply | | `fireworks` | Estimated prepaid balance: configured funding minus rated account costs | Fireworks billing API, grouped by day and model for the last 30 days | | `copilot` | Account-wide AI credit allowance (GitHub quota endpoint, or estimated locally) | the Copilot CLI session store's `assistant_usage_events`, with `~/.copilot/session-state` transcripts as fallback | | `muse` | The client's key endpoint (session + weekly percents), else local metering only — never estimated | native `muse` CLI session files plus opencode sessions on the `meta` provider | @@ -238,6 +239,18 @@ 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. +### Kimchi + +Kimchi reads `KIMCHI_API_KEY` first, then the key in +`~/.config/kimchi/config.json`, and asks the gateway its region names (or +`KIMCHI_BASE_URL`); its sessions follow `KIMCHI_CODING_AGENT_DIR`. The +credits answer says how much remains but not how much was bought, so the +balance meter only fills in when a personal spend budget supplies the funded +figure. Every budget is a limit line that resets when its period ends, +titled Monthly or Weekly when the period is one, with each capped provider +inside it a tick on that line. When the gateway can't be reached, the last +answer stays up, dimmed and dated. + ### Fireworks balance The collector first asks the account's `:getBalance` endpoint for the real @@ -372,7 +385,8 @@ edit `shell.json` directly): omarchy bar set omarchy.agents providers '{ "claude": { "enabled": true }, "codex": { "enabled": false }, - "fireworks": { "enabled": true } + "fireworks": { "enabled": true }, + "kimchi": { "enabled": true } }' --json ``` diff --git a/shell/plugins/agents/assets/kimchi.svg b/shell/plugins/agents/assets/kimchi.svg new file mode 100644 index 00000000000..f2fb29de6eb --- /dev/null +++ b/shell/plugins/agents/assets/kimchi.svg @@ -0,0 +1,4 @@ + + + + diff --git a/shell/plugins/agents/manifest.json b/shell/plugins/agents/manifest.json index f969a7ebb64..3ebbe94d9dd 100644 --- a/shell/plugins/agents/manifest.json +++ b/shell/plugins/agents/manifest.json @@ -34,7 +34,8 @@ "ollama": { "enabled": true }, "minimax": { "enabled": true }, "openrouter": { "enabled": true }, - "deepseek": { "enabled": true } + "deepseek": { "enabled": true }, + "kimchi": { "enabled": true } }, "refreshIntervalSec": 900, "syncMode": "Off", diff --git a/test/shell.d/agent-usage-kimchi-test.sh b/test/shell.d/agent-usage-kimchi-test.sh new file mode 100755 index 00000000000..3dcd72d8acd --- /dev/null +++ b/test/shell.d/agent-usage-kimchi-test.sh @@ -0,0 +1,182 @@ +#!/bin/bash + +set -euo pipefail + +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" + +require_command jq +require_command python3 + +test_tmp=$(mktemp -d) +trap 'rm -rf "$test_tmp"' EXIT + +export HOME="$test_tmp/home" +export XDG_STATE_HOME="$test_tmp/state" +export XDG_CACHE_HOME="$test_tmp/cache" +unset KIMCHI_API_KEY KIMCHI_BASE_URL KIMCHI_REGION KIMCHI_CODING_AGENT_DIR + +period_end=$(python3 -c 'import datetime as dt; print((dt.datetime.now(dt.timezone.utc) + dt.timedelta(days=20)).isoformat())') +period_start=$(python3 -c 'import datetime as dt; print((dt.datetime.now(dt.timezone.utc) - dt.timedelta(days=11)).isoformat())') +export PERIOD_END="$period_end" PERIOD_START="$period_start" + +# A Kimchi login the way the CLI leaves it. +signed_in() { + mkdir -p "$HOME/.config/kimchi" + jq -n --arg key "$1" --arg region "${2:-us}" '{apiKey: $key, region: $region}' >"$HOME/.config/kimchi/config.json" +} + +# The gateway answers per key: credits for every key, a budget only for some, +# and nothing at all for a revoked one. Each request is logged so the test can +# check which gateway was asked. +collect() { + COLLECTOR="$ROOT/bin/omarchy-agent-usage-kimchi" REQUESTS="$test_tmp/requests" 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) + +credits = { + "key-coder": {"serverless": True, "tier": "coder", "is_paid_tier": True, "remaining": "24.95", "has_credits": True}, + "key-budget": {"serverless": True, "tier": "TEAMS", "remaining": 30, "has_credits": True}, + "key-empty": {"serverless": True, "tier": "free", "remaining": "0", "has_credits": False}, +} +budgets = { + "key-budget": {"period": {"startTime": os.environ["PERIOD_START"], "endTime": os.environ["PERIOD_END"]}, "budgets": [ + {"scope": "USER", "scopeId": "u-1", "budgetLimitUsd": "50", "totalSpendUsd": "20", "providerBudgets": [ + {"provider": "anthropic", "limitType": "PROVIDER_LIMIT_TYPE_CAPPED", "budgetLimitUsd": "10", "usageUsd": "2.5"}, + {"provider": "openai", "limitType": "UNLIMITED", "usageUsd": "4"}, + ]}, + ]}, +} + +def urlopen(request, timeout=None): + with open(os.environ["REQUESTS"], "a") as log: + log.write(request.full_url + "\n") + key = request.get_header("Authorization").split(" ", 1)[1] + if key == "key-offline" or os.environ.get("GATEWAY_DOWN"): + raise urllib.error.URLError("offline") + if key == "key-revoked": + raise urllib.error.HTTPError(request.full_url, 401, "Unauthorized", {}, None) + answers = credits if request.full_url.endswith("/v1/credits") else budgets + return io.BytesIO(json.dumps(answers.get(key, {"period": {"startTime": "", "endTime": ""}, "budgets": []})).encode()) + +collector.urllib.request.urlopen = urlopen +sys.argv = ["omarchy-agent-usage-kimchi"] + (os.environ.get("COLLECT_ARGS") or "--force").split() +collector.main() +PY +} + +record=$(collect) +[[ $(jq -c '{ready, tierLabel, limits, hasLocalStats, balance}' <<<"$record") == '{"ready":false,"tierLabel":"","limits":[],"hasLocalStats":false,"balance":null}' ]] || + fail "a machine nobody signed in to Kimchi on gives an empty record" "$record" +[[ ! -s $test_tmp/requests ]] || fail "no key means the gateway is never asked" "$(cat "$test_tmp/requests")" +pass "a machine nobody signed in to Kimchi on gives an empty record" + +signed_in key-coder +record=$(collect) +[[ $(jq -c '{ready, tierLabel, limits, balance, usageStatusText, authHelpText}' <<<"$record") == '{"ready":true,"tierLabel":"Coder","limits":[],"balance":{"remaining":24.95,"funded":0,"spent":0,"currency":"USD","estimated":false},"usageStatusText":"","authHelpText":""}' ]] || + fail "a Coder plan reports its plan and prepaid credits" "$record" +[[ $(sort -u "$test_tmp/requests") == $'https://llm.kimchi.dev/v1/budget\nhttps://llm.kimchi.dev/v1/credits' ]] || + fail "the US region asks the US gateway" "$(cat "$test_tmp/requests")" +pass "a Coder plan reports its plan and prepaid credits" + +: >"$test_tmp/requests" +signed_in key-coder eu +collect >/dev/null +[[ $(sort -u "$test_tmp/requests") == $'https://llm.eu.kimchi.dev/v1/budget\nhttps://llm.eu.kimchi.dev/v1/credits' ]] || + fail "the EU region asks the EU gateway" "$(cat "$test_tmp/requests")" +: >"$test_tmp/requests" +KIMCHI_BASE_URL="https://gateway.example/openai/v1/" collect >/dev/null +[[ $(sort -u "$test_tmp/requests") == $'https://gateway.example/v1/budget\nhttps://gateway.example/v1/credits' ]] || + fail "KIMCHI_BASE_URL names the gateway, without its OpenAI path" "$(cat "$test_tmp/requests")" +pass "the gateway follows Kimchi's region and KIMCHI_BASE_URL" + +record=$(KIMCHI_API_KEY=key-budget collect) +[[ $(jq -c '{tierLabel, limit: (.limits[0] | {label, title, percent}), gauge: (.balance | .funded == 50 and .spent == 20)}' <<<"$record") == '{"tierLabel":"Teams","limit":{"label":"Budget","title":"Monthly","percent":0.4},"gauge":true}' ]] || + fail "a monthly spend budget becomes a Monthly limit and funds the balance gauge" "$record" +[[ $(jq -c '[.limits[1:][] | {title, percent}]' <<<"$record") == '[{"title":"Anthropic Monthly","percent":0.25}]' ]] || + fail "a capped provider is a scoped allowance on the budget's clock, an uncapped one isn't" "$record" +[[ $(jq -r '.limits[0].resetsAt' <<<"$record") == "$period_end" ]] || fail "the budget resets when its period ends" "$record" +pass "a monthly spend budget becomes a Monthly limit with its provider caps" + +record=$(KIMCHI_API_KEY=key-empty collect) +[[ $(jq -c '{tierLabel, empty: (.balance.remaining == 0), usageStatusText}' <<<"$record") == '{"tierLabel":"Community","empty":true,"usageStatusText":"Out of credits, rate limited"}' ]] || + fail "an empty balance says Kimchi is rate limited" "$record" +pass "an empty balance says Kimchi is rate limited" + +record=$(KIMCHI_API_KEY=key-revoked collect) +[[ $(jq -c '{ready, usageStatusText, retry: .retryAdvised}' <<<"$record") == '{"ready":true,"usageStatusText":"Waiting for auth","retry":null}' ]] || + fail "a rejected key asks for a sign-in" "$record" +record=$(KIMCHI_API_KEY=key-offline collect) +[[ $(jq -c '{usageStatusText, retry: .retryAdvised}' <<<"$record") == '{"usageStatusText":"Kimchi credits unavailable","retry":true}' ]] || + fail "an unreachable gateway asks for a sooner retry" "$record" +pass "a rejected key and an unreachable gateway are told apart" + +# A gateway that stops answering leaves the last numbers up, dimmed and dated, +# rather than an empty section. +record=$(KIMCHI_API_KEY=key-budget collect) +fetched=$(jq -r '.limitsFetchedAt' <<<"$record") +(( fetched > 0 )) && [[ $(jq -r '.limitsStale' <<<"$record") == false ]] || fail "a live answer is dated and not stale" "$record" +record=$(KIMCHI_API_KEY=key-budget GATEWAY_DOWN=1 collect) +[[ $(jq -c '{tierLabel, stale: .limitsStale, at: .limitsFetchedAt, limits: (.limits | length), kept: (.balance.remaining == 30), usageStatusText, retry: .retryAdvised}' <<<"$record") == '{"tierLabel":"Teams","stale":true,"at":'"$fetched"',"limits":2,"kept":true,"usageStatusText":"","retry":true}' ]] || + fail "an outage keeps the last credits and budgets, marked stale" "$record" +record=$(KIMCHI_API_KEY=key-coder GATEWAY_DOWN=1 collect) +[[ $(jq -c '{tierLabel, limits}' <<<"$record") == '{"tierLabel":"Coder","limits":[]}' ]] || + fail "another key never shows the last key's numbers" "$record" +pass "an outage keeps the last credits and budgets, marked stale" + +# Tokens come from assistant messages in the pi session files: cached input +# kept apart, a repeated entry counted once, user turns and other entries +# skipped. +sessions="$HOME/.config/kimchi/harness/sessions/--home-me--" +mkdir -p "$sessions" +now=$(python3 -c 'import datetime as dt; print(dt.datetime.now(dt.timezone.utc).isoformat())') +now_ms=$(python3 -c 'import time; print(int(time.time() * 1000))') +reply() { + jq -nc --arg id "$1" --arg at "$2" --arg model "$3" --argjson usage "$4" \ + '{type: "message", id: $id, timestamp: $at, message: {role: "assistant", model: $model, provider: "kimchi-dev", usage: $usage}}' +} +{ + jq -nc '{type: "session", id: "s-1"}' + jq -nc --arg at "$now" '{type: "message", id: "u-1", timestamp: $at, message: {role: "user", content: "hi", usage: {input: 999}}}' + reply a-1 "$now" minimax-m3 '{"input": 100, "output": 20, "cacheRead": 800, "cacheWrite": 5, "totalTokens": 925}' + reply a-1 "$now" minimax-m3 '{"input": 100, "output": 20, "cacheRead": 800, "cacheWrite": 5, "totalTokens": 925}' + reply a-2 "$now" glm-5.3 '{"input": 0, "output": 0, "cacheRead": 0, "cacheWrite": 0, "totalTokens": 75}' + jq -nc --arg at "$now" '{type: "custom", id: "c-1", timestamp: $at, data: {usage: {input: 999}}}' +} >"$sessions/today.jsonl" +{ + reply b-1 "2026-01-02T10:00:00Z" minimax-m3 '{"input": 300, "output": 30, "cacheRead": 0, "cacheWrite": 0}' +} >"$sessions/old.jsonl" +jq -nc --argjson at "$now_ms" '{type: "message", id: "a-3", message: {role: "assistant", model: "glm-5.3", timestamp: $at, usage: {input: 10, output: 5}}}' >>"$sessions/today.jsonl" + +record=$(KIMCHI_API_KEY=key-coder collect) +[[ $(jq -c '{hasLocalStats, todayPrompts, todaySessions, totalPrompts, totalSessions, activeDays, todayTotalTokens, todayTokensByModel, today: .recentDays[6].messageCount}' <<<"$record") == '{"hasLocalStats":true,"todayPrompts":3,"todaySessions":1,"totalPrompts":4,"totalSessions":2,"activeDays":2,"todayTotalTokens":1015,"todayTokensByModel":{"minimax-m3":925,"glm-5.3":90},"today":1015}' ]] || + fail "Kimchi counts sessions, prompts, and tokens from its session files" "$record" +[[ $(jq -cS '.modelUsage' <<<"$record") == '{"glm-5.3":{"cacheCreationInputTokens":0,"cacheReadInputTokens":0,"inputTokens":85,"outputTokens":5},"minimax-m3":{"cacheCreationInputTokens":5,"cacheReadInputTokens":800,"inputTokens":400,"outputTokens":50}}' ]] || + fail "Kimchi's tokens by model keep cached input apart" "$record" +pass "Kimchi counts sessions, prompts, and tokens from its session files" + +# KIMCHI_CODING_AGENT_DIR moves the harness, as it does for the CLI. +mkdir -p "$test_tmp/harness/sessions/--elsewhere--" +reply x-1 "$now" kimi-k2.6 '{"input": 7, "output": 3}' >"$test_tmp/harness/sessions/--elsewhere--/s.jsonl" +record=$(KIMCHI_API_KEY=key-coder KIMCHI_CODING_AGENT_DIR="$test_tmp/harness" collect) +[[ $(jq -c '{totalSessions, todayTokensByModel}' <<<"$record") == '{"totalSessions":1,"todayTokensByModel":{"kimi-k2.6":10}}' ]] || + fail "KIMCHI_CODING_AGENT_DIR names the harness the sessions are read from" "$record" +pass "KIMCHI_CODING_AGENT_DIR names the harness the sessions are read from" + +# A limits-only refresh reuses the last scan rather than reading every session. +KIMCHI_API_KEY=key-coder collect >/dev/null +reply b-2 "$now" minimax-m3 '{"input": 1, "output": 1}' >"$sessions/later.jsonl" +record=$(KIMCHI_API_KEY=key-coder COLLECT_ARGS="--limits-only" collect) +[[ $(jq -r '.totalSessions' <<<"$record") == 2 ]] || fail "a limits-only refresh reuses the session scan" "$record" +pass "a limits-only refresh reuses the session scan" + +# A scan from another day is never reused, so yesterday's sessions don't +# count as today's after midnight. +jq '.day = "2000-01-01"' "$XDG_CACHE_HOME/omarchy/agent-usage/kimchi-stats.json" >"$test_tmp/stats.json" +mv "$test_tmp/stats.json" "$XDG_CACHE_HOME/omarchy/agent-usage/kimchi-stats.json" +record=$(KIMCHI_API_KEY=key-coder COLLECT_ARGS="--limits-only" collect) +[[ $(jq -r '.totalSessions' <<<"$record") == 3 ]] || fail "a scan from another day is made again" "$record" +pass "a scan from another day is made again" From 15c4b9f28d7290614ce07c416212c1a7ee6d8cb5 Mon Sep 17 00:00:00 2001 From: Oleksandr Anosov Date: Sat, 3 Oct 2026 22:42:58 +0300 Subject: [PATCH 041/546] Keep Kimchi's billing endpoints apart and its budgets out of funding Each endpoint now caches and falls back on its own, keyed by gateway and key; a spend cap no longer poses as purchased credit; a rejected key no longer offers the panel's sign-in, which can't sign in to Kimchi; and capped providers show in budget periods of any length. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-kimchi | 112 ++++++++++++++---------- shell/plugins/agents/README.md | 14 +-- test/shell.d/agent-usage-kimchi-test.sh | 41 +++++++-- 3 files changed, 108 insertions(+), 59 deletions(-) diff --git a/bin/omarchy-agent-usage-kimchi b/bin/omarchy-agent-usage-kimchi index 1f9d268efe9..8ad511390c8 100755 --- a/bin/omarchy-agent-usage-kimchi +++ b/bin/omarchy-agent-usage-kimchi @@ -229,19 +229,20 @@ def plan_label(raw): # The credits answer names the plan and what's left of the prepaid balance. -# It says how much remains but not how much was bought, so the gauge has -# no funded figure unless a budget supplies one. -def apply_credits(result, data): +# It says how much remains but not how much was bought, so the gauge shows +# the remaining figure alone. +def credits_part(data): if data.get("serverless") is False: - return + return {} tier = next((data[k] for k in ("tier", "tier_name", "tierName") if k in data), "") - result["tierLabel"] = plan_label(tier) + part = {"tierLabel": plan_label(tier)} remaining = usd(data.get("remaining")) if remaining is not None and remaining >= 0: - result["balance"] = {"remaining": remaining, "funded": 0, "spent": 0, "currency": "USD", "estimated": False} + part["balance"] = {"remaining": remaining, "funded": 0, "spent": 0, "currency": "USD", "estimated": False} has_credits = next((data[k] for k in ("has_credits", "hasCredits") if k in data), None) if has_credits is False or (remaining is not None and remaining <= 0): - result["usageStatusText"] = "Out of credits, rate limited" + part["usageStatusText"] = "Out of credits, rate limited" + return part def parse_time(value): @@ -268,12 +269,14 @@ def period_title(period): # Each spend budget is a limit window: spend against its cap until the # period ends. A provider capped inside it is its own allowance on the same -# clock, which the panel draws as a tick on that budget's meter. A personal -# budget also gives the balance its funded figure. -def apply_budget(result, data): +# clock, which the panel draws as a tick on that budget's meter; in a period +# the panel has no clock name for, it gets a line of its own. A cap is what +# may be spent, not what was bought, so it never funds the balance gauge. +def budget_part(data): period = data.get("period") if isinstance(data.get("period"), dict) else {} resets = str(period.get("endTime") or "") title = period_title(period) + limits = [] for entry in data.get("budgets") or []: if not isinstance(entry, dict): continue @@ -286,8 +289,8 @@ def apply_budget(result, data): limit["title"] = title if resets: limit["resetsAt"] = resets - result["limits"].append(limit) - for provider in (entry.get("providerBudgets") or []) if title else []: + limits.append(limit) + for provider in entry.get("providerBudgets") or []: if not isinstance(provider, dict) or not str(provider.get("limitType") or "").endswith("CAPPED"): continue provider_cap = usd(provider.get("budgetLimitUsd")) @@ -295,60 +298,75 @@ def apply_budget(result, data): name = str(provider.get("provider") or "").strip() if not (provider_cap and used is not None and name): continue - scoped = {"label": f"{name} budget", "title": f"{name[:1].upper()}{name[1:]} {title}", "percent": max(0.0, min(1.0, used / provider_cap))} + name = name[:1].upper() + name[1:] + scoped = {"label": f"{name} budget", "title": f"{name} {title}" if title else f"{name} budget", "percent": max(0.0, min(1.0, used / provider_cap))} if resets: scoped["resetsAt"] = resets - result["limits"].append(scoped) - balance = result.get("balance") - if balance and not balance["funded"] and entry.get("scope") in ("USER", "API_KEY"): - balance.update(funded=cap, spent=spent) + limits.append(scoped) + return {"limits": limits} -# Keyed by the API key, so a machine signed in to someone else never shows -# the last one's numbers. -def billing_cache(key): - return cache_root() / f"kimchi-billing-{hashlib.sha256(key.encode()).hexdigest()[:16]}.json" +# Keyed by the API key and the gateway it was asked of, so a machine signed +# in to someone else, or pointed at another region, never shows the last +# one's numbers. +def billing_cache(key, base): + digest = hashlib.sha256((base + "\n" + key).encode()).hexdigest()[:16] + return cache_root() / f"kimchi-billing-{digest}.json" +# The two endpoints answer on their own, and so does each fall back: a part +# that didn't come back is the last one that did, dimmed and dated, so one +# endpoint failing never blanks or overwrites what the other said. def collect_billing(): result = {"tierLabel": "", "limits": [], "usageStatusText": "", "authHelpText": AUTH_HELP, "signedIn": False, "limitsStale": False, "limitsFetchedAt": 0} key, base = gateway() if not key: return result result["signedIn"] = True - reached = False - for resource, apply in (("credits", apply_credits), ("budget", apply_budget)): + + cache = billing_cache(key, base) + cached = read_json(cache) or {} + parts = {} + rejected = unreachable = False + now = round(time.time() * 1000) + for resource, parse in (("credits", credits_part), ("budget", budget_part)): try: data = get_json(f"{base}/v1/{resource}", key) + parts[resource] = dict(parse(data if isinstance(data, dict) else {}), fetchedAtMs=now, fresh=True) except urllib.error.HTTPError as error: - if error.code in (401, 403): - result["usageStatusText"] = "Waiting for auth" - continue + rejected = rejected or error.code in (401, 403) except Exception: - continue - reached = True - if isinstance(data, dict): - apply(result, data) - - cache = billing_cache(key) - if reached: + unreachable = True + if resource not in parts and isinstance(cached.get(resource), dict) and cached[resource].get("fetchedAtMs"): + parts[resource] = dict(cached[resource], fresh=False) + + fresh = {resource: {k: v for k, v in part.items() if k != "fresh"} for resource, part in parts.items() if part["fresh"]} + if fresh: + write_json(cache, {**{k: v for k, v in cached.items() if k in ("credits", "budget")}, **fresh}) + + credits = parts.get("credits") or {} + result["tierLabel"] = credits.get("tierLabel") or "" + if credits.get("balance"): + result["balance"] = credits["balance"] + result["limits"] = (parts.get("budget") or {}).get("limits") or [] + if parts: + result["limitsStale"] = not all(part["fresh"] for part in parts.values()) + result["limitsFetchedAt"] = min(part["fetchedAtMs"] for part in parts.values()) + + # A rejected key can't be fixed from the panel: Kimchi signs in through + # its own CLI, not omarchy-agent-account-add, so this status deliberately + # isn't one the panel offers its Sign-in link for. + if rejected: + result["usageStatusText"] = "Kimchi key rejected" + elif credits.get("usageStatusText"): + result["usageStatusText"] = credits["usageStatusText"] + elif not parts: + result["usageStatusText"] = "Kimchi credits unavailable" + if parts and not rejected: result["authHelpText"] = "" - result["limitsFetchedAt"] = round(time.time() * 1000) - write_json(cache, {"fetchedAtMs": result["limitsFetchedAt"], "tierLabel": result["tierLabel"], "limits": result["limits"], "balance": result.get("balance")}) - return result - - # Nothing reached the gateway: the last answer stands, dimmed and dated, - # rather than a blank section. - cached = read_json(cache) or {} - if cached.get("fetchedAtMs"): - result.update(tierLabel=cached.get("tierLabel") or "", limits=cached.get("limits") or [], limitsStale=True, limitsFetchedAt=cached["fetchedAtMs"]) - if cached.get("balance"): - result["balance"] = cached["balance"] - if not result["usageStatusText"]: + if unreachable: # Usually the seconds after login before the network is up; the panel # tries once more soon instead of waiting out its refresh interval. - if not cached.get("fetchedAtMs"): - result["usageStatusText"] = "Kimchi credits unavailable" result["retryAdvised"] = True return result diff --git a/shell/plugins/agents/README.md b/shell/plugins/agents/README.md index 23568287753..ac8357b42ae 100644 --- a/shell/plugins/agents/README.md +++ b/shell/plugins/agents/README.md @@ -245,11 +245,15 @@ Kimchi reads `KIMCHI_API_KEY` first, then the key in `~/.config/kimchi/config.json`, and asks the gateway its region names (or `KIMCHI_BASE_URL`); its sessions follow `KIMCHI_CODING_AGENT_DIR`. The credits answer says how much remains but not how much was bought, so the -balance meter only fills in when a personal spend budget supplies the funded -figure. Every budget is a limit line that resets when its period ends, -titled Monthly or Weekly when the period is one, with each capped provider -inside it a tick on that line. When the gateway can't be reached, the last -answer stays up, dimmed and dated. +balance shows the remaining figure without a funded meter; a spend budget's +cap is what may be spent, not what was bought, so it never stands in. Every +budget is a limit line that resets when its period ends, titled Monthly or +Weekly when the period is one, with each capped provider inside it a tick on +that line (or a line of its own in a period of another length). Each +endpoint falls back on its own: when one can't be reached, its last answer +stays up, dimmed and dated, beside the other's fresh one. A rejected key +reads as such rather than offering the panel's sign-in, since Kimchi signs +in through its own CLI. ### Fireworks balance diff --git a/test/shell.d/agent-usage-kimchi-test.sh b/test/shell.d/agent-usage-kimchi-test.sh index 3dcd72d8acd..be4853d17d2 100755 --- a/test/shell.d/agent-usage-kimchi-test.sh +++ b/test/shell.d/agent-usage-kimchi-test.sh @@ -41,8 +41,14 @@ credits = { "key-coder": {"serverless": True, "tier": "coder", "is_paid_tier": True, "remaining": "24.95", "has_credits": True}, "key-budget": {"serverless": True, "tier": "TEAMS", "remaining": 30, "has_credits": True}, "key-empty": {"serverless": True, "tier": "free", "remaining": "0", "has_credits": False}, + "key-sprint": {"serverless": True, "tier": "teams", "remaining": 12, "has_credits": True}, } budgets = { + "key-sprint": {"period": {"startTime": "2026-10-01T00:00:00Z", "endTime": "2026-10-15T00:00:00Z"}, "budgets": [ + {"scope": "TEAM_POOLED", "scopeId": "t-1", "budgetLimitUsd": "100", "totalSpendUsd": "10", "providerBudgets": [ + {"provider": "openai", "limitType": "CAPPED", "budgetLimitUsd": "20", "usageUsd": "5"}, + ]}, + ]}, "key-budget": {"period": {"startTime": os.environ["PERIOD_START"], "endTime": os.environ["PERIOD_END"]}, "budgets": [ {"scope": "USER", "scopeId": "u-1", "budgetLimitUsd": "50", "totalSpendUsd": "20", "providerBudgets": [ {"provider": "anthropic", "limitType": "PROVIDER_LIMIT_TYPE_CAPPED", "budgetLimitUsd": "10", "usageUsd": "2.5"}, @@ -55,7 +61,8 @@ def urlopen(request, timeout=None): with open(os.environ["REQUESTS"], "a") as log: log.write(request.full_url + "\n") key = request.get_header("Authorization").split(" ", 1)[1] - if key == "key-offline" or os.environ.get("GATEWAY_DOWN"): + down = os.environ.get("GATEWAY_DOWN", "") + if key == "key-offline" or down == "all" or (down and request.full_url.endswith("/v1/" + down)): raise urllib.error.URLError("offline") if key == "key-revoked": raise urllib.error.HTTPError(request.full_url, 401, "Unauthorized", {}, None) @@ -94,11 +101,14 @@ KIMCHI_BASE_URL="https://gateway.example/openai/v1/" collect >/dev/null pass "the gateway follows Kimchi's region and KIMCHI_BASE_URL" record=$(KIMCHI_API_KEY=key-budget collect) -[[ $(jq -c '{tierLabel, limit: (.limits[0] | {label, title, percent}), gauge: (.balance | .funded == 50 and .spent == 20)}' <<<"$record") == '{"tierLabel":"Teams","limit":{"label":"Budget","title":"Monthly","percent":0.4},"gauge":true}' ]] || - fail "a monthly spend budget becomes a Monthly limit and funds the balance gauge" "$record" +[[ $(jq -c '{tierLabel, limit: (.limits[0] | {label, title, percent}), unfunded: (.balance | .funded == 0 and .spent == 0 and .remaining == 30)}' <<<"$record") == '{"tierLabel":"Teams","limit":{"label":"Budget","title":"Monthly","percent":0.4},"unfunded":true}' ]] || + fail "a monthly spend budget becomes a Monthly limit, and its cap never poses as funding" "$record" [[ $(jq -c '[.limits[1:][] | {title, percent}]' <<<"$record") == '[{"title":"Anthropic Monthly","percent":0.25}]' ]] || fail "a capped provider is a scoped allowance on the budget's clock, an uncapped one isn't" "$record" [[ $(jq -r '.limits[0].resetsAt' <<<"$record") == "$period_end" ]] || fail "the budget resets when its period ends" "$record" +record=$(KIMCHI_API_KEY=key-sprint collect) +[[ $(jq -c '[.limits[] | {title, percent}]' <<<"$record") == '[{"title":null,"percent":0.1},{"title":"Openai budget","percent":0.25}]' ]] || + fail "a period with no clock name still shows its capped providers" "$record" pass "a monthly spend budget becomes a Monthly limit with its provider caps" record=$(KIMCHI_API_KEY=key-empty collect) @@ -107,8 +117,8 @@ record=$(KIMCHI_API_KEY=key-empty collect) pass "an empty balance says Kimchi is rate limited" record=$(KIMCHI_API_KEY=key-revoked collect) -[[ $(jq -c '{ready, usageStatusText, retry: .retryAdvised}' <<<"$record") == '{"ready":true,"usageStatusText":"Waiting for auth","retry":null}' ]] || - fail "a rejected key asks for a sign-in" "$record" +[[ $(jq -c '{ready, usageStatusText, retry: .retryAdvised}' <<<"$record") == '{"ready":true,"usageStatusText":"Kimchi key rejected","retry":null}' ]] || + fail "a rejected key says so without the panel's sign-in link, which can't sign in to Kimchi" "$record" record=$(KIMCHI_API_KEY=key-offline collect) [[ $(jq -c '{usageStatusText, retry: .retryAdvised}' <<<"$record") == '{"usageStatusText":"Kimchi credits unavailable","retry":true}' ]] || fail "an unreachable gateway asks for a sooner retry" "$record" @@ -119,14 +129,31 @@ pass "a rejected key and an unreachable gateway are told apart" record=$(KIMCHI_API_KEY=key-budget collect) fetched=$(jq -r '.limitsFetchedAt' <<<"$record") (( fetched > 0 )) && [[ $(jq -r '.limitsStale' <<<"$record") == false ]] || fail "a live answer is dated and not stale" "$record" -record=$(KIMCHI_API_KEY=key-budget GATEWAY_DOWN=1 collect) +record=$(KIMCHI_API_KEY=key-budget GATEWAY_DOWN=all collect) [[ $(jq -c '{tierLabel, stale: .limitsStale, at: .limitsFetchedAt, limits: (.limits | length), kept: (.balance.remaining == 30), usageStatusText, retry: .retryAdvised}' <<<"$record") == '{"tierLabel":"Teams","stale":true,"at":'"$fetched"',"limits":2,"kept":true,"usageStatusText":"","retry":true}' ]] || fail "an outage keeps the last credits and budgets, marked stale" "$record" -record=$(KIMCHI_API_KEY=key-coder GATEWAY_DOWN=1 collect) +record=$(KIMCHI_API_KEY=key-coder GATEWAY_DOWN=all collect) [[ $(jq -c '{tierLabel, limits}' <<<"$record") == '{"tierLabel":"Coder","limits":[]}' ]] || fail "another key never shows the last key's numbers" "$record" pass "an outage keeps the last credits and budgets, marked stale" +# One endpoint failing keeps that endpoint's last answer and takes the other's +# fresh one, without letting the half that came back replace the cache. +record=$(KIMCHI_API_KEY=key-budget collect) +record=$(KIMCHI_API_KEY=key-budget GATEWAY_DOWN=budget collect) +[[ $(jq -c '{stale: .limitsStale, limits: (.limits | length), credits: (.balance.remaining == 30), retry: .retryAdvised}' <<<"$record") == '{"stale":true,"limits":2,"credits":true,"retry":true}' ]] || + fail "a failed budget endpoint keeps its last limits beside fresh credits" "$record" +record=$(KIMCHI_API_KEY=key-budget GATEWAY_DOWN=credits collect) +[[ $(jq -c '{stale: .limitsStale, tierLabel, limits: (.limits | length), credits: (.balance.remaining == 30)}' <<<"$record") == '{"stale":true,"tierLabel":"Teams","limits":2,"credits":true}' ]] || + fail "a failed credits endpoint keeps its last plan and balance beside fresh limits" "$record" +pass "one endpoint failing keeps its own last answer" + +# The same key asked of another gateway is another answer. +record=$(KIMCHI_API_KEY=key-budget KIMCHI_BASE_URL="https://gateway.example" GATEWAY_DOWN=all collect) +[[ $(jq -c '{tierLabel, limits, usageStatusText}' <<<"$record") == '{"tierLabel":"","limits":[],"usageStatusText":"Kimchi credits unavailable"}' ]] || + fail "another gateway never shows the last gateway's numbers" "$record" +pass "another gateway never shows the last gateway's numbers" + # Tokens come from assistant messages in the pi session files: cached input # kept apart, a repeated entry counted once, user turns and other entries # skipped. From 4195d789a931e1f9ea40756857706b6f7ef3abd1 Mon Sep 17 00:00:00 2001 From: Oleksandr Anosov Date: Sun, 4 Oct 2026 19:07:46 +0300 Subject: [PATCH 042/546] Mark missing Kimchi billing stale and keep the latest answer An endpoint that fails with nothing cached now marks the record stale and asks for a sooner retry, server errors included. Each endpoint caches to its own file and only replaces an answer fetched before it, so overlapping refreshes can't roll the cache back. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-kimchi | 52 +++++++++++++++---------- test/shell.d/agent-usage-kimchi-test.sh | 27 +++++++++++++ 2 files changed, 58 insertions(+), 21 deletions(-) diff --git a/bin/omarchy-agent-usage-kimchi b/bin/omarchy-agent-usage-kimchi index 8ad511390c8..9a7da373392 100755 --- a/bin/omarchy-agent-usage-kimchi +++ b/bin/omarchy-agent-usage-kimchi @@ -308,15 +308,25 @@ def budget_part(data): # Keyed by the API key and the gateway it was asked of, so a machine signed # in to someone else, or pointed at another region, never shows the last -# one's numbers. -def billing_cache(key, base): +# one's numbers. One file per endpoint, so a run only ever writes the answer +# it actually got. +def billing_cache(key, base, resource): digest = hashlib.sha256((base + "\n" + key).encode()).hexdigest()[:16] - return cache_root() / f"kimchi-billing-{digest}.json" + return cache_root() / f"kimchi-{resource}-{digest}.json" + + +# Refreshes can overlap (the panel's own and a limits retry), so an answer +# only replaces one fetched after it. +def remember(path, part): + current = read_json(path) or {} + if float(current.get("fetchedAtMs") or 0) <= part["fetchedAtMs"]: + write_json(path, part) # The two endpoints answer on their own, and so does each fall back: a part # that didn't come back is the last one that did, dimmed and dated, so one -# endpoint failing never blanks or overwrites what the other said. +# endpoint failing never blanks or overwrites what the other said. Any part +# that didn't come back fresh marks the record stale, cached or not. def collect_billing(): result = {"tierLabel": "", "limits": [], "usageStatusText": "", "authHelpText": AUTH_HELP, "signedIn": False, "limitsStale": False, "limitsFetchedAt": 0} key, base = gateway() @@ -324,33 +334,32 @@ def collect_billing(): return result result["signedIn"] = True - cache = billing_cache(key, base) - cached = read_json(cache) or {} parts = {} - rejected = unreachable = False - now = round(time.time() * 1000) + failed = set() + rejected = False for resource, parse in (("credits", credits_part), ("budget", budget_part)): + cache = billing_cache(key, base, resource) try: data = get_json(f"{base}/v1/{resource}", key) - parts[resource] = dict(parse(data if isinstance(data, dict) else {}), fetchedAtMs=now, fresh=True) + parts[resource] = dict(parse(data if isinstance(data, dict) else {}), fetchedAtMs=round(time.time() * 1000)) + remember(cache, parts[resource]) + continue except urllib.error.HTTPError as error: rejected = rejected or error.code in (401, 403) except Exception: - unreachable = True - if resource not in parts and isinstance(cached.get(resource), dict) and cached[resource].get("fetchedAtMs"): - parts[resource] = dict(cached[resource], fresh=False) - - fresh = {resource: {k: v for k, v in part.items() if k != "fresh"} for resource, part in parts.items() if part["fresh"]} - if fresh: - write_json(cache, {**{k: v for k, v in cached.items() if k in ("credits", "budget")}, **fresh}) + pass + failed.add(resource) + cached = read_json(cache) + if isinstance(cached, dict) and cached.get("fetchedAtMs"): + parts[resource] = cached credits = parts.get("credits") or {} result["tierLabel"] = credits.get("tierLabel") or "" if credits.get("balance"): result["balance"] = credits["balance"] result["limits"] = (parts.get("budget") or {}).get("limits") or [] + result["limitsStale"] = bool(failed) if parts: - result["limitsStale"] = not all(part["fresh"] for part in parts.values()) result["limitsFetchedAt"] = min(part["fetchedAtMs"] for part in parts.values()) # A rejected key can't be fixed from the panel: Kimchi signs in through @@ -360,13 +369,14 @@ def collect_billing(): result["usageStatusText"] = "Kimchi key rejected" elif credits.get("usageStatusText"): result["usageStatusText"] = credits["usageStatusText"] - elif not parts: + elif "credits" not in parts: result["usageStatusText"] = "Kimchi credits unavailable" if parts and not rejected: result["authHelpText"] = "" - if unreachable: - # Usually the seconds after login before the network is up; the panel - # tries once more soon instead of waiting out its refresh interval. + if failed and not rejected: + # Usually the seconds after login before the network is up, or a + # gateway hiccup; the panel tries once more soon instead of waiting out + # its refresh interval. result["retryAdvised"] = True return result diff --git a/test/shell.d/agent-usage-kimchi-test.sh b/test/shell.d/agent-usage-kimchi-test.sh index be4853d17d2..1aeb5ddf61f 100755 --- a/test/shell.d/agent-usage-kimchi-test.sh +++ b/test/shell.d/agent-usage-kimchi-test.sh @@ -64,6 +64,8 @@ def urlopen(request, timeout=None): down = os.environ.get("GATEWAY_DOWN", "") if key == "key-offline" or down == "all" or (down and request.full_url.endswith("/v1/" + down)): raise urllib.error.URLError("offline") + if os.environ.get("GATEWAY_ERROR") and request.full_url.endswith("/v1/" + os.environ["GATEWAY_ERROR"]): + raise urllib.error.HTTPError(request.full_url, 500, "Internal Server Error", {}, None) if key == "key-revoked": raise urllib.error.HTTPError(request.full_url, 401, "Unauthorized", {}, None) answers = credits if request.full_url.endswith("/v1/credits") else budgets @@ -148,6 +150,31 @@ record=$(KIMCHI_API_KEY=key-budget GATEWAY_DOWN=credits collect) fail "a failed credits endpoint keeps its last plan and balance beside fresh limits" "$record" pass "one endpoint failing keeps its own last answer" +# A part that fails before anything was cached is missing, not current: the +# record is stale and asks for a sooner retry, a server error included. +rm -f "$XDG_CACHE_HOME"/omarchy/agent-usage/kimchi-budget-*.json +record=$(KIMCHI_API_KEY=key-sprint GATEWAY_ERROR=budget collect) +[[ $(jq -c '{stale: .limitsStale, limits, credits: (.balance.remaining == 12), retry: .retryAdvised, usageStatusText}' <<<"$record") == '{"stale":true,"limits":[],"credits":true,"retry":true,"usageStatusText":""}' ]] || + fail "a budget endpoint failing with nothing cached leaves the record stale and retrying" "$record" +rm -f "$XDG_CACHE_HOME"/omarchy/agent-usage/kimchi-credits-*.json +record=$(KIMCHI_API_KEY=key-sprint GATEWAY_ERROR=credits collect) +[[ $(jq -c '{stale: .limitsStale, balance, usageStatusText, retry: .retryAdvised}' <<<"$record") == '{"stale":true,"balance":null,"usageStatusText":"Kimchi credits unavailable","retry":true}' ]] || + fail "a credits endpoint failing with nothing cached says credits are unavailable" "$record" +pass "a part that fails before anything was cached is missing, not current" + +# Refreshes can overlap; an answer fetched earlier never replaces one fetched +# later, so a slow run finishing last can't roll the cache back. +KIMCHI_API_KEY=key-budget collect >/dev/null +credits_cache=$(ls -t "$XDG_CACHE_HOME"/omarchy/agent-usage/kimchi-credits-*.json | head -1) +later=$(python3 -c 'import time; print(int(time.time() * 1000) + 60000)') +jq --argjson at "$later" '.fetchedAtMs = $at | .tierLabel = "Enterprise"' "$credits_cache" >"$test_tmp/credits.json" +mv "$test_tmp/credits.json" "$credits_cache" +KIMCHI_API_KEY=key-budget collect >/dev/null +[[ $(jq -r '.tierLabel' "$credits_cache") == "Enterprise" ]] || fail "an earlier answer never replaces a later one" "$(cat "$credits_cache")" +record=$(KIMCHI_API_KEY=key-budget GATEWAY_DOWN=all collect) +[[ $(jq -r '.tierLabel' <<<"$record") == "Enterprise" ]] || fail "an outage falls back to the latest answer" "$record" +pass "overlapping refreshes keep the latest answer" + # The same key asked of another gateway is another answer. record=$(KIMCHI_API_KEY=key-budget KIMCHI_BASE_URL="https://gateway.example" GATEWAY_DOWN=all collect) [[ $(jq -c '{tierLabel, limits, usageStatusText}' <<<"$record") == '{"tierLabel":"","limits":[],"usageStatusText":"Kimchi credits unavailable"}' ]] || From a3ae14f3258ada381bd53f18ba4bcb940c50ce7e Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 22:48:58 -0400 Subject: [PATCH 043/546] Document the contract every agent usage collector follows Co-Authored-By: Claude Opus 5.5 --- docs/agent-usage-collectors.md | 71 ++++++++++++++++++++++++++++++++++ 1 file changed, 71 insertions(+) create mode 100644 docs/agent-usage-collectors.md diff --git a/docs/agent-usage-collectors.md b/docs/agent-usage-collectors.md new file mode 100644 index 00000000000..d7067e58395 --- /dev/null +++ b/docs/agent-usage-collectors.md @@ -0,0 +1,71 @@ +# Agent Usage Collectors + +Every agent or provider in the agents panel is a collector: `bin/omarchy-agent-usage-`, a Python script that prints one JSON record. `omarchy-agent-usage-update` finds collectors by globbing `bin/omarchy-agent-usage-*`, runs them, and the panel (`shell/plugins/agents/Main.qml`, `Panel.qml`) reads the records. This is the contract every collector follows, so that each agent behaves the same in the panel. `bin/omarchy-agent-usage-claude` is the reference implementation. + +## The command + +- `#!/usr/bin/python3`, two-space indentation, and the header lines `# omarchy:summary=...`, `# omarchy:args=[--force] [--limits-only]`, and `# omarchy:hidden=true`. +- argparse accepts `--force` and `--limits-only`, and both mean something: + - A normal run reuses a local scan for `SCAN_REUSE_SECONDS = 20` and a limits probe for `PROBE_MIN_INTERVAL_SECONDS = 15`. + - `--limits-only` (the panel's near-limit refresh) reuses a local scan for `LIMITS_ONLY_REUSE_SECONDS = 900`, but asks for limits like a normal run. + - `--force` (a person pressing refresh) reuses nothing. +- The record is printed on stdout. Diagnostics go to stderr and never carry a credential. +- A collector never writes to an agent's own files. Credentials are only read, and a token is refreshed only where the agent's own CLI does the same under its lock, never by rewriting another harness's store. + +## The record + +| Field | Meaning | +|---|---| +| `schemaVersion` | `1` | +| `id`, `name` | The id matches the filename; `name` is the product name | +| `updatedAt` | ISO time the record was made | +| `ready` | `true` once the agent is in use on this machine (a credential or local usage) | +| `tierLabel` | The plan as the service names it, or `""`; never a placeholder or a guess. `"Prepaid"` is special to the panel | +| `usageStatusText`, `authHelpText` | See [Status](#status) | +| `limits` | `[{label, percent, resetsAt, title?}]`, `percent` in 0..1, `resetsAt` ISO or `""`, never `null`. Add `title` (Session, Weekly, Monthly) when the label doesn't parse as one | +| `limitsStale`, `limitsFetchedAt` | Whether `limits` were kept past a failed check, and when they were measured (epoch ms, `0` when never) | +| `retryAdvised` | `true` only when the service couldn't be reached at all (DNS, connection, timeout), so the panel retries this agent in 30 seconds | +| `balance` | `{remaining, funded, spent, currency, estimated}` for prepaid accounts | +| `accounts`, `accountSwitch` | Only when two or more accounts are registered; each account carries `limits`, `stale`, `fetchedAt`, `usageStatusText`, `authHelpText` | +| `hasLocalStats` | Whether the stats below come from this machine | +| `hasPromptStats` | `false` when the source has no real prompt or session counts (an API that counts requests, not prompts) | +| `scope` | `"account"` when the stats are account-wide rather than this machine's | +| Stats | `todayPrompts`, `todaySessions`, `todayTotalTokens`, `todayTokensByModel`, `recentDays` (seven `{date, messageCount}`, where `messageCount` holds tokens), `totalPrompts`, `totalSessions`, `activeDays`, `activeDates` (every date), and `modelUsage` buckets of `inputTokens`, `outputTokens`, `cacheReadInputTokens`, `cacheCreationInputTokens` | + +The full stats block is always present, zeroed when there is nothing. An agent the machine doesn't use prints `ready: false` with empty stats, no limits and no balance, and the panel skips it. That's the normal answer for a missing credential on an agent with no local usage; it isn't an error. + +## Status + +The panel acts on a small vocabulary in `usageStatusText`, so the same situation reads the same for every agent: + +| Status | When | +|---|---| +| `""` | Healthy, or a failed check that kept the last good limits (`limitsStale` says the rest) | +| `"Waiting for auth"` | The agent is in use but has no credential for its limits | +| `"Sign-in expired"` | The credential was refused, or a refresh token lapsed; signing in again is needed | +| `"Limits paused"` | An access token lapsed that the agent's own CLI renews the next time it runs; signing in again is not needed | +| `" limits unavailable"` | A check failed with nothing kept to show | +| `"Balance unavailable"` | A prepaid balance couldn't be read and none was kept | + +The detail goes in `authHelpText`: what failed and the one thing to do about it. It is `""` when healthy. A plan the service doesn't offer limits for (pay-as-you-go, no Go plan, a free tier) isn't trouble: the status stays `""` and local stats show alone. Any other status shows under the name in the urgent color, so collectors don't invent new ones. + +## Limits + +- Each credential or account caches its own last good answer at `cache_root() / "-limits-.json"`, `{fetchedAtMs, limits, ...}`, where `cache_root()` is `$XDG_CACHE_HOME/omarchy/agent-usage` and `` is a short hash of who the credential signs in as (a JWT subject, an email, an account id) or of the credential itself. Another sign-in never shows this one's numbers. +- The cache is read whatever its file time, so a clock that went back doesn't lose it, and every age check is `0 <= age < N`. +- A failed check keeps the cached limits with `limitsStale: true` and their `limitsFetchedAt`. The panel shows their age on hover. +- A kept window whose `resetsAt` has passed comes back at `percent: 0` with `resetsAt: ""`, never dropped: a reset allowance is untouched, not unknown. A window with no reset time is kept as it is. Reset times are only ever the service's own, never computed. +- A meter for a balance or allowance is a real ratio of the service's numbers, never estimated from local usage unless its label says `(est.)`. + +## Local stats + +- Tokens are counted once, under the provider that bills them. pi, omp, and opencode sessions are split by provider into the matching collector; there is no harness tab that counts them again. +- pi and omp sessions live under `~/.pi/agent/sessions` and `~/.omp/agent/sessions`, plus `profiles//agent/sessions` under each, and `PI_CODING_AGENT_DIR` moves pi's. They are found with `rg --no-ignore`, since a `$HOME` that is a git checkout would otherwise hide them. A message is deduplicated by `(id, timestamp)` when it has both, so a forked session doesn't count its parent's messages again, and a line without an id always counts. +- Timestamps and token fields from disk are untrusted: a zero, huge, negative, `NaN`, boolean or non-numeric value skips that field or row, never the record. +- A scan is cached per day (`scanDate`) and rebuilt under `flock`, and a scan that couldn't read all its sources isn't cached. + +## Assets, docs and tests + +- `shell/plugins/agents/assets/.svg`, plus `-light.svg` when the mark is white and needs a dark twin for light surfaces. +- A provider entry in `shell/plugins/agents/manifest.json`, a row in the README's collector table, and a `###` section in the README naming where its credentials and data come from. +- `test/shell.d/agent-usage--*-test.sh` covering at least: the record an unused machine prints, fresh limits, stale limits after a failed check, an elapsed kept window at zero, the probe reuse and `--force`, and `retryAdvised` on an unreachable service. From 473739e1690ac1196a4087229020732fdfa2dd65 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 22:55:41 -0400 Subject: [PATCH 044/546] Offer the panel's sign-in link only to agents it can sign in The link runs omarchy-agent-account-add, which knows Claude, Codex, and Grok. Any other agent reporting "Sign-in expired" or "Waiting for auth" got a link that signed the default agent in instead; it now keeps the status as text. Co-authored-by: Josh LeBlanc Co-Authored-By: Claude Opus 5.5 --- shell/plugins/agents/Panel.qml | 11 ++++++++++- test/shell.d/agents-panel-test.sh | 14 ++++++++++++++ 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/shell/plugins/agents/Panel.qml b/shell/plugins/agents/Panel.qml index 330a7de3296..88ada928fc3 100644 --- a/shell/plugins/agents/Panel.qml +++ b/shell/plugins/agents/Panel.qml @@ -421,7 +421,16 @@ Panel { // A lapsed or missing sign-in is something you can fix from here; any other // trouble is only reported. + // Only the agents `omarchy-agent-account-add` signs in get the link; any + // other agent's sign-in trouble stays text, with its own help to follow. + function supportsSignIn(p) { + for (var i = 0; i < addProviders.length; i++) + if (p && addProviders[i].providerId === p.providerId) return true + return false + } + function needsSignIn(item) { + if (item && item.providerId && !supportsSignIn(item)) return false var status = String(item && item.usageStatusText || "") return status === "Sign-in expired" || status === "Waiting for auth" } @@ -438,7 +447,7 @@ Panel { // Sign an account that's already here in again, following along in the // add view just like adding one. function signInAgain(p, account) { - if (!p || addStage === "running") return + if (!supportsSignIn(p) || addStage === "running") return addProvider = p.providerId addPrivate = !!account && account.primary !== true addStatus = Commons.I18n.t("agents.starting") diff --git a/test/shell.d/agents-panel-test.sh b/test/shell.d/agents-panel-test.sh index d82cc494e79..5daa8487f66 100755 --- a/test/shell.d/agents-panel-test.sh +++ b/test/shell.d/agents-panel-test.sh @@ -8,6 +8,20 @@ run_node_test <<'JS' const fs = require('fs') const panelSource = fs.readFileSync(root + '/shell/plugins/agents/Panel.qml', 'utf8') +const signInFunctions = ['supportsSignIn', 'needsSignIn', 'otherTrouble'].map(name => { + const match = panelSource.match(new RegExp('function ' + name + '\\([^]*?\\n \\}')) + return match ? match[0] : '' +}).join('\n') +const signIn = new Function('addProviders', signInFunctions + '\nreturn {needsSignIn, otherTrouble}')( + [{providerId: 'claude'}, {providerId: 'codex'}, {providerId: 'grok'}] +) +for (const id of ['claude', 'codex', 'grok']) + assert(signIn.needsSignIn({providerId: id, usageStatusText: 'Sign-in expired'}), id + ' keeps its sign-in link') +assert(signIn.needsSignIn({usageStatusText: 'Sign-in expired'}), 'account rows keep their sign-in link') +for (const status of ['Sign-in expired', 'Waiting for auth']) { + assert(!signIn.needsSignIn({providerId: 'agy', usageStatusText: status}), 'an agent the panel cannot sign in gets no link for ' + status) + assert(signIn.otherTrouble({providerId: 'agy', usageStatusText: status}) === status, 'an agent the panel cannot sign in still shows ' + status) +} assert(/function launchAgent\(\)/.test(panelSource), 'agents panel launches the default agent') assert(/root\.bar\.run\("omarchy-agent --pick"\)/.test(panelSource), 'agents panel uses the desktop agent launcher') assert(/if \(buttonCode === Qt\.RightButton\) root\.launchAgent\(\)/.test(panelSource), 'agents right click launches the agent') From 2923e83161de99a838a27ddfc0ee2c7beddd31c0 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 22:58:16 -0400 Subject: [PATCH 045/546] Bring the Kimi collector in line with the usage collector contract Waiting for auth and Sign-in expired replace the catch-all Kimi unavailable, an unused machine stays hidden, a failed probe keeps the last limits without flagging them and zeroes elapsed windows, a reply cut off while reading advises a retry, probes are reused for 15 seconds, and the cache moves under cache_root() keyed by fetchedAtMs. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-kimi | 193 +++++++++++++----- .../agent-usage-kimi-regression-test.sh | 107 ++++++++-- 2 files changed, 234 insertions(+), 66 deletions(-) diff --git a/bin/omarchy-agent-usage-kimi b/bin/omarchy-agent-usage-kimi index f7d251743d1..f02bebdd138 100755 --- a/bin/omarchy-agent-usage-kimi +++ b/bin/omarchy-agent-usage-kimi @@ -9,6 +9,7 @@ from __future__ import annotations import argparse import fcntl import hashlib +import http.client import json import math import os @@ -39,7 +40,9 @@ OAUTH_CLIENT_ID = "17e5f671-d194-4dfb-9706-5516cb48c098" ALLOWED_DOMAINS = ("kimi.com", "kimi.ai") REFRESH_SKEW_S = 60 LOCK_WAIT_SECONDS = 15 -PROBE_REUSE_SECONDS = 60 +# A panel opened and shut repeatedly reuses a probe this recent instead of +# asking again; --force skips it. +PROBE_MIN_INTERVAL_SECONDS = 15 def credentials_dir() -> Path: @@ -57,9 +60,15 @@ CREDENTIALS_PATH = CREDENTIALS_DIR / "kimi-code.json" class KimiError(Exception): - def __init__(self, message: str, retry: bool = False): + # `retry` marks a service that couldn't be reached at all, so the panel + # tries again soon; `refused` marks a credential Kimi turned down, which + # only signing in again (or a new key) repairs; `missing` marks no + # credential at all. + def __init__(self, message: str, retry: bool = False, refused: bool = False, missing: bool = False): super().__init__(message) self.retry = retry + self.refused = refused + self.missing = missing class NoRedirectHandler(urllib.request.HTTPRedirectHandler): @@ -181,14 +190,16 @@ def request_refresh(refresh_token: str) -> dict[str, Any]: decoded = json.load(response) except urllib.error.HTTPError as error: if error.code in (400, 401, 403): - raise KimiError("Kimi rejected the sign-in; sign in again with `kimi login`") from error + raise KimiError("Kimi rejected the sign-in; sign in again with `kimi login`", refused=True) from error raise KimiError(f"Kimi token refresh failed (HTTP {error.code})") from error except urllib.error.URLError as error: raise KimiError("Could not reach the Kimi login service", retry=True) from error - except TimeoutError as error: - raise KimiError("Kimi token refresh timed out", retry=True) from error except json.JSONDecodeError as error: raise KimiError("Kimi returned an invalid token response") from error + except (OSError, http.client.HTTPException) as error: + # urllib only wraps failures while sending; a reset, timeout, or cut-off + # reply while reading the answer arrives raw, and is just as unreachable. + raise KimiError("Could not reach the Kimi login service", retry=True) from error if not isinstance(decoded, dict) or any( not isinstance(decoded.get(key), str) or not decoded[key] for key in ("access_token", "refresh_token") ): @@ -269,7 +280,7 @@ def refresh_credentials(auth: dict[str, str] | None = None) -> dict[str, Any]: return credentials refresh_token = str(credentials.get("refresh_token") or "") if not refresh_token: - raise KimiError("Kimi sign-in is missing; sign in with `kimi login`") + raise KimiError("Kimi sign-in expired; sign in again with `kimi login`", refused=True) granted = request_refresh(refresh_token) check() latest = read_credentials() @@ -305,7 +316,7 @@ def api_key(auth: dict[str, str] | None = None) -> str: if auth is not None: auth["key"] = str(credentials.get("access_token") or "") if not credentials.get("access_token"): - raise KimiError(AUTH_HELP) + raise KimiError(AUTH_HELP, missing=True) if not token_fresh(credentials): credentials = refresh_credentials(auth) return str(credentials.get("access_token") or "") @@ -326,16 +337,17 @@ def quota(key: str) -> dict[str, Any]: except urllib.error.HTTPError as error: if error.code in (401, 403): remedy = "Update or unset KIMI_API_KEY." if os.environ.get("KIMI_API_KEY", "").strip() else "Sign in again with `kimi login`." - raise KimiError("Kimi rejected the credentials. " + remedy) from error + raise KimiError("Kimi rejected the credentials. " + remedy, refused=True) from error if error.code == 404: raise KimiError("Kimi usage endpoint not available. Try Kimi for Coding.") from error raise KimiError(f"Kimi API returned HTTP {error.code}") from error except urllib.error.URLError as error: raise KimiError("Could not reach the Kimi API", retry=True) from error - except TimeoutError as error: - raise KimiError("Kimi usage request timed out", retry=True) from error except json.JSONDecodeError as error: raise KimiError("Kimi returned an invalid usage response") from error + except (OSError, http.client.HTTPException) as error: + # A reset, timeout, or cut-off reply while reading arrives unwrapped. + raise KimiError("Could not reach the Kimi API", retry=True) from error if not isinstance(decoded, dict): raise KimiError("Kimi returned an empty usage response") return decoded @@ -448,74 +460,151 @@ def to_limits(data: dict[str, Any]) -> list[dict[str, Any]]: return limits +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 probe_cache(key: str) -> Path: identity = str(CREDENTIALS_PATH) + "\0" + USAGE_URL + "\0" + key digest = hashlib.sha256(identity.encode()).hexdigest() - root = Path(os.environ.get("XDG_CACHE_HOME") or Path.home() / ".cache") / "omarchy" / "agent-usage" - return root / f"kimi-limits-{digest}.json" + return cache_root() / f"kimi-limits-{digest}.json" + + +def fetched_at_ms(data: dict[str, Any]) -> int: + # Older caches stored the same epoch milliseconds under `fetchedAt`; the + # next good probe rewrites them under `fetchedAtMs`. + raw = data.get("fetchedAtMs", data.get("fetchedAt")) + if isinstance(raw, bool): + return 0 + try: + stamp = float(raw) + except (TypeError, ValueError): + return 0 + return int(stamp) if math.isfinite(stamp) and stamp > 0 else 0 def read_probe(path: Path) -> dict[str, Any]: try: data = json.loads(path.read_text()) - if isinstance(data, dict) and isinstance(data.get("limits"), list): - stamp = float(data.get("fetchedAt")) - if math.isfinite(stamp) and stamp > 0: - return data - except (OSError, ValueError, TypeError): + except (OSError, ValueError): return {} + if isinstance(data, dict) and isinstance(data.get("limits"), list) and fetched_at_ms(data): + return data return {} +# A kept window whose reset has passed is an untouched allowance, not an +# unknown one: it comes back at zero, with no reset time until Kimi names +# the next one. +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 ValueError: + 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 if limit_window_open(entry, now) else dict(entry, percent=0.0, resetsAt="") + for entry in entries if isinstance(entry, dict) + ] + + +def has_credential() -> bool: + return bool(os.environ.get("KIMI_API_KEY", "").strip() or read_credentials().get("access_token")) + + +# Kimi has no local stats, so it counts as in use on this machine when it has +# a credential or a Kimi home of its own: a home without a sign-in is waiting +# for one, while a machine with neither never ran Kimi at all. +def in_use() -> bool: + if has_credential() or CREDENTIALS_PATH.exists(): + return True + if home := os.environ.get("KIMI_CODE_HOME"): + return Path(home).expanduser().is_dir() + return (Path.home() / ".kimi-code").is_dir() or (Path.home() / ".kimi").is_dir() + + +def failure_record(error: Exception, cached: dict[str, Any]) -> dict[str, Any]: + fallback = usable_cached_limits(cached) + help_text = str(error) if isinstance(error, KimiError) else "Kimi usage scan failed" + if isinstance(error, KimiError) and error.missing or not has_credential(): + status, help_text = "Waiting for auth", AUTH_HELP + elif isinstance(error, KimiError) and error.refused: + status = "Sign-in expired" + elif fallback: + # A failed check that kept the last good limits isn't trouble to flag; + # limitsStale and the age on hover say the rest. + status = "" + else: + status = "Kimi limits unavailable" + record = base_record( + ready=True, + usageStatusText=status, + authHelpText=help_text, + limits=fallback, + limitsStale=bool(fallback), + limitsFetchedAt=fetched_at_ms(cached) if fallback else 0, + ) + if isinstance(error, KimiError) and error.retry: + record["retryAdvised"] = True + return record + + def main() -> int: parser = argparse.ArgumentParser(description="Print the Kimi usage record as JSON") - # Flags accepted so the collector answers the same invocation as the others. + # Kimi keeps no local history, so there is no scan for --limits-only to + # reuse; --force skips the probe reuse window. parser.add_argument("--force", action="store_true") parser.add_argument("--limits-only", action="store_true") args = parser.parse_args() + if not in_use(): + print(json.dumps(base_record(), separators=(",", ":"))) + return 0 + cached = {} auth = {} try: key = api_key(auth) if not key: - record = base_record(usageStatusText="Kimi unavailable", authHelpText=AUTH_HELP) - else: - path = probe_cache(key) - cached = read_probe(path) - if not cached and auth.get("previous_key"): - cached = read_probe(probe_cache(auth["previous_key"])) - if cached: - try: - write_json(path, cached) - except OSError: - pass - age = time.time() * 1000 - float(cached.get("fetchedAt") or 0) - if not args.force and 0 <= age < PROBE_REUSE_SECONDS * 1000: - limits, fetched_at = cached["limits"], cached["fetchedAt"] - else: - limits, fetched_at = to_limits(quota(key)), int(time.time() * 1000) + raise KimiError(AUTH_HELP, missing=True) + path = probe_cache(key) + cached = read_probe(path) + if not cached and auth.get("previous_key"): + cached = read_probe(probe_cache(auth["previous_key"])) + if cached: try: - write_json(path, {"limits": limits, "fetchedAt": fetched_at}) + write_json(path, cached) except OSError: pass - record = base_record( - ready=True, - limits=limits, - limitsFetchedAt=fetched_at, - ) - except KimiError as error: - if not cached: - cached = read_probe(probe_cache(auth.get("key", ""))) - record = base_record(usageStatusText="Kimi unavailable", authHelpText=str(error), - limits=cached.get("limits", []), limitsStale=bool(cached.get("limits")), - limitsFetchedAt=cached.get("fetchedAt", 0), retryAdvised=error.retry) - except Exception: - if not cached: - cached = read_probe(probe_cache(auth.get("key", ""))) - record = base_record(usageStatusText="Kimi unavailable", authHelpText="Kimi usage scan failed", - limits=cached.get("limits", []), limitsStale=bool(cached.get("limits")), - limitsFetchedAt=cached.get("fetchedAt", 0)) + fallback = usable_cached_limits(cached) + age = time.time() - fetched_at_ms(cached) / 1000 + if fallback and not args.force and 0 <= age < PROBE_MIN_INTERVAL_SECONDS: + limits, fetched_at = fallback, fetched_at_ms(cached) + else: + limits, fetched_at = to_limits(quota(key)), int(time.time() * 1000) + try: + write_json(path, {"limits": limits, "fetchedAtMs": fetched_at}) + except OSError: + pass + record = base_record(ready=True, limits=limits, limitsFetchedAt=fetched_at) + except Exception as error: + if not cached and auth.get("key"): + cached = read_probe(probe_cache(auth["key"])) + record = failure_record(error, cached) print(json.dumps(record, separators=(",", ":"))) return 0 diff --git a/test/shell.d/agent-usage-kimi-regression-test.sh b/test/shell.d/agent-usage-kimi-regression-test.sh index 9b5f9f92f1f..a4c2d7c05fa 100644 --- a/test/shell.d/agent-usage-kimi-regression-test.sh +++ b/test/shell.d/agent-usage-kimi-regression-test.sh @@ -37,7 +37,11 @@ class Handler(BaseHTTPRequestHandler): def do_GET(self): requests.append((self.path, self.headers.get("Authorization"), None)) - if state.get("redirect"): + if state.get("drop"): + # Hang up without an answer, the way an unreachable network looks. + self.close_connection = True + self.connection.shutdown(2) + elif state.get("redirect"): self.send_response(302) self.send_header("Location", "/sink") self.end_headers() @@ -57,6 +61,17 @@ class Handler(BaseHTTPRequestHandler): else: self.reply(200, state.get("grant", {"access_token": "AT_new", "refresh_token": "RT_new", "expires_in": 3600})) +def run(*flags): + saved = sys.argv + sys.argv = [saved[1], *flags] + output = io.StringIO() + try: + with contextlib.redirect_stdout(output): + assert scanner.main() == 0 + finally: + sys.argv = saved + return json.loads(output.getvalue()) + server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) threading.Thread(target=server.serve_forever, daemon=True).start() local_url = f"http://127.0.0.1:{server.server_port}" @@ -90,15 +105,7 @@ def credentials(token="AT_old", expiry=None): "expires_at": time.time() - 10 if expiry is None else expiry})) def record(): - saved = sys.argv - sys.argv = [saved[1], "--force", "--limits-only"] - output = io.StringIO() - try: - with contextlib.redirect_stdout(output): - assert scanner.main() == 0 - finally: - sys.argv = saved - return json.loads(output.getvalue()) + return run("--force", "--limits-only") def destination(root): for url in ("https://outside.example/api/oauth/token", local_url + "/api/oauth/token"): @@ -209,6 +216,7 @@ def stale_limits(root): stale = record() assert good["limits"] and stale["limits"] == good["limits"], "failed probe hid the provider" assert stale["limitsStale"] and stale["limitsFetchedAt"] == good["limitsFetchedAt"] + assert stale["usageStatusText"] == "" and "retryAdvised" not in stale, "kept limits were flagged as trouble" os.environ["KIMI_API_KEY"] = "AT_other_account" assert record()["limits"] == [], "cached quota crossed credential identities" @@ -245,7 +253,9 @@ def concurrent_signin(root): def concurrent_logout(root): credentials() state["token_hook"] = scanner.CREDENTIALS_PATH.unlink - assert not record()["ready"] and not scanner.CREDENTIALS_PATH.exists(), "refresh restored a signed-out account" + data = record() + assert not scanner.CREDENTIALS_PATH.exists(), "refresh restored a signed-out account" + assert data["usageStatusText"] == "Waiting for auth", "a signed-out home did not ask for a sign-in" def invalid_grant(root): for grant in ({"access_token": "AT_new", "expires_in": 3600}, @@ -254,7 +264,8 @@ def invalid_grant(root): credentials() before = scanner.CREDENTIALS_PATH.read_text() state["grant"] = grant - assert not record()["ready"], "invalid rotated token was accepted" + data = record() + assert data["usageStatusText"] == "Kimi limits unavailable" and not data["limits"], "invalid rotated token was accepted" assert scanner.CREDENTIALS_PATH.read_text() == before, "invalid grant replaced usable credentials" def key_help(root): @@ -270,16 +281,84 @@ def error_redaction(root): def redirects(root): os.environ["KIMI_API_KEY"] = "AT_env" state["redirect"] = True - assert not record()["ready"] + data = record() + assert data["usageStatusText"] == "Kimi limits unavailable" and not data["limits"] assert len(requests) == 1 and requests[0][0] == "/coding/v1/usages", "usage followed a redirect" +def unused(root): + data = run() + assert data["ready"] is False and data["usageStatusText"] == "" and data["authHelpText"] == "", "an unused machine was flagged" + assert data["limits"] == [] and data["recentDays"] == [] and not requests, "an unused machine probed" + +def waiting_for_auth(root): + (root / "code home").mkdir() + data = run() + assert data["ready"] and data["usageStatusText"] == "Waiting for auth", "a home without a sign-in did not ask for one" + assert "kimi login" in data["authHelpText"] and not requests + +def refused(root): + credentials("AT_live", time.time() + 3600) + good = record() + state["usage_status"] = 401 + data = record() + assert data["usageStatusText"] == "Sign-in expired" and "kimi login" in data["authHelpText"], "a refused token read as healthy" + assert data["limits"] == good["limits"] and data["limitsStale"], "a refused token hid the last good limits" + +def refused_refresh(root): + credentials() + state["token_status"] = 401 + assert record()["usageStatusText"] == "Sign-in expired", "a lapsed refresh token did not ask for a sign-in" + +def elapsed_window(root): + os.environ["KIMI_API_KEY"] = "AT_env" + path = scanner.probe_cache("AT_env") + path.parent.mkdir(parents=True, exist_ok=True) + # The pre-fetchedAtMs cache shape, read once and carried forward. + path.write_text(json.dumps({"fetchedAt": 1_700_000_000_000, "limits": [ + {"label": "Session (5-hour)", "percent": 0.9, "resetsAt": "2020-01-01T00:00:00Z"}, + {"label": "Weekly limit", "percent": 0.4, "resetsAt": "2999-01-01T00:00:00Z"}, + {"label": "Monthly total", "percent": 0.2, "resetsAt": ""}]})) + state["usage_status"] = 503 + data = record() + assert data["limits"] == [ + {"label": "Session (5-hour)", "percent": 0.0, "resetsAt": ""}, + {"label": "Weekly limit", "percent": 0.4, "resetsAt": "2999-01-01T00:00:00Z"}, + {"label": "Monthly total", "percent": 0.2, "resetsAt": ""}], "an elapsed window was not zeroed" + assert data["limitsStale"] and data["limitsFetchedAt"] == 1_700_000_000_000 + +def probe_reuse(root): + os.environ["KIMI_API_KEY"] = "AT_env" + first = run() + second = run("--limits-only") + assert len(requests) == 1 and second["limits"] == first["limits"], "a recent probe was not reused" + assert second["limitsFetchedAt"] == first["limitsFetchedAt"] and not second["limitsStale"] + cache = json.loads(scanner.probe_cache("AT_env").read_text()) + assert cache["fetchedAtMs"] == first["limitsFetchedAt"] and "fetchedAt" not in cache + run("--force") + assert len(requests) == 2, "--force reused a recent probe" + assert scanner.PROBE_MIN_INTERVAL_SECONDS == 15 + +def unreachable(root): + os.environ["KIMI_API_KEY"] = "AT_env" + good = record() + state["drop"] = True + data = record() + assert data.get("retryAdvised") is True, "an unreachable service did not ask for a retry" + assert data["limits"] == good["limits"] and data["limitsStale"] and data["usageStatusText"] == "" + state["drop"] = False + state["usage_status"] = 500 + assert "retryAdvised" not in record(), "a server error asked for a retry" + cases = {"destination": destination, "home": home_isolation, "static": static_token, "cli_lock": cli_lock, "lock_failure": lock_failure, "heartbeat": heartbeat, "private_write": private_write, "quota": modern_quota, "stale": stale_limits, "rotation_cache": rotated_cache, "cache_identity": cache_identity, "concurrent_signin": concurrent_signin, "concurrent_logout": concurrent_logout, "invalid_grant": invalid_grant, - "key_help": key_help, "redaction": error_redaction, "redirects": redirects} + "key_help": key_help, "redaction": error_redaction, "redirects": redirects, + "unused": unused, "waiting_for_auth": waiting_for_auth, "refused": refused, + "refused_refresh": refused_refresh, "elapsed_window": elapsed_window, + "probe_reuse": probe_reuse, "unreachable": unreachable} failed = False for name, check in cases.items(): selected = os.environ.get("OMARCHY_TEST_KIMI_CASE") From 24452d7c43764b7ebc7724cf50ca0b7929e6eb84 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 22:58:16 -0400 Subject: [PATCH 046/546] Add Kimi's mark to the agents panel The K and dot from Moonshot's own kimi.com icon set, in white with a dark twin for light surfaces. Co-Authored-By: Claude Opus 5.5 --- shell/plugins/agents/assets/kimi-light.svg | 1 + shell/plugins/agents/assets/kimi.svg | 1 + 2 files changed, 2 insertions(+) create mode 100644 shell/plugins/agents/assets/kimi-light.svg create mode 100644 shell/plugins/agents/assets/kimi.svg diff --git a/shell/plugins/agents/assets/kimi-light.svg b/shell/plugins/agents/assets/kimi-light.svg new file mode 100644 index 00000000000..c7bfa27f3a9 --- /dev/null +++ b/shell/plugins/agents/assets/kimi-light.svg @@ -0,0 +1 @@ +Kimi diff --git a/shell/plugins/agents/assets/kimi.svg b/shell/plugins/agents/assets/kimi.svg new file mode 100644 index 00000000000..4222df4ff97 --- /dev/null +++ b/shell/plugins/agents/assets/kimi.svg @@ -0,0 +1 @@ +Kimi From ca1ce3ddc46fa6a77df114a8085c5df9fe3b786e Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:03:52 -0400 Subject: [PATCH 047/546] Bring the Kimchi collector in line with the usage collector contract A rejected key reads Sign-in expired and missing credits Balance unavailable, local sessions without a key wait for auth, retries are advised only when the gateway can't be reached, answers are reused for 15 seconds unless forced, elapsed kept budgets come back at zero, and the budget cache joins the shared kimchi-limits name. Sessions are listed with rg --no-ignore, forks are deduplicated by id and timestamp across files, untrusted stamps and token fields skip the field, and the scan is cached per day under a lock, only when every file was read. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-kimchi | 333 +++++++++++++++++------- test/shell.d/agent-usage-kimchi-test.sh | 102 +++++++- 2 files changed, 327 insertions(+), 108 deletions(-) diff --git a/bin/omarchy-agent-usage-kimchi b/bin/omarchy-agent-usage-kimchi index 9a7da373392..3704a4da8c4 100755 --- a/bin/omarchy-agent-usage-kimchi +++ b/bin/omarchy-agent-usage-kimchi @@ -12,9 +12,13 @@ and token usage. The agents panel only ever reads the JSON this prints. """ import argparse +import fcntl import hashlib +import http.client import json +import math import os +import subprocess import tempfile import time import urllib.error @@ -25,6 +29,7 @@ from pathlib import Path AGENT_ID = "kimchi" AGENT_NAME = "Kimchi" AUTH_HELP = "Start Kimchi and log in, or set KIMCHI_API_KEY." +REJECTED_HELP = "Kimchi rejected its API key. Log in again in Kimchi, or update KIMCHI_API_KEY." REGION_LLM_BASE = {"us": "https://llm.kimchi.dev", "eu": "https://llm.eu.kimchi.dev"} PLAN_LABELS = {"community": "Community", "coder": "Coder", "teams": "Teams", "enterprise": "Enterprise"} BUDGET_LABELS = { @@ -35,8 +40,14 @@ BUDGET_LABELS = { "ORGANIZATION_SOFT": "Org budget", "ORGANIZATION_HARD": "Org budget", } +SCAN_REUSE_SECONDS = 20 # A limits-only refresh needs fresh credits, not a fresh history scan. -LIMITS_ONLY_SCAN_REUSE_SECONDS = 900 +LIMITS_ONLY_REUSE_SECONDS = 900 +# A panel opened and shut repeatedly reuses answers this recent instead of +# asking the gateway again; --force skips it. +PROBE_MIN_INTERVAL_SECONDS = 15 +# A token count past this is a corrupt field, not a reply. +MAX_TOKENS = 1_000_000_000_000 def read_json(path): @@ -62,14 +73,30 @@ def write_json(path, payload): def cache_root(): - return Path(os.environ.get("XDG_CACHE_HOME") or (Path.home() / ".cache")) / "omarchy" / "agent-usage" + root = Path(os.environ.get("XDG_CACHE_HOME") or (Path.home() / ".cache")) / "omarchy" / "agent-usage" + root.mkdir(parents=True, exist_ok=True) + return root +# A finite, positive number, or 0: cache stamps are read as found. def number(value): + if isinstance(value, bool): + return 0 try: - return max(0, int(value or 0)) - except Exception: + parsed = float(value or 0) + except (TypeError, ValueError): + return 0 + return int(parsed) if math.isfinite(parsed) and parsed > 0 else 0 + + +# Token counts from session files are untrusted: anything but a plain, +# plausible count skips the field. +def tokens(value): + if isinstance(value, bool) or not isinstance(value, (int, float)): return 0 + if not math.isfinite(value) or value <= 0 or value >= MAX_TOKENS: + return 0 + return int(value) # Kimchi sends dollar amounts as strings, sometimes with a sign or unit. @@ -77,34 +104,84 @@ def usd(value): if isinstance(value, bool): return None if isinstance(value, (int, float)): - return float(value) - text = str(value or "").lower().replace("$", "").replace("usd", "").replace(",", ".").strip() + parsed = float(value) + else: + text = str(value or "").lower().replace("$", "").replace("usd", "").replace(",", ".").strip() + try: + parsed = float(text) if text else None + except ValueError: + return None + return parsed if parsed is not None and math.isfinite(parsed) else None + + +def parse_time(value): + if not isinstance(value, str) or not value.strip(): + return None try: - return float(text) if text else None + parsed = datetime.fromisoformat(value.strip().replace("Z", "+00:00")) except ValueError: return None + return parsed if parsed.tzinfo else parsed.replace(tzinfo=timezone.utc) # pi stamps entries with ISO strings and messages with epoch milliseconds. +# A stamp that isn't a plausible moment (zero, negative, NaN, far future, +# not a time at all) leaves the reply undated rather than counted as today. def local_day(value): - if isinstance(value, (int, float)) and not isinstance(value, bool): - return datetime.fromtimestamp(value / 1000 if value > 10_000_000_000 else value).astimezone().date().isoformat() + if isinstance(value, bool): + return None + if isinstance(value, (int, float)): + if not math.isfinite(value) or value <= 0: + return None + seconds = value / 1000 if value > 10_000_000_000 else value + else: + parsed = parse_time(value) + if parsed is None: + return None + seconds = parsed.timestamp() + # Nothing wrote a Kimchi session before 2000, or from the future. + if not 946_684_800 <= seconds <= time.time() + 2 * 86_400: + return None try: - parsed = datetime.fromisoformat(str(value or "").replace("Z", "+00:00")) - except ValueError: + return datetime.fromtimestamp(seconds).astimezone().date().isoformat() + except (OverflowError, OSError, ValueError): return None - return (parsed if parsed.tzinfo else parsed.replace(tzinfo=timezone.utc)).astimezone().date().isoformat() # The harness directory Kimchi keeps its sessions under: KIMCHI_CODING_AGENT_DIR -# when set, as the CLI itself reads it, else ~/.config/kimchi/harness. +# when set, as the CLI itself reads it, else ~/.config/kimchi/harness. Kimchi +# has no profiles that move it. def harness_dir(): return Path(os.environ.get("KIMCHI_CODING_AGENT_DIR") or (Path.home() / ".config" / "kimchi" / "harness")) -# Every session file under /sessions//. Only assistant -# messages carry usage, and a resumed session can repeat an entry, so each -# entry id counts once per file. +# Listed with rg --no-ignore, since a $HOME kept as a git checkout would +# otherwise hide ~/.config from the listing. The flag says whether every +# file could be listed. +def session_files(sessions_root): + root = Path(sessions_root) + if not root.is_dir(): + return [], True + try: + listed = subprocess.run( + ["rg", "--files", "--no-ignore", "--hidden", "--glob", "*.jsonl", str(root)], + capture_output=True, text=True, timeout=60, + ) + except (OSError, subprocess.TimeoutExpired): + return [], False + # rg exits 1 when it finds nothing, and 2 when part of the tree was + # unreadable while still listing the rest. + return sorted(Path(line) for line in listed.stdout.splitlines() if line), listed.returncode in (0, 1) + + +def empty_bucket(): + return {"inputTokens": 0, "outputTokens": 0, "cacheReadInputTokens": 0, "cacheCreationInputTokens": 0} + + +# Every session file under /sessions. Only assistant messages carry +# usage. A forked or resumed session repeats its parent's entries, so a reply +# with both an id and a timestamp counts once across every file; one without +# an id always counts. Returns the stats and whether every source was read. def local_stats(sessions_root): now = datetime.now().astimezone() today = now.date().isoformat() @@ -114,12 +191,13 @@ def local_stats(sessions_root): by_day = {day: 0 for day in week} today_by_model = {} models = {} - files = sorted(Path(sessions_root).glob("*/*.jsonl")) + seen = set() + files, complete = session_files(sessions_root) for path in files: - seen = set() try: lines = path.read_text(encoding="utf-8", errors="replace").splitlines() except OSError: + complete = False continue for line in lines: if '"usage"' not in line: @@ -131,39 +209,45 @@ def local_stats(sessions_root): message = entry.get("message") if isinstance(entry, dict) and entry.get("type") == "message" else None if not isinstance(message, dict) or message.get("role") != "assistant": continue - entry_id = str(entry.get("id") or "") - if entry_id and entry_id in seen: - continue - seen.add(entry_id) - - usage = message.get("usage") or {} - uncached = number(usage.get("input")) - output = number(usage.get("output")) - cached = number(usage.get("cacheRead")) - created = number(usage.get("cacheWrite")) + stamp = entry.get("timestamp") or message.get("timestamp") + entry_id = entry.get("id") + if entry_id and stamp: + identity = (str(entry_id), str(stamp)) + if identity in seen: + continue + seen.add(identity) + + usage = message.get("usage") if isinstance(message.get("usage"), dict) else {} + uncached = tokens(usage.get("input")) + output = tokens(usage.get("output")) + cached = tokens(usage.get("cacheRead")) + created = tokens(usage.get("cacheWrite")) if not (uncached or output or cached or created): - uncached = number(usage.get("totalTokens")) - tokens = uncached + output + cached + created - if not tokens: + uncached = tokens(usage.get("totalTokens")) + total = uncached + output + cached + created + if not total: continue - day = local_day(entry.get("timestamp") or message.get("timestamp")) or today model = str(message.get("model") or AGENT_ID) prompts += 1 - sessions.add(path.name) - days.add(day) - totals = models.setdefault(model, {"inputTokens": 0, "outputTokens": 0, "cacheReadInputTokens": 0, "cacheCreationInputTokens": 0}) + sessions.add(str(path)) + totals = models.setdefault(model, empty_bucket()) totals["inputTokens"] += uncached totals["outputTokens"] += output totals["cacheReadInputTokens"] += cached totals["cacheCreationInputTokens"] += created + # An undated reply still counts in the totals, just on no day. + day = local_day(stamp) + if day is None: + continue + days.add(day) if day in by_day: - by_day[day] += tokens + by_day[day] += total if day == today: today_prompts += 1 - today_sessions.add(path.name) - today_tokens += tokens - today_by_model[model] = today_by_model.get(model, 0) + tokens + today_sessions.add(str(path)) + today_tokens += total + today_by_model[model] = today_by_model.get(model, 0) + total return { "hasLocalStats": bool(sessions), @@ -177,26 +261,45 @@ def local_stats(sessions_root): "activeDays": len(days), "activeDates": sorted(days), "modelUsage": models, - } + }, complete -# A scan from before midnight would count yesterday's sessions as today's. +# A scan is reused on the day it was made only, since one from before +# midnight would count yesterday's sessions as today's. It's rebuilt under a +# lock so overlapping refreshes scan once, and a scan that couldn't read +# every source isn't kept. def cached_local_stats(sessions_root, max_age): - cache = cache_root() / "kimchi-stats.json" - cached = read_json(cache) or {} + digest = hashlib.sha256(str(sessions_root).encode()).hexdigest()[:16] + cache = cache_root() / f"kimchi-scan-{digest}.json" today = datetime.now().astimezone().date().isoformat() - if cached.get("root") == str(sessions_root) and cached.get("day") == today and time.time() - float(cached.get("at") or 0) < max_age: - return cached["stats"] - stats = local_stats(sessions_root) - write_json(cache, {"root": str(sessions_root), "day": today, "at": time.time(), "stats": stats}) - return stats + + def fresh(): + cached = read_json(cache) + if not isinstance(cached, dict) or cached.get("scanDate") != today or not isinstance(cached.get("stats"), dict): + return None + age = time.time() - number(cached.get("scannedAt")) + return cached["stats"] if 0 <= age < max_age else None + + stats = fresh() + if stats is not None: + return stats + with (cache_root() / f"kimchi-scan-{digest}.lock").open("w") as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + stats = fresh() + if stats is not None: + return stats + stats, complete = local_stats(sessions_root) + if complete: + write_json(cache, {"scanDate": today, "scannedAt": time.time(), "stats": stats}) + return stats # The key and gateway Kimchi itself would use: KIMCHI_API_KEY, KIMCHI_BASE_URL, # and KIMCHI_REGION over ~/.config/kimchi/config.json, the region's gateway # when nothing names one. def gateway(): - config = read_json(Path.home() / ".config" / "kimchi" / "config.json") or {} + config = read_json(Path.home() / ".config" / "kimchi" / "config.json") + config = config if isinstance(config, dict) else {} key = os.environ.get("KIMCHI_API_KEY") or config.get("apiKey") or config.get("api_key") or "" base = os.environ.get("KIMCHI_BASE_URL") or config.get("llmEndpoint") or "" if not base: @@ -240,19 +343,13 @@ def credits_part(data): if remaining is not None and remaining >= 0: part["balance"] = {"remaining": remaining, "funded": 0, "spent": 0, "currency": "USD", "estimated": False} has_credits = next((data[k] for k in ("has_credits", "hasCredits") if k in data), None) + # The gateway throttles an account with no credits left, and with no funded + # figure the balance gauge has no way to raise that alarm itself. if has_credits is False or (remaining is not None and remaining <= 0): part["usageStatusText"] = "Out of credits, rate limited" return part -def parse_time(value): - try: - parsed = datetime.fromisoformat(str(value or "").replace("Z", "+00:00")) - except ValueError: - return None - return parsed if parsed.tzinfo else parsed.replace(tzinfo=timezone.utc) - - # The panel titles a window by its clock, so a budget period that spans a # month or a week says so rather than leaving the panel to guess from a label. def period_title(period): @@ -274,7 +371,8 @@ def period_title(period): # may be spent, not what was bought, so it never funds the balance gauge. def budget_part(data): period = data.get("period") if isinstance(data.get("period"), dict) else {} - resets = str(period.get("endTime") or "") + # The gateway's own period end, or none: a reset time is never made up. + resets = str(period.get("endTime") or "") if parse_time(period.get("endTime")) else "" title = period_title(period) limits = [] for entry in data.get("budgets") or []: @@ -282,13 +380,11 @@ def budget_part(data): continue cap = usd(entry.get("budgetLimitUsd")) spent = usd(entry.get("totalSpendUsd")) - if not cap or spent is None: + if not cap or cap <= 0 or spent is None: continue - limit = {"label": BUDGET_LABELS.get(entry.get("scope"), "Budget"), "percent": max(0.0, min(1.0, spent / cap))} + limit = {"label": BUDGET_LABELS.get(entry.get("scope"), "Budget"), "percent": max(0.0, min(1.0, spent / cap)), "resetsAt": resets} if title: limit["title"] = title - if resets: - limit["resetsAt"] = resets limits.append(limit) for provider in entry.get("providerBudgets") or []: if not isinstance(provider, dict) or not str(provider.get("limitType") or "").endswith("CAPPED"): @@ -296,39 +392,72 @@ def budget_part(data): provider_cap = usd(provider.get("budgetLimitUsd")) used = usd(provider.get("usageUsd")) name = str(provider.get("provider") or "").strip() - if not (provider_cap and used is not None and name): + if not (provider_cap and provider_cap > 0 and used is not None and name): continue name = name[:1].upper() + name[1:] - scoped = {"label": f"{name} budget", "title": f"{name} {title}" if title else f"{name} budget", "percent": max(0.0, min(1.0, used / provider_cap))} - if resets: - scoped["resetsAt"] = resets - limits.append(scoped) + limits.append({ + "label": f"{name} budget", + "title": f"{name} {title}" if title else f"{name} budget", + "percent": max(0.0, min(1.0, used / provider_cap)), + "resetsAt": resets, + }) return {"limits": limits} # Keyed by the API key and the gateway it was asked of, so a machine signed # in to someone else, or pointed at another region, never shows the last # one's numbers. One file per endpoint, so a run only ever writes the answer -# it actually got. +# it actually got: the budget's limits, and the credits' plan and balance. +CACHE_NAMES = {"budget": "limits", "credits": "credits"} + + def billing_cache(key, base, resource): digest = hashlib.sha256((base + "\n" + key).encode()).hexdigest()[:16] - return cache_root() / f"kimchi-{resource}-{digest}.json" + return cache_root() / f"kimchi-{CACHE_NAMES[resource]}-{digest}.json" + + +def read_part(path, resource): + cached = read_json(path) + if not isinstance(cached, dict) and resource == "budget": + # Budgets were cached as kimchi-budget-.json before they joined the + # limits cache name every collector shares. + cached = read_json(path.with_name(path.name.replace("kimchi-limits-", "kimchi-budget-", 1))) + return cached if isinstance(cached, dict) and number(cached.get("fetchedAtMs")) else None + + +# A kept budget whose period ended is an untouched allowance, not an unknown +# one: it comes back at zero, with no reset time until the gateway names the +# next period. +def limit_window_open(entry, now): + resets_at = parse_time(entry.get("resetsAt")) + return resets_at is None or resets_at > now + + +def usable_cached_limits(part): + entries = part.get("limits") + if not isinstance(entries, list): + return [] + now = datetime.now(timezone.utc) + return [ + entry if limit_window_open(entry, now) else dict(entry, percent=0.0, resetsAt="") + for entry in entries if isinstance(entry, dict) + ] # Refreshes can overlap (the panel's own and a limits retry), so an answer # only replaces one fetched after it. def remember(path, part): - current = read_json(path) or {} - if float(current.get("fetchedAtMs") or 0) <= part["fetchedAtMs"]: + current = read_json(path) + if not isinstance(current, dict) or number(current.get("fetchedAtMs")) <= part["fetchedAtMs"]: write_json(path, part) # The two endpoints answer on their own, and so does each fall back: a part -# that didn't come back is the last one that did, dimmed and dated, so one -# endpoint failing never blanks or overwrites what the other said. Any part -# that didn't come back fresh marks the record stale, cached or not. -def collect_billing(): - result = {"tierLabel": "", "limits": [], "usageStatusText": "", "authHelpText": AUTH_HELP, "signedIn": False, "limitsStale": False, "limitsFetchedAt": 0} +# that didn't come back is the last one that did, so one endpoint failing +# never blanks or overwrites what the other said. Any part that didn't come +# back fresh marks the record stale, cached or not. +def collect_billing(force): + result = {"tierLabel": "", "limits": [], "usageStatusText": "", "authHelpText": "", "signedIn": False, "limitsStale": False, "limitsFetchedAt": 0} key, base = gateway() if not key: return result @@ -336,9 +465,16 @@ def collect_billing(): parts = {} failed = set() - rejected = False + rejected = unreachable = False for resource, parse in (("credits", credits_part), ("budget", budget_part)): cache = billing_cache(key, base, resource) + cached = read_part(cache, resource) + if cached is not None: + cached = dict(cached, limits=usable_cached_limits(cached)) if resource == "budget" else cached + age = time.time() - number(cached.get("fetchedAtMs")) / 1000 + if not force and 0 <= age < PROBE_MIN_INTERVAL_SECONDS: + parts[resource] = cached + continue try: data = get_json(f"{base}/v1/{resource}", key) parts[resource] = dict(parse(data if isinstance(data, dict) else {}), fetchedAtMs=round(time.time() * 1000)) @@ -346,11 +482,13 @@ def collect_billing(): continue except urllib.error.HTTPError as error: rejected = rejected or error.code in (401, 403) + except (urllib.error.URLError, OSError, http.client.HTTPException): + # No answer at all: no route, no DNS, a timeout, or a reply cut off. + unreachable = True except Exception: pass failed.add(resource) - cached = read_json(cache) - if isinstance(cached, dict) and cached.get("fetchedAtMs"): + if cached is not None: parts[resource] = cached credits = parts.get("credits") or {} @@ -360,23 +498,22 @@ def collect_billing(): result["limits"] = (parts.get("budget") or {}).get("limits") or [] result["limitsStale"] = bool(failed) if parts: - result["limitsFetchedAt"] = min(part["fetchedAtMs"] for part in parts.values()) + result["limitsFetchedAt"] = min(number(part.get("fetchedAtMs")) for part in parts.values()) - # A rejected key can't be fixed from the panel: Kimchi signs in through - # its own CLI, not omarchy-agent-account-add, so this status deliberately - # isn't one the panel offers its Sign-in link for. if rejected: - result["usageStatusText"] = "Kimchi key rejected" + result["usageStatusText"] = "Sign-in expired" + result["authHelpText"] = REJECTED_HELP elif credits.get("usageStatusText"): result["usageStatusText"] = credits["usageStatusText"] elif "credits" not in parts: - result["usageStatusText"] = "Kimchi credits unavailable" - if parts and not rejected: - result["authHelpText"] = "" - if failed and not rejected: - # Usually the seconds after login before the network is up, or a - # gateway hiccup; the panel tries once more soon instead of waiting out - # its refresh interval. + result["usageStatusText"] = "Balance unavailable" + result["authHelpText"] = ( + "Couldn't reach Kimchi's gateway. Local Kimchi stats are still shown." if unreachable + else "Kimchi's gateway didn't return the credits. Local Kimchi stats are still shown." + ) + if unreachable and not rejected: + # Usually the seconds after login before the network is up; the panel + # tries once more soon instead of waiting out its refresh interval. result["retryAdvised"] = True return result @@ -387,16 +524,24 @@ def main(): parser.add_argument("--limits-only", action="store_true") args = parser.parse_args() - billing = collect_billing() + scan_age = 0 if args.force else (LIMITS_ONLY_REUSE_SECONDS if args.limits_only else SCAN_REUSE_SECONDS) + stats = cached_local_stats(harness_dir() / "sessions", scan_age) + billing = collect_billing(args.force) + signed_in = billing.pop("signedIn") + # Sessions on disk with no key to ask the gateway with: Kimchi is in use + # here, its plan and credits just can't be read yet. + if not signed_in and stats.get("hasLocalStats"): + billing["usageStatusText"] = "Waiting for auth" + billing["authHelpText"] = AUTH_HELP record = { "schemaVersion": 1, "id": AGENT_ID, "name": AGENT_NAME, "updatedAt": datetime.now(timezone.utc).isoformat(), - "ready": billing.pop("signedIn"), + "ready": signed_in or bool(stats.get("hasLocalStats")), } record.update(billing) - record.update(cached_local_stats(harness_dir() / "sessions", LIMITS_ONLY_SCAN_REUSE_SECONDS if args.limits_only and not args.force else 0)) + record.update(stats) print(json.dumps(record, separators=(",", ":"))) diff --git a/test/shell.d/agent-usage-kimchi-test.sh b/test/shell.d/agent-usage-kimchi-test.sh index 1aeb5ddf61f..fc8eb40ace8 100755 --- a/test/shell.d/agent-usage-kimchi-test.sh +++ b/test/shell.d/agent-usage-kimchi-test.sh @@ -78,7 +78,7 @@ PY } record=$(collect) -[[ $(jq -c '{ready, tierLabel, limits, hasLocalStats, balance}' <<<"$record") == '{"ready":false,"tierLabel":"","limits":[],"hasLocalStats":false,"balance":null}' ]] || +[[ $(jq -c '{ready, tierLabel, limits, hasLocalStats, balance, usageStatusText, authHelpText, limitsStale, limitsFetchedAt, totalPrompts, days: (.recentDays | length), modelUsage}' <<<"$record") == '{"ready":false,"tierLabel":"","limits":[],"hasLocalStats":false,"balance":null,"usageStatusText":"","authHelpText":"","limitsStale":false,"limitsFetchedAt":0,"totalPrompts":0,"days":7,"modelUsage":{}}' ]] || fail "a machine nobody signed in to Kimchi on gives an empty record" "$record" [[ ! -s $test_tmp/requests ]] || fail "no key means the gateway is never asked" "$(cat "$test_tmp/requests")" pass "a machine nobody signed in to Kimchi on gives an empty record" @@ -119,15 +119,16 @@ record=$(KIMCHI_API_KEY=key-empty collect) pass "an empty balance says Kimchi is rate limited" record=$(KIMCHI_API_KEY=key-revoked collect) -[[ $(jq -c '{ready, usageStatusText, retry: .retryAdvised}' <<<"$record") == '{"ready":true,"usageStatusText":"Kimchi key rejected","retry":null}' ]] || - fail "a rejected key says so without the panel's sign-in link, which can't sign in to Kimchi" "$record" +[[ $(jq -c '{ready, usageStatusText, retry: .retryAdvised}' <<<"$record") == '{"ready":true,"usageStatusText":"Sign-in expired","retry":null}' ]] || + fail "a rejected key asks for a sign-in" "$record" +[[ $(jq -r '.authHelpText' <<<"$record") == *"Log in again in Kimchi"* ]] || fail "a rejected key says to log in again in Kimchi" "$record" record=$(KIMCHI_API_KEY=key-offline collect) -[[ $(jq -c '{usageStatusText, retry: .retryAdvised}' <<<"$record") == '{"usageStatusText":"Kimchi credits unavailable","retry":true}' ]] || +[[ $(jq -c '{usageStatusText, retry: .retryAdvised}' <<<"$record") == '{"usageStatusText":"Balance unavailable","retry":true}' ]] || fail "an unreachable gateway asks for a sooner retry" "$record" pass "a rejected key and an unreachable gateway are told apart" -# A gateway that stops answering leaves the last numbers up, dimmed and dated, -# rather than an empty section. +# A gateway that stops answering leaves the last numbers up, marked stale and +# dated, rather than an empty section. record=$(KIMCHI_API_KEY=key-budget collect) fetched=$(jq -r '.limitsFetchedAt' <<<"$record") (( fetched > 0 )) && [[ $(jq -r '.limitsStale' <<<"$record") == false ]] || fail "a live answer is dated and not stale" "$record" @@ -151,14 +152,15 @@ record=$(KIMCHI_API_KEY=key-budget GATEWAY_DOWN=credits collect) pass "one endpoint failing keeps its own last answer" # A part that fails before anything was cached is missing, not current: the -# record is stale and asks for a sooner retry, a server error included. -rm -f "$XDG_CACHE_HOME"/omarchy/agent-usage/kimchi-budget-*.json +# record is stale, but a server that answered with an error isn't retried +# sooner. +rm -f "$XDG_CACHE_HOME"/omarchy/agent-usage/kimchi-limits-*.json record=$(KIMCHI_API_KEY=key-sprint GATEWAY_ERROR=budget collect) -[[ $(jq -c '{stale: .limitsStale, limits, credits: (.balance.remaining == 12), retry: .retryAdvised, usageStatusText}' <<<"$record") == '{"stale":true,"limits":[],"credits":true,"retry":true,"usageStatusText":""}' ]] || - fail "a budget endpoint failing with nothing cached leaves the record stale and retrying" "$record" +[[ $(jq -c '{stale: .limitsStale, limits, credits: (.balance.remaining == 12), retry: .retryAdvised, usageStatusText}' <<<"$record") == '{"stale":true,"limits":[],"credits":true,"retry":null,"usageStatusText":""}' ]] || + fail "a budget endpoint failing with nothing cached leaves the record stale" "$record" rm -f "$XDG_CACHE_HOME"/omarchy/agent-usage/kimchi-credits-*.json record=$(KIMCHI_API_KEY=key-sprint GATEWAY_ERROR=credits collect) -[[ $(jq -c '{stale: .limitsStale, balance, usageStatusText, retry: .retryAdvised}' <<<"$record") == '{"stale":true,"balance":null,"usageStatusText":"Kimchi credits unavailable","retry":true}' ]] || +[[ $(jq -c '{stale: .limitsStale, balance, usageStatusText, retry: .retryAdvised}' <<<"$record") == '{"stale":true,"balance":null,"usageStatusText":"Balance unavailable","retry":null}' ]] || fail "a credits endpoint failing with nothing cached says credits are unavailable" "$record" pass "a part that fails before anything was cached is missing, not current" @@ -177,7 +179,7 @@ pass "overlapping refreshes keep the latest answer" # The same key asked of another gateway is another answer. record=$(KIMCHI_API_KEY=key-budget KIMCHI_BASE_URL="https://gateway.example" GATEWAY_DOWN=all collect) -[[ $(jq -c '{tierLabel, limits, usageStatusText}' <<<"$record") == '{"tierLabel":"","limits":[],"usageStatusText":"Kimchi credits unavailable"}' ]] || +[[ $(jq -c '{tierLabel, limits, usageStatusText}' <<<"$record") == '{"tierLabel":"","limits":[],"usageStatusText":"Balance unavailable"}' ]] || fail "another gateway never shows the last gateway's numbers" "$record" pass "another gateway never shows the last gateway's numbers" @@ -229,8 +231,80 @@ pass "a limits-only refresh reuses the session scan" # A scan from another day is never reused, so yesterday's sessions don't # count as today's after midnight. -jq '.day = "2000-01-01"' "$XDG_CACHE_HOME/omarchy/agent-usage/kimchi-stats.json" >"$test_tmp/stats.json" -mv "$test_tmp/stats.json" "$XDG_CACHE_HOME/omarchy/agent-usage/kimchi-stats.json" +scan_cache=$(ls -t "$XDG_CACHE_HOME"/omarchy/agent-usage/kimchi-scan-*.json | head -1) +jq '.scanDate = "2000-01-01"' "$scan_cache" >"$test_tmp/stats.json" +mv "$test_tmp/stats.json" "$scan_cache" record=$(KIMCHI_API_KEY=key-coder COLLECT_ARGS="--limits-only" collect) [[ $(jq -r '.totalSessions' <<<"$record") == 3 ]] || fail "a scan from another day is made again" "$record" pass "a scan from another day is made again" + +# A recent answer is reused rather than asking the gateway again, unless a +# person pressed refresh. (The overlap test above left a credits answer +# stamped in the future, which is never reused.) +rm -f "$XDG_CACHE_HOME"/omarchy/agent-usage/kimchi-credits-*.json +KIMCHI_API_KEY=key-budget collect >/dev/null +: >"$test_tmp/requests" +record=$(KIMCHI_API_KEY=key-budget COLLECT_ARGS="--limits-only" collect) +[[ ! -s $test_tmp/requests ]] || fail "a recent gateway answer is reused" "$(cat "$test_tmp/requests")" +[[ $(jq -c '{stale: .limitsStale, tierLabel, limits: (.limits | length)}' <<<"$record") == '{"stale":false,"tierLabel":"Teams","limits":2}' ]] || + fail "a reused answer is current, not stale" "$record" +KIMCHI_API_KEY=key-budget collect >/dev/null +(( $(wc -l <"$test_tmp/requests") == 2 )) || fail "--force asks the gateway again" "$(cat "$test_tmp/requests")" +pass "a recent gateway answer is reused, and --force asks again" + +# A kept budget whose period ended comes back at zero with no reset time, +# never dropped, and a budget cached under its old name is still read. +limits_cache=$(ls -t "$XDG_CACHE_HOME"/omarchy/agent-usage/kimchi-limits-*.json | head -1) +legacy_cache=${limits_cache/kimchi-limits-/kimchi-budget-} +jq -c '.limits[0].resetsAt = "2020-01-01T00:00:00Z" | .limits[0].percent = 0.9 | .limits[1].resetsAt = ""' "$limits_cache" >"$legacy_cache" +rm "$limits_cache" +record=$(KIMCHI_API_KEY=key-budget GATEWAY_DOWN=budget collect) +[[ $(jq -c '[.limits[] | {percent: (.percent * 100 | round), resetsAt}]' <<<"$record") == '[{"percent":0,"resetsAt":""},{"percent":25,"resetsAt":""}]' ]] || + fail "an elapsed kept budget is zeroed, and one with no reset time is kept as it is" "$record" +[[ $(jq -r '.limitsStale' <<<"$record") == true ]] || fail "an elapsed kept budget is still stale" "$record" +pass "an elapsed kept budget comes back at zero" + +# Sessions on disk without a key: Kimchi is in use here and waits for a sign-in. +rm "$HOME/.config/kimchi/config.json" +record=$(collect) +[[ $(jq -c '{ready, usageStatusText, authHelpText, hasLocalStats}' <<<"$record") == '{"ready":true,"usageStatusText":"Waiting for auth","authHelpText":"Start Kimchi and log in, or set KIMCHI_API_KEY.","hasLocalStats":true}' ]] || + fail "local sessions without a key wait for auth" "$record" +pass "local sessions without a key wait for auth" + +# A forked session repeats its parent's replies with the same id and time, and +# counts them once; a reply without an id always counts. Untrusted stamps and +# token fields skip the field, never the record. +fork="$HOME/.config/kimchi/harness/sessions/--fork--" +mkdir -p "$fork" +{ + reply a-1 "$now" minimax-m3 '{"input": 100, "output": 20, "cacheRead": 800, "cacheWrite": 5, "totalTokens": 925}' + jq -nc --arg at "$now" '{type: "message", timestamp: $at, message: {role: "assistant", model: "glm-5.3", usage: {input: 4}}}' + jq -nc --arg at "$now" '{type: "message", timestamp: $at, message: {role: "assistant", model: "glm-5.3", usage: {input: 4}}}' + reply z-1 "not a time" glm-5.3 '{"input": 6, "output": true, "cacheRead": "lots", "cacheWrite": 1e300}' + jq -nc '{type: "message", id: "z-2", timestamp: -5, message: {role: "assistant", model: "glm-5.3", usage: {input: 2}}}' + jq -nc '{type: "message", id: "z-3", timestamp: 99999999999999999, message: {role: "assistant", model: "glm-5.3", usage: {input: 1}}}' +} >"$fork/fork.jsonl" +record=$(KIMCHI_API_KEY=key-coder collect) +[[ $(jq -c '{totalPrompts, totalSessions, today: .todayTokensByModel["glm-5.3"], glm: .modelUsage["glm-5.3"].inputTokens, out: .modelUsage["glm-5.3"].outputTokens}' <<<"$record") == '{"totalPrompts":10,"totalSessions":4,"today":98,"glm":102,"out":5}' ]] || + fail "forks count once, unidentified replies always count, and bad fields are skipped" "$record" +pass "forks count once and untrusted fields are skipped" + +# A $HOME kept as a git checkout that ignores everything still has its sessions +# found. +mkdir -p "$HOME/.git" +printf '*\n' >"$HOME/.gitignore" +record=$(KIMCHI_API_KEY=key-coder collect) +[[ $(jq -r '.totalSessions' <<<"$record") == 4 ]] || fail "an ignoring git checkout in HOME doesn't hide sessions" "$record" +pass "an ignoring git checkout in HOME doesn't hide sessions" + +# A scan that couldn't read every session isn't kept for the next run. +if (( EUID != 0 )); then + chmod 000 "$fork/fork.jsonl" + rm -f "$XDG_CACHE_HOME"/omarchy/agent-usage/kimchi-scan-*.json + record=$(KIMCHI_API_KEY=key-coder collect) + chmod 644 "$fork/fork.jsonl" + [[ $(jq -r '.totalSessions' <<<"$record") == 3 ]] || fail "an unreadable session is skipped" "$record" + ! compgen -G "$XDG_CACHE_HOME/omarchy/agent-usage/kimchi-scan-*.json" >/dev/null || + fail "a partial scan is not cached" "$(ls "$XDG_CACHE_HOME/omarchy/agent-usage")" + pass "a scan that couldn't read every session isn't cached" +fi From 08043cb138dfe071bfa03f9e11892f1986819fcc Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:08:32 -0400 Subject: [PATCH 048/546] Bring the Muse collector in line with the usage collector contract Pay-as-you-go no longer invents an Account tier, a sign-in counts as in use, the Muse CLI's own home without a login waits for auth, probes are reused for 15 seconds, and untrusted timestamps and token fields skip the field instead of landing on today or counting as tokens. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-muse | 131 +++++++++++------- test/shell.d/agent-usage-muse-limits-test.sh | 44 +++++- test/shell.d/agent-usage-muse-scanner-test.sh | 11 +- 3 files changed, 124 insertions(+), 62 deletions(-) diff --git a/bin/omarchy-agent-usage-muse b/bin/omarchy-agent-usage-muse index 0639df6759a..6e53f1cf666 100755 --- a/bin/omarchy-agent-usage-muse +++ b/bin/omarchy-agent-usage-muse @@ -23,8 +23,8 @@ neither it nor the returned API key is persisted, and the login file is only ever read, never refreshed or rewritten. A failed check keeps the last good windows, marked stale. Pay-as-you-go accounts and active plans that return no `subs_usage` (e.g. Muse Code Power Usage) -have no subscription windows, so limits stay empty — tier plus local -metering still show. No estimated meters from manual caps: without an +have no subscription windows, so limits stay empty — the plan name (none for +pay-as-you-go) plus local metering still show. No estimated meters from manual caps: without an authoritative probe the panel shows measured usage only. The agents panel only ever reads the JSON this prints. """ @@ -34,6 +34,7 @@ import base64 import fcntl import hashlib import json +import math import os import re import sqlite3 @@ -58,14 +59,17 @@ MAX_SESSION_BYTES = 512 * 1024 * 1024 # caches are rescanned instead of trusted. CACHE_SCHEMA_VERSION = 2 +# A token count past this is a corrupt field, not a model call. +MAX_TOKENS = 1_000_000_000_000 + # Same cache semantics as the codex collector: a scan this recent only dedups # concurrent collector runs; --limits-only may reuse a scan for 15 minutes. SCAN_REUSE_SECONDS = 20 LIMITS_ONLY_REUSE_SECONDS = 900 -# Subscription probe reuse: polling mints the same key (idempotent) but the -# answer moves slowly, so repeated panel opens reuse a recent probe. -PROBE_MIN_INTERVAL_SECONDS = 60 +# Subscription probe reuse: minting is idempotent and cheap (the same key +# comes back every call), so this only absorbs repeated panel opens. +PROBE_MIN_INTERVAL_SECONDS = 15 def key_endpoint() -> str: @@ -81,8 +85,24 @@ def number(value): return 0 +# Token fields come from local files: a bool, string, NaN, negative or +# implausibly huge value skips the field. def token_number(value): - return max(0, number(value)) + if isinstance(value, bool) or not isinstance(value, (int, float)): + return 0 + if not math.isfinite(value) or value <= 0 or value >= MAX_TOKENS: + return 0 + return int(value) + + +def percent_number(value): + if isinstance(value, bool): + return -1.0 + try: + parsed = float(value) + except (TypeError, ValueError): + return -1.0 + return min(1.0, max(0.0, parsed / 100.0)) if math.isfinite(parsed) else -1.0 def model_name(raw): @@ -98,39 +118,46 @@ def model_name(raw): def local_day(value): - """Normalize assorted timestamps to a local YYYY-MM-DD date string.""" - # bool is an int subclass: True would otherwise file under 1970. + """Normalize assorted timestamps to a local YYYY-MM-DD date string. + + A stamp that isn't a plausible moment (missing, zero, negative, NaN, a + bool, far future, not a time at all) gives None, so the caller can fall + back or leave the row undated rather than file it under today. + """ if value is None or isinstance(value, bool): - return datetime.now().astimezone().date().isoformat() + return None if isinstance(value, (int, float)): seconds = float(value) + if not math.isfinite(seconds) or seconds <= 0: + return None # Native muse recorded_at is microseconds; opencode created is millis; # codex-style seconds also possible. if seconds > 10_000_000_000_000: seconds = seconds / 1_000_000 elif seconds > 10_000_000_000: seconds = seconds / 1000 + else: + text = str(value).strip() try: - return datetime.fromtimestamp(seconds).astimezone().date().isoformat() - except Exception: - return datetime.now().astimezone().date().isoformat() - text = str(value).strip() - try: - # Some writers emit the epoch as a string: route micro/milli/second - # magnitudes through the numeric branch instead of misfiling them as today. - return local_day(float(text)) - except ValueError: - pass - try: - if text.endswith("Z"): - dt = datetime.fromisoformat(text[:-1] + "+00:00") - else: - dt = datetime.fromisoformat(text) - if dt.tzinfo is not None: + # Some writers emit the epoch as a string: route micro/milli/second + # magnitudes through the numeric branch instead of misfiling them. + return local_day(float(text)) + except ValueError: + pass + try: + dt = datetime.fromisoformat(text[:-1] + "+00:00" if text.endswith("Z") else text) + except ValueError: + return None + if dt.tzinfo is None: dt = dt.astimezone() - return dt.strftime("%Y-%m-%d") - except Exception: - return datetime.now().astimezone().date().isoformat() + seconds = dt.timestamp() + # Muse didn't exist before 2000 and can't write from the future. + if not 946_684_800 <= seconds <= time.time() + 2 * 86_400: + return None + try: + return datetime.fromtimestamp(seconds).astimezone().date().isoformat() + except (OverflowError, OSError, ValueError): + return None now = datetime.now().astimezone() @@ -148,6 +175,7 @@ total_prompts = 0 total_sessions = set() +# An undated row (day None) counts in the totals only. def add_usage(day, session_key, model, input_tokens, output_tokens, cache_read, cache_write): global today_prompts, today_total_tokens, total_prompts total = input_tokens + output_tokens + cache_read + cache_write @@ -155,7 +183,8 @@ def add_usage(day, session_key, model, input_tokens, output_tokens, cache_read, return total_prompts += 1 total_sessions.add(session_key) - active_days.add(day) + if day is not None: + active_days.add(day) bucket = model_usage.setdefault(model, { "inputTokens": 0, @@ -315,7 +344,8 @@ def scan_native_muse_sessions(): if event_id in seen_event_ids: continue seen_event_ids.add(event_id) - day = local_day(recorded_at) if recorded_at else mtime_day + # An unusable stamp falls back on the file's own time. + day = local_day(recorded_at) or mtime_day # Subagent trees roll up into their parent session: normalize both the # main session.jsonl and any subagent/.../session.jsonl to the session # directory, so one user session counts once however many subagents @@ -380,9 +410,9 @@ def add_opencode_row(message_id, session_id, entry, row_type, created, seen): return seen.add(message_id) # A malformed time object must not drop the row's valid tokens: non-dict - # shapes fall back to the row's own time, then today. + # shapes fall back to the row's own time, then leave the row undated. stamped = entry.get("time") - day = local_day((stamped.get("created") if isinstance(stamped, dict) else None) or created) + day = local_day(stamped.get("created") if isinstance(stamped, dict) else None) or local_day(created) add_usage(day, "opencode:" + str(session_id), model, input_tokens, output_tokens, cache_read, cache_write) @@ -484,7 +514,7 @@ def read_access_token() -> str: def normalize_reset_at(value) -> str: - if value is None: + if value is None or isinstance(value, bool): return "" try: ts = float(value) @@ -554,20 +584,14 @@ def probe_key(access_token: str) -> dict: limits = [] window = usage.get("window") or {} if isinstance(window, dict) and window.get("used_percent") is not None: - try: - percent = min(1.0, max(0.0, float(window.get("used_percent")) / 100.0)) - except Exception: - percent = -1.0 + percent = percent_number(window.get("used_percent")) if percent >= 0: mins = number(window.get("window_duration_mins")) label = "Session (5-hour)" if mins == 300 else (f"Session ({mins}m)" if mins > 0 else "Session") limits.append({"label": label, "percent": percent, "resetsAt": normalize_reset_at(window.get("resets_at"))}) weekly = usage.get("weekly") or {} if isinstance(weekly, dict) and weekly.get("used_percent") is not None: - try: - percent = min(1.0, max(0.0, float(weekly.get("used_percent")) / 100.0)) - except Exception: - percent = -1.0 + percent = percent_number(weekly.get("used_percent")) if percent >= 0: limits.append({"label": "Weekly (7-day)", "percent": percent, "resetsAt": normalize_reset_at(weekly.get("resets_at"))}) @@ -629,9 +653,14 @@ def collect_limits(force: bool) -> dict: access_token = read_access_token() if access_token == "": - # No login: PAYG or signed-out. Never estimated, never an auth nag — - # local metering is valid on its own. + # No login. Muse Spark used only through opencode's meta provider signs + # in with a Model API key, not `muse login`, so that alone is no reason + # to ask for one; the Muse CLI's own home on this machine is. + if muse_config_dir().is_dir() or muse_sessions_root().is_dir(): + result["usageStatusText"] = "Waiting for auth" + result["authHelpText"] = "Run `muse login` to show Muse's session and weekly limits." return result + result["signedIn"] = True probe_cache = None cached = {} @@ -644,7 +673,7 @@ def collect_limits(force: bool) -> dict: cached = {} fallback = usable_cached_limits(cached) fetched_at = number(cached.get("fetchedAtMs")) - result["tierLabel"] = str(cached.get("tierLabel") or "") or "Account" + result["tierLabel"] = str(cached.get("tierLabel") or "") # The last good windows, marked stale, with when they were measured. def keep_cached(): @@ -672,9 +701,10 @@ def collect_limits(force: bool) -> dict: if probe.get("transport"): result["retryAdvised"] = True if probe.get("payg"): - # Authenticated PAYG: no subscription windows exist. Keep local metering - # only, silent — not an error, not expired. - result["tierLabel"] = "Account" + # Authenticated PAYG: no subscription windows exist, and the service + # names no plan. Keep local metering only, silent — not an error, not + # expired. + result["tierLabel"] = "" return result if probe.get("silent"): # Active plan with no windows: same silence as PAYG, but keep the plan @@ -713,7 +743,7 @@ def read_fresh_json(path, max_age_seconds): return None try: age = time.time() - path.stat().st_mtime - if 0 <= age <= max_age_seconds: + if 0 <= age < max_age_seconds: return json.loads(path.read_text(encoding="utf-8")) except Exception: return None @@ -820,8 +850,9 @@ def main(): "id": AGENT_ID, "name": AGENT_NAME, "updatedAt": datetime.now(timezone.utc).isoformat(), - # A signed-in plan shows even before its first session or window. - "ready": number(stats.get("totalPrompts")) > 0 or len(limits["limits"]) > 0 or limits["tierLabel"] not in ("", "Account"), + # A sign-in shows even before its first session or window, and so does + # a Muse CLI waiting for one. + "ready": number(stats.get("totalPrompts")) > 0 or len(limits["limits"]) > 0 or bool(limits.get("signedIn")) or limits["usageStatusText"] != "", "hasLocalStats": True, # Prompts and sessions are counted, so the panel shows them in tooltips. "hasPromptStats": True, diff --git a/test/shell.d/agent-usage-muse-limits-test.sh b/test/shell.d/agent-usage-muse-limits-test.sh index 959807dc773..12e2c40b869 100755 --- a/test/shell.d/agent-usage-muse-limits-test.sh +++ b/test/shell.d/agent-usage-muse-limits-test.sh @@ -161,10 +161,11 @@ pass "Muse collector resets a kept window once it has ended" pass "Muse collector never writes the login file" # Pay-as-you-go: key mints fine but there is no subscription window to -# meter. Silent local metering — not an error, not expired. +# meter, and no plan name to invent. Silent local metering — not an error, +# not expired. set_payload '{"body":{"api_key":"minted","is_subs_active":false,"subs_tier_name":null}}' payg=$(run_collector --force) -[[ $(jq -c '[.limits,.tierLabel,.usageStatusText,.authHelpText]' <<<"$payg") == '[[],"Account","",""]' ]] || +[[ $(jq -c '[.limits,.tierLabel,.usageStatusText,.authHelpText]' <<<"$payg") == '[[],"","",""]' ]] || fail "Muse collector stays silent for pay-as-you-go accounts" "$payg" pass "Muse collector stays silent for pay-as-you-go accounts" @@ -180,9 +181,8 @@ nosnap=$(run_collector --force) fail "Muse collector stays silent for an active plan without usage windows" "$nosnap" pass "Muse collector stays silent for an active plan without usage windows" -# A signed-in plan shows before its first session; pay-as-you-go with no -# usage has nothing to show. -[[ $(jq -r '.ready' <<<"$nosnap") == "true" && $(jq -r '.ready' <<<"$payg") == "false" ]] || +# A sign-in shows before its first session, pay-as-you-go included. +[[ $(jq -r '.ready' <<<"$nosnap") == "true" && $(jq -r '.ready' <<<"$payg") == "true" ]] || fail "Muse collector shows a signed-in plan without sessions" "$nosnap" pass "Muse collector shows a signed-in plan without sessions" @@ -204,12 +204,13 @@ expired=$(run_collector --force) fail "Muse collector reports an expired sign-in" "$expired" pass "Muse collector reports an expired sign-in" -# A malformed login file is a missing credential, not a crash. +# A malformed login file is a missing credential, not a crash: the Muse CLI +# is set up here, so it waits for a sign-in. for malformed in '[1,2]' '"nope"' '{"providers":[]}' '{"providers":{"meta":7}}'; do printf '%s' "$malformed" >"$TEST_HOME/.config/muse/auth.json" garbled=$(run_collector --force) || fail "Muse collector survives a malformed login file" "$malformed" - [[ $(jq -c '[.limits,.usageStatusText]' <<<"$garbled") == '[[],""]' ]] || + [[ $(jq -c '[.limits,.usageStatusText,.ready]' <<<"$garbled") == '[[],"Waiting for auth",true]' ]] || fail "Muse collector survives a malformed login file" "$malformed -> $garbled" done pass "Muse collector survives a malformed login file" @@ -225,3 +226,32 @@ custom=$(run_collector --force) fail "Muse collector honors MUSE_AUTH_PATH over the default login" "$(cat "$STUB_SEEN_FILE")" pass "Muse collector honors MUSE_AUTH_PATH over the default login" unset MUSE_AUTH_PATH + +# A probe this recent is reused, and --force asks again. +run_collector --force >/dev/null +: >"$STUB_SEEN_FILE" +run_collector --limits-only >/dev/null +[[ ! -s $STUB_SEEN_FILE ]] || fail "Muse collector reuses a probe under 15 seconds old" "$(cat "$STUB_SEEN_FILE")" +run_collector --force >/dev/null +[[ -s $STUB_SEEN_FILE ]] || fail "Muse collector re-probes on --force" "nothing seen" +pass "Muse collector reuses a recent probe, and --force asks again" + +# A server that answered with an error isn't retried sooner; one that can't be +# reached at all is. +printf '{"status":500}' >"$STUB_PAYLOAD_FILE" +errored=$(run_collector --force) +[[ $(jq 'has("retryAdvised")' <<<"$errored") == "false" ]] || + fail "Muse collector doesn't retry a server error sooner" "$errored" +closed=$(python3 -c 'import socket; s=socket.socket(); s.bind(("127.0.0.1",0)); print(s.getsockname()[1])') +unreachable=$(HOME="$TEST_HOME" XDG_DATA_HOME="$TEST_HOME/.local/share" XDG_CACHE_HOME="$TEST_HOME/.cache" XDG_CONFIG_HOME="$TEST_HOME/.config" \ + MUSE_KEY_ENDPOINT="http://127.0.0.1:$closed/key" "$ROOT/bin/omarchy-agent-usage-muse" --force) +[[ $(jq -c '{retry: .retryAdvised, stale: .limitsStale, kept: (.limits | length > 0), status: .usageStatusText}' <<<"$unreachable") == '{"retry":true,"stale":true,"kept":true,"status":""}' ]] || + fail "Muse collector retries an unreachable endpoint sooner, keeping the last limits" "$unreachable" +pass "Muse collector advises a retry only when the endpoint can't be reached" + +# A machine that never ran Muse prints an unused record. +unused=$(HOME="$TEST_HOME/nobody" XDG_DATA_HOME="$TEST_HOME/nobody/data" XDG_CACHE_HOME="$TEST_HOME/nobody/cache" XDG_CONFIG_HOME="$TEST_HOME/nobody/config" \ + MUSE_KEY_ENDPOINT="http://127.0.0.1:$port/key" "$ROOT/bin/omarchy-agent-usage-muse") +[[ $(jq -c '{ready, tierLabel, usageStatusText, authHelpText, limits, totalPrompts}' <<<"$unused") == '{"ready":false,"tierLabel":"","usageStatusText":"","authHelpText":"","limits":[],"totalPrompts":0}' ]] || + fail "Muse collector prints an unused record on a machine without Muse" "$unused" +pass "Muse collector prints an unused record on a machine without Muse" diff --git a/test/shell.d/agent-usage-muse-scanner-test.sh b/test/shell.d/agent-usage-muse-scanner-test.sh index 887e9418a43..9f84dd62598 100755 --- a/test/shell.d/agent-usage-muse-scanner-test.sh +++ b/test/shell.d/agent-usage-muse-scanner-test.sh @@ -320,10 +320,11 @@ pass "Muse collector honors --limits-only cache reuse" # Test 17: Odd timestamp and token shapes cannot misfile or drop usage # A microsecond epoch emitted as a string still lands on its real day (8 days -# ago, outside the 7-day window); a boolean timestamp falls back to today -# instead of 1970; a non-dict opencode time object keeps its row's tokens -# under today; a non-dict tokens object skips just that row; and a non-dict -# cache object keeps the row's input. Today: 5 + 44 + 11 = 60. +# ago, outside the 7-day window); a boolean timestamp falls back to the +# file's own time (today) instead of 1970; a non-dict opencode time object +# keeps its row's tokens in the totals but on no day, rather than guessing +# today; a non-dict tokens object skips just that row; and a non-dict cache +# object keeps the row's input. Today: 5 + 11 = 16. odd_home="$TEST_HOME/odd-shapes" mkdir -p "$odd_home/muse-data/sessions/sess-odd" "$odd_home/data/opencode" cat >"$odd_home/muse-data/sessions/sess-odd/session.jsonl" < Date: Wed, 14 Oct 2026 23:09:15 -0400 Subject: [PATCH 049/546] Read Muse's current session logs and session index Completions logged as record_type event with no payload kind count, cache tokens never claim more than the input they are part of, a session window of whole hours reads in hours, real prompt and session counts come from the CLI's session-index.db (falling back on completions without it), and lines without usage skip the JSON parse. Co-authored-by: Nicolas Desjardins Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-muse | 114 ++++++++++++++---- test/shell.d/agent-usage-muse-limits-test.sh | 7 ++ test/shell.d/agent-usage-muse-scanner-test.sh | 44 +++++++ 3 files changed, 142 insertions(+), 23 deletions(-) diff --git a/bin/omarchy-agent-usage-muse b/bin/omarchy-agent-usage-muse index 6e53f1cf666..7421187ff9b 100755 --- a/bin/omarchy-agent-usage-muse +++ b/bin/omarchy-agent-usage-muse @@ -8,6 +8,8 @@ Local metering (always available, no network): token counts, prompt and session counts, per-day and per-model breakdowns. Sources (all local): - native `muse` CLI sessions in $XDG_DATA_HOME/muse/sessions, defaulting to ~/.local/share/muse/sessions (`model_completed` events carry usage), + with the real prompt and session counts from the CLI's own + session-index.db beside them when it is there, - opencode sessions with providerID == "meta" (any Muse Spark model). Muse Spark served through OpenCode Zen (providerID == "opencode") is Zen usage, and the opencode record counts it. @@ -56,8 +58,9 @@ MAX_MODEL_NAME_LEN = 120 MAX_SESSION_BYTES = 512 * 1024 * 1024 # Cache schema version. Bump when the cached stats shape changes so old -# caches are rescanned instead of trusted. -CACHE_SCHEMA_VERSION = 2 +# caches are rescanned instead of trusted. 3: prompts and sessions from the +# session index. +CACHE_SCHEMA_VERSION = 3 # A token count past this is a corrupt field, not a model call. MAX_TOKENS = 1_000_000_000_000 @@ -175,14 +178,18 @@ total_prompts = 0 total_sessions = set() -# An undated row (day None) counts in the totals only. -def add_usage(day, session_key, model, input_tokens, output_tokens, cache_read, cache_write): +# `counted` is False for native completions once the session index has +# supplied the real prompt and session counts: their tokens still count, but +# a completion is a model call, not a prompt. An undated row (day None) +# counts in the totals only. +def add_usage(day, session_key, model, input_tokens, output_tokens, cache_read, cache_write, counted=True): global today_prompts, today_total_tokens, total_prompts total = input_tokens + output_tokens + cache_read + cache_write if total <= 0: return - total_prompts += 1 - total_sessions.add(session_key) + if counted: + total_prompts += 1 + total_sessions.add(session_key) if day is not None: active_days.add(day) @@ -200,19 +207,66 @@ def add_usage(day, session_key, model, input_tokens, output_tokens, cache_read, if day in recent: recent[day] += total if day == today: - today_prompts += 1 - today_sessions.add(session_key) + if counted: + today_prompts += 1 + today_sessions.add(session_key) today_total_tokens += total today_tokens_by_model[model] = today_tokens_by_model.get(model, 0) + total -def muse_sessions_root(): +def muse_data_dir(): override = (os.environ.get("MUSE_DATA_DIR") or "").strip() if override: - base = Path(os.path.expandvars(os.path.expanduser(override))) - else: - base = Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) / "muse" - return base / "sessions" + return Path(os.path.expandvars(os.path.expanduser(override))) + return Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) / "muse" + + +# The Muse CLI keeps one row per session in session-index.db with the number +# of prompts typed into it, the real counts a completion tally can only +# approximate (one prompt can run many model calls). Returns whether the +# index supplied them, and whether it could be read: an index that's there +# but unreadable falls back on completions without caching that scan. A +# session nobody prompted isn't counted, and a row whose creation time isn't +# a plausible moment counts in the totals only. +def scan_session_index(): + global today_prompts, total_prompts + database = muse_data_dir() / "session-index.db" + if not database.is_file(): + return False, True + try: + conn = sqlite3.connect(database.resolve().as_uri() + "?mode=ro", uri=True, timeout=2) + except sqlite3.Error: + return False, False + try: + conn.execute("PRAGMA query_only = ON") + rows = conn.execute("SELECT session_id, created_at_us, prompt_count FROM sessions").fetchall() + except sqlite3.Error: + return False, False + finally: + try: + conn.close() + except Exception: + pass + for session_id, created_at_us, prompt_count in rows: + prompts = token_number(prompt_count) + if not prompts: + continue + session_key = "muse-index:" + str(session_id) + total_prompts += prompts + total_sessions.add(session_key) + # created_at_us is microseconds, which local_day recognises by size. + day = local_day(created_at_us) + if day is None: + continue + active_days.add(day) + if day == today: + today_prompts += prompts + today_sessions.add(session_key) + return True, True + + +def muse_sessions_root(): + return muse_data_dir() / "sessions" def iter_native_events(path): @@ -224,6 +278,12 @@ def iter_native_events(path): try: with path.open(errors="replace") as handle: for raw in handle: + # Session logs run to hundreds of megabytes, mostly lines without + # usage. Only completions ("usage", escaped or not inside a + # retained frame) and the model records ("model_id") matter, so + # everything else skips the JSON parse. + if "usage" not in raw and "model_id" not in raw: + continue raw = raw.strip() if not raw: continue @@ -257,7 +317,7 @@ def iter_native_events(path): return -def scan_native_muse_sessions(): +def scan_native_muse_sessions(counted=True): root = muse_sessions_root() if not root.exists(): return @@ -309,7 +369,10 @@ def scan_native_muse_sessions(): if isinstance(record, dict) and record.get("model_id"): current_model = model_name(record.get("model_id")) continue - if kind != "run": + # Older logs wrap a completion as payload kind "run"; current ones + # (record_type "event", payload_type "runtime.session") carry the event + # with no payload kind at all. + if kind not in ("run", None): continue event = payload.get("event") or {} if not isinstance(event, dict): @@ -330,11 +393,12 @@ def scan_native_muse_sessions(): input_tokens = token_number(usage.get("input_tokens")) # Native Muse reports reasoning as a subset of output_tokens. output_tokens = token_number(usage.get("output_tokens")) - cache_read = token_number(usage.get("cache_read_tokens", usage.get("cached_tokens"))) - cache_write = token_number(usage.get("cache_write_tokens")) # Cached tokens are a subset of input: keep the split without double - # counting (same convention as the codex collector). - input_tokens = max(0, input_tokens - cache_read - cache_write) + # counting (same convention as the codex collector), and never let a + # cache figure claim more than the input it is part of. + cache_read = min(input_tokens, token_number(usage.get("cache_read_tokens", usage.get("cached_tokens")))) + cache_write = min(input_tokens - cache_read, token_number(usage.get("cache_write_tokens"))) + input_tokens = input_tokens - cache_read - cache_write if not (input_tokens or output_tokens or cache_read or cache_write): continue # Native record IDs survive copies and retained-frame wrapping. Claim @@ -352,7 +416,7 @@ def scan_native_muse_sessions(): # it spawned. session_key = "muse:" + re.sub(r"(/subagent/.*)?/session\.jsonl$", "", str(path)) add_usage(day, session_key, model, - input_tokens, output_tokens, cache_read, cache_write) + input_tokens, output_tokens, cache_read, cache_write, counted) def opencode_db(): @@ -587,7 +651,10 @@ def probe_key(access_token: str) -> dict: percent = percent_number(window.get("used_percent")) if percent >= 0: mins = number(window.get("window_duration_mins")) - label = "Session (5-hour)" if mins == 300 else (f"Session ({mins}m)" if mins > 0 else "Session") + if mins >= 60 and mins % 60 == 0: + label = f"Session ({mins // 60}-hour)" + else: + label = f"Session ({mins}m)" if mins > 0 else "Session" limits.append({"label": label, "percent": percent, "resetsAt": normalize_reset_at(window.get("resets_at"))}) weekly = usage.get("weekly") or {} if isinstance(weekly, dict) and weekly.get("used_percent") is not None: @@ -805,9 +872,10 @@ def local_stats(): def run_local_scans(): - scan_native_muse_sessions() + indexed, index_read = scan_session_index() + scan_native_muse_sessions(counted=not indexed) complete = scan_opencode_sessions() - return local_stats(), complete + return local_stats(), complete and index_read def cached_local_stats(max_age): diff --git a/test/shell.d/agent-usage-muse-limits-test.sh b/test/shell.d/agent-usage-muse-limits-test.sh index 12e2c40b869..34bba696a24 100755 --- a/test/shell.d/agent-usage-muse-limits-test.sh +++ b/test/shell.d/agent-usage-muse-limits-test.sh @@ -227,6 +227,13 @@ custom=$(run_collector --force) pass "Muse collector honors MUSE_AUTH_PATH over the default login" unset MUSE_AUTH_PATH +# A session window that is a whole number of hours says so in hours. +set_payload '{"body":{"subs_tier_name":"T","subs_usage":{"window":{"used_percent":5,"window_duration_mins":120,"resets_at":4102444800},"weekly":{"used_percent":1,"resets_at":4103049600}}}}' +hours=$(run_collector --force) +[[ $(jq -r '.limits[0].label' <<<"$hours") == "Session (2-hour)" ]] || + fail "Muse collector labels a whole-hour session window in hours" "$hours" +pass "Muse collector labels a whole-hour session window in hours" + # A probe this recent is reused, and --force asks again. run_collector --force >/dev/null : >"$STUB_SEEN_FILE" diff --git a/test/shell.d/agent-usage-muse-scanner-test.sh b/test/shell.d/agent-usage-muse-scanner-test.sh index 9f84dd62598..2e4fc6bb971 100755 --- a/test/shell.d/agent-usage-muse-scanner-test.sh +++ b/test/shell.d/agent-usage-muse-scanner-test.sh @@ -370,3 +370,47 @@ result_v2=$(HOME="$v2_home" XDG_DATA_HOME="$v2_home/data" XDG_CACHE_HOME="$v2_ho [[ $(jq -c '{todayTotalTokens, totalPrompts, totalSessions}' <<<"$result_v2") == '{"todayTotalTokens":115,"totalPrompts":2,"totalSessions":1}' ]] || fail "Muse collector reads OpenCode V2 records once each" "$result_v2" pass "Muse collector reads OpenCode V2 records once each" + +# Current Muse logs carry a completion as record_type "event" with no payload +# kind, and report cached tokens inside input_tokens; a cache figure larger +# than its input never claims more than the input. +index_home="$TEST_HOME/index" +mkdir -p "$index_home/muse-data/sessions/2026/10/12/demo" +cat >"$index_home/muse-data/sessions/2026/10/12/demo/session.jsonl" <"$index_home/muse-data/session-index.db" +rm -rf "$index_home/cache" +result_badindex=$("${index_env[@]}" "$ROOT/bin/omarchy-agent-usage-muse") +[[ $(jq -c '{totalPrompts, totalSessions}' <<<"$result_badindex") == '{"totalPrompts":3,"totalSessions":1}' ]] || + fail "Muse collector counts completions past an unreadable session index" "$result_badindex" +! compgen -G "$index_home/cache/omarchy/agent-usage/muse-scan-*.json" >/dev/null || + fail "Muse collector doesn't cache a scan past an unreadable session index" "$(ls "$index_home/cache/omarchy/agent-usage")" +pass "Muse collector falls back on completions past an unreadable session index" From 587ddb616cd5cb7d55272627592f6900921f61e5 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:02:22 -0400 Subject: [PATCH 050/546] Bring the Cursor collector in line with the usage collector contract A limits answer is reused for 15 seconds and a token scan for 20, both bypassed by --force. An unreachable dashboard advises a retry, a refused sign-in reads "Sign-in expired", a lapsed one "Limits paused" (and gives way to a live editor sign-in), and any other failure keeps the last limits, plan and tokens. A healthy record clears its help text, unknown resets are "", the meters are titled after the monthly billing period, and a machine without Cursor installed prints an unused record. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-cursor | 275 +++++++++++++++++------- test/shell.d/agent-usage-cursor-test.sh | 189 +++++++++++++++- 2 files changed, 375 insertions(+), 89 deletions(-) diff --git a/bin/omarchy-agent-usage-cursor b/bin/omarchy-agent-usage-cursor index 3f72f39b092..5a67bfc445d 100755 --- a/bin/omarchy-agent-usage-cursor +++ b/bin/omarchy-agent-usage-cursor @@ -17,7 +17,9 @@ from __future__ import annotations import argparse import base64 import hashlib +import http.client import json +import math import os import sqlite3 import sys @@ -36,39 +38,60 @@ API_BASE_URL = "https://api2.cursor.sh" SERVICE = "aiserver.v1.DashboardService" REQUEST_TIMEOUT = 20 RECENT_DAYS = 7 +# Repeated panel opens reuse a limits answer this recent instead of asking +# again; --force (a person pressing refresh) always asks. +PROBE_MIN_INTERVAL_SECONDS = 15 +# A normal run reuses a token scan this recent, so runs started together +# don't ask the dashboard for the same week twice. +SCAN_REUSE_SECONDS = 20 # How long --limits-only serves the last token scan before asking again. LIMITS_ONLY_REUSE_SECONDS = 900 class CursorError(Exception): - def __init__(self, message: str, auth: bool = False) -> None: + def __init__(self, message: str, auth: bool = False, transport: bool = False) -> None: super().__init__(message) self.auth = auth + # Nothing answered at all, so the panel should try again soon. + self.transport = transport -def number(value: Any) -> int: +def finite(value: Any) -> float: + if isinstance(value, bool): + return 0.0 try: - return max(0, round(float(value or 0))) + parsed = float(value or 0) except (TypeError, ValueError): - return 0 + return 0.0 + return parsed if math.isfinite(parsed) else 0.0 + + +def number(value: Any) -> int: + return max(0, round(finite(value))) def fraction(value: Any) -> float: # Cursor reports whole percentages; the panel's meters take 0..1. - try: - return max(0.0, float(value or 0) / 100) - except (TypeError, ValueError): - return 0.0 + return min(1.0, max(0.0, finite(value) / 100)) + + +def config_home() -> Path: + return Path(os.environ.get("XDG_CONFIG_HOME") or (Path.home() / ".config")) def default_auth_path() -> Path: - config_home = Path(os.environ.get("XDG_CONFIG_HOME") or (Path.home() / ".config")) - return config_home / "cursor" / "auth.json" + return config_home() / "cursor" / "auth.json" def default_state_db() -> Path: - config_home = Path(os.environ.get("XDG_CONFIG_HOME") or (Path.home() / ".config")) - return config_home / "Cursor" / "User" / "globalStorage" / "state.vscdb" + return config_home() / "Cursor" / "User" / "globalStorage" / "state.vscdb" + + +# The Cursor CLI keeps its config beside auth.json and in ~/.cursor, the editor +# under ~/.config/Cursor. Any of them means Cursor is used here, so a missing +# sign-in is worth saying; without them the machine doesn't use Cursor at all. +def cursor_in_use(auth_path: Path) -> bool: + return any(path.is_dir() for path in (auth_path.parent, config_home() / "Cursor", Path.home() / ".cursor")) def read_ide_token(state_db: Path) -> str: @@ -96,26 +119,61 @@ def read_ide_token(state_db: Path) -> str: return str(value or "").strip() +def read_cli_token(path: Path) -> str: + try: + data = json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError): + return "" + return str(data.get("accessToken") or "").strip() if isinstance(data, dict) else "" + + +def jwt_claims(token: str) -> dict[str, Any]: + parts = token.split(".") + if len(parts) != 3: + return {} + try: + claims = json.loads(base64.urlsafe_b64decode(parts[1] + "=" * (-len(parts[1]) % 4))) + except (ValueError, UnicodeDecodeError): + return {} + return claims if isinstance(claims, dict) else {} + + +# A sign-in token whose expiry has passed. The CLI and the editor each renew +# their own the next time they run; an opaque API key never lapses here. +def token_lapsed(token: str) -> bool: + expires = finite(jwt_claims(token).get("exp")) + return 0 < expires <= datetime.now(timezone.utc).timestamp() + + def read_token(path: Path, state_db: Path | None = None) -> str: """The key the user exported, the credential the Cursor CLI stores, or the - editor's sign-in.""" + editor's sign-in, in that order. + + A sign-in that has lapsed gives way to the next one that hasn't: someone who + ran the CLI once and now lives in the editor would otherwise wait on a + token nothing renews. With all of them lapsed, the first one decides. The + Grok Bot collector reads the same account the same way. + """ key = str(os.environ.get("CURSOR_API_KEY") or "").strip() if key: return key - token = "" - if path.is_file(): - try: - token = str(json.loads(path.read_text()).get("accessToken") or "").strip() - except (OSError, ValueError, AttributeError): - token = "" - return token or read_ide_token(state_db or default_state_db()) + cli = read_cli_token(path) + if cli and not token_lapsed(cli): + return cli + ide = read_ide_token(state_db or default_state_db()) + if ide and not token_lapsed(ide): + return ide + return cli or ide -def iso(epoch_millis: Any) -> str | None: +def iso(epoch_millis: Any) -> str: millis = number(epoch_millis) if millis <= 0: - return None - return datetime.fromtimestamp(millis / 1000, timezone.utc).isoformat() + return "" + try: + return datetime.fromtimestamp(millis / 1000, timezone.utc).isoformat() + except (OverflowError, OSError, ValueError): + return "" def empty_bucket() -> dict[str, int]: @@ -191,23 +249,30 @@ class CursorClient: except urllib.error.HTTPError as error: if error.code in (401, 403): raise CursorError(AUTH_HELP, auth=True) from error - raise CursorError(f"Cursor API returned {error.code}") from error - except (urllib.error.URLError, TimeoutError) as error: - raise CursorError("Cursor API unreachable") from error + raise CursorError(f"Cursor's dashboard returned status {error.code}.") from error + except http.client.HTTPException as error: + raise CursorError("Cursor's dashboard cut its reply short.") from error + # No route, no DNS, refused, reset or timed out: nothing answered. + except OSError as error: + raise CursorError("Couldn't reach Cursor's dashboard. Retrying shortly.", transport=True) from error try: - return json.loads(body or "{}") + parsed = json.loads(body or "{}") except ValueError as error: - raise CursorError("Cursor API returned malformed JSON") from error + raise CursorError("Cursor's dashboard returned malformed JSON.") from error + if not isinstance(parsed, dict): + raise CursorError("Cursor's dashboard returned malformed JSON.") + return parsed def period(self) -> dict[str, Any]: return self.call("GetCurrentPeriodUsage", {}) - def plan_name(self) -> str: + # None when the plan couldn't be asked, so the last known one stands. + def plan_name(self) -> str | None: try: plan = self.call("GetPlanInfo", {}).get("planInfo") except CursorError: - return "" - return str((plan or {}).get("planName") or "") + return None + return str(plan.get("planName") or "") if isinstance(plan, dict) else "" def aggregations(self, start_millis: int, end_millis: int) -> list[Any]: payload = {"startDate": str(start_millis), "endDate": str(end_millis)} @@ -301,10 +366,10 @@ def carried_stats(key: str, today: date) -> dict[str, Any]: return stats -def carried_stats_fresh(key: str, today: date) -> bool: +def carried_stats_fresh(key: str, today: date, max_age_seconds: float) -> bool: cached = read_json(stats_cache(key)) age = now_ms() - number(cached.get("scannedAtMs")) - return cached.get("scanDate") == today.isoformat() and 0 <= age < LIMITS_ONLY_REUSE_SECONDS * 1000 + return cached.get("scanDate") == today.isoformat() and 0 <= age < max_age_seconds * 1000 def collect_stats(client: CursorClient, today: date, key: str) -> dict[str, Any]: @@ -312,16 +377,37 @@ def collect_stats(client: CursorClient, today: date, key: str) -> dict[str, Any] with ThreadPoolExecutor(max_workers=4) as pool: rows = list(pool.map(lambda window: client.aggregations(window[1], window[2]), windows)) stats = summarize_usage([(window[0], summarize_day(day_rows)) for window, day_rows in zip(windows, rows)]) - write_json(stats_cache(key), {"scanDate": today.isoformat(), "scannedAtMs": now_ms(), "stats": stats}) + # The cache only spares later runs a fetch; failing to keep it never costs + # this run's answer. + try: + write_json(stats_cache(key), {"scanDate": today.isoformat(), "scannedAtMs": now_ms(), "stats": stats}) + except OSError: + pass return stats +# This week's tokens: a scan of this account recent enough to reuse, else a +# fresh one, else the last one when the fetch fails. +def week_stats(client: CursorClient, today: date, key: str, max_age_seconds: float) -> dict[str, Any]: + if max_age_seconds > 0 and carried_stats_fresh(key, today, max_age_seconds): + return carried_stats(key, today) + try: + return collect_stats(client, today, key) + except Exception: + return carried_stats(key, today) + + def summarize_limits(period: dict[str, Any]) -> list[dict[str, Any]]: """Cursor's own percentages, never a percentage derived from spend. The dollar figures do not divide into the allowance: bonus usage the providers hand out is spend without an allowance behind it, so totalSpend/limit reads well past 100% while Cursor's own meters do not. + + Every meter runs on the billing cycle. The included total is the Monthly + row; the model meters and the on-demand budget are titled after it, so the + panel attaches them to that row rather than showing four rows that each + read "Monthly". """ plan = period.get("planUsage") plan = plan if isinstance(plan, dict) else {} @@ -330,19 +416,21 @@ def summarize_limits(period: dict[str, Any]) -> list[dict[str, Any]]: resets_at = iso(period.get("billingCycleEnd")) limits = [] - for label, key in ( - ("Included total", "totalPercentUsed"), - ("Auto models", "autoPercentUsed"), - ("Named models", "apiPercentUsed"), + for label, title, key in ( + ("Included total", "Monthly", "totalPercentUsed"), + ("Auto models", "Auto models Monthly", "autoPercentUsed"), + ("Named models", "Named models Monthly", "apiPercentUsed"), ): if key in plan: - limits.append({"label": label, "percent": fraction(plan.get(key)), "resetsAt": resets_at}) + limits.append({"label": label, "title": title, "percent": fraction(plan.get(key)), "resetsAt": resets_at}) limit_cents = number(spend.get("individualLimit")) if limit_cents > 0: used_cents = number(spend.get("individualUsed")) + label = f"On-demand ${used_cents / 100:.2f} / ${limit_cents / 100:.2f}" limits.append({ - "label": f"On-demand ${used_cents / 100:.2f} / ${limit_cents / 100:.2f}", + "label": label, + "title": f"{label} Monthly", "percent": min(1.0, used_cents / limit_cents), "resetsAt": resets_at, }) @@ -379,14 +467,7 @@ def write_json(path: Path, payload: dict[str, Any]) -> None: # someone else never shows the last one's limits or tokens, and a refreshed # token keeps the same caches. An opaque API key keys by its own hash. def account_key(token: str) -> str: - subject = "" - parts = token.split(".") - if len(parts) == 3: - try: - claims = json.loads(base64.urlsafe_b64decode(parts[1] + "=" * (-len(parts[1]) % 4))) - subject = str(claims.get("sub") or "") if isinstance(claims, dict) else "" - except (ValueError, UnicodeDecodeError): - subject = "" + subject = str(jwt_claims(token).get("sub") or "") return hashlib.sha256((subject or token).encode("utf-8")).hexdigest()[:16] @@ -403,10 +484,12 @@ def open_limits(limits: Any) -> list[dict[str, Any]]: if not isinstance(limit, dict): continue try: - resets_at = datetime.fromisoformat(str(limit.get("resetsAt") or "")) + resets_at = datetime.fromisoformat(str(limit.get("resetsAt") or "").replace("Z", "+00:00")) except ValueError: resets_at = None - if resets_at is None or resets_at.tzinfo is None or resets_at > now: + if resets_at is not None and resets_at.tzinfo is None: + resets_at = resets_at.replace(tzinfo=timezone.utc) + if resets_at is None or resets_at > now: kept.append(limit) else: kept.append(dict(limit, percent=0.0, resetsAt="")) @@ -417,59 +500,85 @@ def now_ms() -> int: return round(datetime.now(timezone.utc).timestamp() * 1000) -def scan(api_base_url: str, auth_path: Path, limits_only: bool, state_db: Path | None = None) -> dict[str, Any]: +def scan(api_base_url: str, auth_path: Path, limits_only: bool, state_db: Path | None = None, force: bool = False) -> dict[str, Any]: token = read_token(auth_path, state_db) if not token: - return base_record(usageStatusText="Cursor unavailable", authHelpText=AUTH_HELP) + if cursor_in_use(auth_path): + return base_record(ready=True, usageStatusText="Waiting for auth", authHelpText=AUTH_HELP) + return base_record() client = CursorClient(token, api_base_url) key = account_key(token) cache = limits_cache(key) + cached = read_json(cache) + kept = open_limits(cached.get("limits")) + fetched_at = number(cached.get("fetchedAtMs")) + kept_tier = str(cached.get("tierLabel") or "") today = datetime.now().astimezone().date() - try: - period = client.period() - except CursorError as error: - # A failed check keeps the last good limits, marked stale so the panel - # gives their age on hover, and the week of tokens already on disk. - cached = read_json(cache) - limits = open_limits(cached.get("limits")) - status = "Waiting for auth" if error.auth else ("" if limits else "Cursor limits unavailable") + # --limits-only runs while a meter is near its limit, which can last hours, + # so the tokens it carries are only so old before they are fetched again. + scan_age = 0 if force else (LIMITS_ONLY_REUSE_SECONDS if limits_only else SCAN_REUSE_SECONDS) + + # A failed check keeps the last good limits, marked stale so the panel + # gives their age on hover, and the week of tokens already on disk. + def failed(status: str, detail: str, retry: bool = False) -> dict[str, Any]: record = base_record( ready=True, - tierLabel=str(cached.get("tierLabel") or ""), + tierLabel=kept_tier, usageStatusText=status, - authHelpText=str(error), - limits=limits, + authHelpText=detail + (" Showing the last known limits." if kept else ""), + limits=kept, limitsStale=True, - limitsFetchedAt=number(cached.get("fetchedAtMs")) if limits else 0, + limitsFetchedAt=fetched_at if kept else 0, ) + if retry: + record["retryAdvised"] = True record.update(carried_stats(key, today)) return record + if token_lapsed(token): + return failed( + "Limits paused", + "Cursor's access token has lapsed. Open the Cursor editor, or run `cursor-agent`, to renew it; signing in again is not needed.", + ) + + # A panel opened and shut repeatedly must not turn into a request per + # flick, so an answer this recent is the answer. + age = now_ms() - fetched_at + if not force and fetched_at > 0 and isinstance(cached.get("limits"), list) and 0 <= age < PROBE_MIN_INTERVAL_SECONDS * 1000: + record = base_record(ready=True, tierLabel=kept_tier, limits=kept, limitsFetchedAt=fetched_at) + record.update(week_stats(client, today, key, scan_age)) + return record + + try: + limits = summarize_limits(client.period()) + except CursorError as error: + if error.auth: + return failed("Sign-in expired", "Cursor refused the saved sign-in. " + AUTH_HELP) + return failed("" if kept else "Cursor limits unavailable", str(error), retry=error.transport) + except Exception as error: + print(f"omarchy-agent-usage-cursor: {type(error).__name__}", file=sys.stderr) + return failed("" if kept else "Cursor limits unavailable", "Cursor's usage check failed.") + + plan = client.plan_name() record = base_record( ready=True, - tierLabel=client.plan_name(), - authHelpText=AUTH_HELP, - limits=summarize_limits(period), + tierLabel=kept_tier if plan is None else plan, + limits=limits, limitsFetchedAt=now_ms(), ) - write_json(cache, {"fetchedAtMs": record["limitsFetchedAt"], "tierLabel": record["tierLabel"], "limits": record["limits"]}) - # --limits-only runs while a meter is near its limit, which can last hours, - # so the tokens it carries are only so old before they are fetched again. - if limits_only and carried_stats_fresh(key, today): - record.update(carried_stats(key, today)) - else: - try: - record.update(collect_stats(client, today, key)) - except CursorError: - record.update(carried_stats(key, today)) + # The cache only backs later runs; failing to keep it never costs this + # run's live answer. + try: + write_json(cache, {"fetchedAtMs": record["limitsFetchedAt"], "tierLabel": record["tierLabel"], "limits": limits}) + except OSError: + pass + record.update(week_stats(client, today, key, scan_age)) return record def main() -> int: parser = argparse.ArgumentParser(description="Print the Cursor usage record as JSON") - # Every call is a live API call, so there is no cache to force past; the - # flag exists 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("--auth-path", default=os.environ.get("CURSOR_AUTH_PATH", "")) @@ -478,11 +587,11 @@ def main() -> int: auth_path = Path(args.auth_path).expanduser() if args.auth_path else default_auth_path() try: - record = scan(args.api_base_url, auth_path, args.limits_only) - except CursorError as error: - record = base_record(usageStatusText="Cursor unavailable", authHelpText=str(error)) + record = scan(args.api_base_url, auth_path, args.limits_only, force=args.force) + # scan() keeps the last limits through any failed check; only a failure + # before it knows the account ends up here. except Exception as error: - record = base_record(usageStatusText="Cursor unavailable", authHelpText="Cursor usage scan failed") + record = base_record(usageStatusText="Cursor limits unavailable", authHelpText="Cursor's usage check failed.") print(f"omarchy-agent-usage-cursor: {type(error).__name__}", file=sys.stderr) print(json.dumps(record, separators=(",", ":"))) return 0 diff --git a/test/shell.d/agent-usage-cursor-test.sh b/test/shell.d/agent-usage-cursor-test.sh index acca6685c1d..29985249c3c 100755 --- a/test/shell.d/agent-usage-cursor-test.sh +++ b/test/shell.d/agent-usage-cursor-test.sh @@ -144,7 +144,7 @@ limits_cache = collector.limits_cache(key) good = json.loads(limits_cache.read_text()) limits_cache.write_text(json.dumps(dict(good, limits=[dict(good["limits"][0], resetsAt="2020-01-01T00:00:00+00:00")] + good["limits"][1:]))) collector.CursorClient = DownForOther -lapsed = collector.scan("https://api.example", Path(os.environ["XDG_CONFIG_HOME"]) / "cursor" / "auth.json", False) +lapsed = collector.scan("https://api.example", Path(os.environ["XDG_CONFIG_HOME"]) / "cursor" / "auth.json", False, force=True) limits_cache.write_text(json.dumps(good)) collector.CursorClient = StubClient @@ -161,9 +161,9 @@ class RefusedClient(StubClient): collector.CursorClient = DownClient -down = collector.scan("https://api.example", Path(os.environ["XDG_CONFIG_HOME"]) / "cursor" / "auth.json", False) +down = collector.scan("https://api.example", Path(os.environ["XDG_CONFIG_HOME"]) / "cursor" / "auth.json", False, force=True) collector.CursorClient = RefusedClient -refused = collector.scan("https://api.example", Path(os.environ["XDG_CONFIG_HOME"]) / "cursor" / "auth.json", False) +refused = collector.scan("https://api.example", Path(os.environ["XDG_CONFIG_HOME"]) / "cursor" / "auth.json", False, force=True) print(json.dumps({ "other": other, @@ -229,9 +229,9 @@ pass "Cursor collector stamps fresh limits" fail "Cursor collector keeps the last good limits, stale, after a failed check" "$result" pass "Cursor collector keeps the last good limits, stale, after a failed check" -[[ $(jq -c '{stale: .refused.limitsStale, status: .refused.usageStatusText, kept: (.refused.limits | length)}' <<<"$result") == '{"stale":true,"status":"Waiting for auth","kept":4}' ]] || - fail "Cursor collector says a refused sign-in needs attention" "$result" -pass "Cursor collector says a refused sign-in needs attention" +[[ $(jq -c '{stale: .refused.limitsStale, status: .refused.usageStatusText, kept: (.refused.limits | length)}' <<<"$result") == '{"stale":true,"status":"Sign-in expired","kept":4}' ]] || + fail "Cursor collector says a refused sign-in has expired" "$result" +pass "Cursor collector says a refused sign-in has expired" [[ $(jq -c '.other | {limits, tierLabel, todayTotalTokens, recentDays, modelUsage}' <<<"$result") == '{"limits":[],"tierLabel":"","todayTotalTokens":0,"recentDays":[],"modelUsage":{}}' ]] || fail "Cursor collector never carries another account's numbers" "$result" @@ -305,3 +305,180 @@ PY [[ $(jq -c . <<<"$ide") == '{"ide":"ide_token","cliWins":"cursor_test","untouched":true}' ]] || fail "Cursor collector falls back to the editor's sign-in, read-only" "$ide" pass "Cursor collector falls back to the editor's sign-in, read-only" + +# Cursor installed but signed in nowhere: the agent is in use, so it says it +# is waiting for a sign-in. A machine without Cursor at all stays unused. +mkdir -p "$TEST_HOME/unsigned/cursor" +unsigned=$(HOME="$TEST_HOME/unsigned-home" XDG_CONFIG_HOME="$TEST_HOME/unsigned" XDG_CACHE_HOME="$TEST_HOME/unsigned-cache" \ + CURSOR_API_KEY="" "$ROOT/bin/omarchy-agent-usage-cursor") +[[ $(jq -c '{ready, usageStatusText, limits, help: (.authHelpText != "")}' <<<"$unsigned") == '{"ready":true,"usageStatusText":"Waiting for auth","limits":[],"help":true}' ]] || + fail "Cursor collector waits for a sign-in where Cursor is installed" "$unsigned" +[[ $(jq -c '{ready, usageStatusText, authHelpText, limits, todayPrompts, recentDays, modelUsage}' <<<"$no_token") == '{"ready":false,"usageStatusText":"","authHelpText":"","limits":[],"todayPrompts":0,"recentDays":[],"modelUsage":{}}' ]] || + fail "Cursor collector prints an unused record where Cursor isn't installed" "$no_token" +pass "Cursor collector waits for a sign-in only where Cursor is installed" + +contract=$(python3 - "$ROOT/bin/omarchy-agent-usage-cursor" "$TEST_HOME/contract" <<'PY' +import base64 +import importlib.machinery +import importlib.util +import json +import os +import sqlite3 +import sys +import time +from pathlib import Path + +root = Path(sys.argv[2]) +os.environ["HOME"] = str(root / "home") +os.environ["XDG_CONFIG_HOME"] = str(root / "config") +os.environ["XDG_CACHE_HOME"] = str(root / "cache") +os.environ.pop("CURSOR_API_KEY", None) +os.environ["TZ"] = "UTC" +time.tzset() + +loader = importlib.machinery.SourceFileLoader("cursor_collector", sys.argv[1]) +spec = importlib.util.spec_from_loader(loader.name, loader) +collector = importlib.util.module_from_spec(spec) +loader.exec_module(collector) +RealClient = collector.CursorClient + + +def jwt(sub, exp): + claims = base64.urlsafe_b64encode(json.dumps({"sub": sub, "exp": exp}).encode()).decode().rstrip("=") + return f"h.{claims}.s" + + +auth = root / "config" / "cursor" / "auth.json" +auth.parent.mkdir(parents=True) +auth.write_text(json.dumps({"accessToken": jwt("u-main", time.time() + 3600)})) + +PERIOD = {"planUsage": {"totalPercentUsed": 40, "autoPercentUsed": 20, "apiPercentUsed": 60}, + "spendLimitUsage": {"individualLimit": "2000", "individualUsed": "500"}} +calls = [] +answers = {"GetPlanInfo": {"planInfo": {"planName": "Pro"}}, "GetCurrentPeriodUsage": PERIOD} + + +class Stub(RealClient): + def call(self, method, payload): + calls.append(method) + answer = answers.get(method, {"aggregations": [{"modelIntent": "default", "inputTokens": "3"}]}) + if isinstance(answer, Exception): + raise answer + return answer + + +def run(client=Stub, **kwargs): + collector.CursorClient = client + del calls[:] + record = collector.scan("https://api.example", auth, kwargs.pop("limits_only", False), **kwargs) + return record, list(calls) + + +out = {} +# Without a billing cycle end there is no reset time, as "" rather than null. +fresh, _ = run(force=True) +out["fresh"] = {key: fresh[key] for key in ("tierLabel", "usageStatusText", "authHelpText", "limitsStale")} +out["titles"] = [limit["title"] for limit in fresh["limits"]] +out["resets"] = sorted({limit["resetsAt"] for limit in fresh["limits"]}) + +# A second run straight after asks nothing: the limits answer and the token +# scan are both this recent. +reused, asked = run() +out["reused"] = {"asked": asked, "same": reused["limits"] == fresh["limits"], "stale": reused["limitsStale"], + "help": reused["authHelpText"], "tokens": reused["todayTotalTokens"]} +_, asked = run(limits_only=True) +out["limitsOnlyReused"] = asked +_, asked = run(force=True) +out["forced"] = sorted(set(asked)) + +# GetPlanInfo failing leaves the plan the account last had. +answers["GetPlanInfo"] = collector.CursorError("Cursor's dashboard returned status 500.") +planless, _ = run(force=True) +out["planKept"] = planless["tierLabel"] +answers["GetPlanInfo"] = {"planInfo": {"planName": "Pro"}} + +# A failure that isn't Cursor's error keeps the limits and the tokens too. +answers["GetCurrentPeriodUsage"] = RuntimeError("boom") +broken, _ = run(force=True) +out["broken"] = {"kept": len(broken["limits"]), "stale": broken["limitsStale"], "status": broken["usageStatusText"], + "tokens": broken["todayTotalTokens"], "retry": broken.get("retryAdvised", False)} +answers["GetCurrentPeriodUsage"] = PERIOD + + +# retryAdvised only when nothing answered at all. +def urlopen_raising(error): + def urlopen(request, timeout=None): + raise error + return urlopen + + +transport = {} +for name, error in ( + ("dns", collector.urllib.error.URLError("no route")), + ("timeout", TimeoutError()), + ("status", collector.urllib.error.HTTPError("https://api.example", 500, "", {}, None)), +): + collector.urllib.request.urlopen = urlopen_raising(error) + record, _ = run(RealClient, force=True) + transport[name] = [record.get("retryAdvised", False), record["limitsStale"], len(record["limits"]), record["usageStatusText"]] +out["transport"] = transport + +# A lapsed sign-in isn't sent: the CLI or editor renews it when it runs. +collector.urllib.request.urlopen = urlopen_raising(AssertionError("a lapsed token was sent")) +auth.write_text(json.dumps({"accessToken": jwt("u-main", time.time() - 60)})) +paused, _ = run(RealClient, force=True) +out["paused"] = [paused["usageStatusText"], paused["limitsStale"], len(paused["limits"])] + +# ...but a live editor sign-in stands in for a lapsed CLI one. +db = root / "config" / "Cursor" / "User" / "globalStorage" / "state.vscdb" +db.parent.mkdir(parents=True) +conn = sqlite3.connect(db) +conn.execute("CREATE TABLE ItemTable (key TEXT UNIQUE ON CONFLICT REPLACE, value BLOB)") +conn.execute("INSERT INTO ItemTable VALUES ('cursorAuth/accessToken', ?)", (jwt("u-editor", time.time() + 3600),)) +conn.commit() +conn.close() +out["fallthrough"] = collector.account_key(collector.read_token(auth)) == collector.account_key(jwt("u-editor", 0)) + +# Untrusted numbers never break the record. +out["robust"] = [collector.fraction("nan"), collector.fraction("inf"), collector.fraction(True), collector.number("1e400"), collector.iso("9" * 30)] +print(json.dumps(out)) +PY +) + +[[ $(jq -c '.fresh' <<<"$contract") == '{"tierLabel":"Pro","usageStatusText":"","authHelpText":"","limitsStale":false}' ]] || + fail "Cursor collector clears the help text when healthy" "$contract" +pass "Cursor collector clears the help text when healthy" + +[[ $(jq -c '.titles' <<<"$contract") == '["Monthly","Auto models Monthly","Named models Monthly","On-demand $5.00 / $20.00 Monthly"]' ]] || + fail "Cursor collector titles its meters after the billing period" "$contract" +pass "Cursor collector titles its meters after the billing period" + +[[ $(jq -c '.resets' <<<"$contract") == '[""]' ]] || + fail "Cursor collector gives an unknown reset as an empty string" "$contract" +pass "Cursor collector gives an unknown reset as an empty string" + +[[ $(jq -c '.reused' <<<"$contract") == '{"asked":[],"same":true,"stale":false,"help":"","tokens":3}' && $(jq -c '.limitsOnlyReused' <<<"$contract") == '[]' ]] || + fail "Cursor collector reuses a recent limits answer and token scan" "$contract" +[[ $(jq -c '.forced' <<<"$contract") == '["GetAggregatedUsageEvents","GetCurrentPeriodUsage","GetPlanInfo"]' ]] || + fail "Cursor collector asks again on --force" "$contract" +pass "Cursor collector reuses a recent probe, and --force asks again" + +[[ $(jq -r '.planKept' <<<"$contract") == "Pro" ]] || + fail "Cursor collector keeps the plan when GetPlanInfo fails" "$contract" +pass "Cursor collector keeps the plan when GetPlanInfo fails" + +[[ $(jq -c '.broken' <<<"$contract") == '{"kept":4,"stale":true,"status":"","tokens":3,"retry":false}' ]] || + fail "Cursor collector keeps limits and tokens through an unexpected failure" "$contract" +pass "Cursor collector keeps limits and tokens through an unexpected failure" + +[[ $(jq -c '.transport' <<<"$contract") == '{"dns":[true,true,4,""],"timeout":[true,true,4,""],"status":[false,true,4,""]}' ]] || + fail "Cursor collector advises a retry only when the dashboard can't be reached" "$contract" +pass "Cursor collector advises a retry only when the dashboard can't be reached" + +[[ $(jq -c '.paused' <<<"$contract") == '["Limits paused",true,4]' && $(jq -r '.fallthrough' <<<"$contract") == "true" ]] || + fail "Cursor collector pauses on a lapsed sign-in, preferring a live one" "$contract" +pass "Cursor collector pauses on a lapsed sign-in, preferring a live one" + +[[ $(jq -c '.robust' <<<"$contract") == '[0.0,0.0,0.0,0,""]' ]] || + fail "Cursor collector shrugs off untrusted numbers" "$contract" +pass "Cursor collector shrugs off untrusted numbers" From b6e8753e923de6cd1606b89df30c9d6666dc1eb0 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:04:28 -0400 Subject: [PATCH 051/546] Bring the Grok Bot collector in line with the usage collector contract It now finds the sign-in exactly as the Cursor collector does, reuses an answer for 15 seconds unless forced, advises a retry only when the dashboard can't be reached, and says "Sign-in expired" or "Limits paused" for a refused or lapsed sign-in. A kept week that ended comes back at zero instead of being dropped, one without a reset time is kept as it is, and every record carries the full empty stats block. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-grokbot | 212 +++++++++++++++++------ test/shell.d/agent-usage-grokbot-test.sh | 121 ++++++++++--- 2 files changed, 255 insertions(+), 78 deletions(-) diff --git a/bin/omarchy-agent-usage-grokbot b/bin/omarchy-agent-usage-grokbot index e263d3726b0..4e7ff530470 100755 --- a/bin/omarchy-agent-usage-grokbot +++ b/bin/omarchy-agent-usage-grokbot @@ -8,10 +8,10 @@ Grok Bot runs on a Cursor account, but its weekly allowance is metered apart from Cursor's own plan: the dashboard's GetSandUsageStatus RPC (Sand is Grok Bot's internal name) says how much of this week's included usage is gone and when it resets. It is asked as the same account the Cursor collector reads, -so the two never describe different people. Grok Bot keeps -nothing countable on disk, so the record carries the limit alone, and an -account without a Grok Bot allowance gives an empty record, which the panel -skips. The agents panel only ever reads the JSON this prints. +found the same way, so the two never describe different people. Grok Bot +keeps nothing countable on disk, so the record carries the limit alone, and +an account without a Grok Bot allowance gives an unused record, which the +panel skips. The agents panel only ever reads the JSON this prints. """ import argparse @@ -19,6 +19,7 @@ import base64 import hashlib import http.client import json +import math import os import sqlite3 import tempfile @@ -31,6 +32,9 @@ AGENT_ID = "grokbot" AGENT_NAME = "Grok Bot" AUTH_HELP = "Sign in to the Cursor editor, or run `cursor-agent login`, with your Grok Bot account." USAGE_URL = "https://api2.cursor.sh/aiserver.v1.DashboardService/GetSandUsageStatus" +# Repeated panel opens reuse an answer this recent instead of asking again; +# --force (a person pressing refresh) always asks. +PROBE_MIN_INTERVAL_SECONDS = 15 def config_home(): @@ -41,6 +45,20 @@ def cache_root(): return Path(os.environ.get("XDG_CACHE_HOME") or (Path.home() / ".cache")) / "omarchy" / "agent-usage" +def now_ms(): + return round(datetime.now(timezone.utc).timestamp() * 1000) + + +def finite(value): + if isinstance(value, bool): + return 0.0 + try: + parsed = float(value or 0) + except (TypeError, ValueError): + return 0.0 + return parsed if math.isfinite(parsed) else 0.0 + + def read_json(path): try: data = json.loads(Path(path).read_text(encoding="utf-8")) @@ -81,6 +99,7 @@ def read_ide_token(state_db): except sqlite3.Error: return "" try: + conn.execute("PRAGMA query_only = ON") row = conn.execute("SELECT value FROM ItemTable WHERE key = 'cursorAuth/accessToken' LIMIT 1").fetchone() except sqlite3.Error: return "" @@ -92,25 +111,51 @@ def read_ide_token(state_db): return str(value or "").strip() -# The same account the Cursor collector reads, in the same order: the key the -# user exported, the credential the Cursor CLI stores, or the editor's sign-in. +def read_cli_token(path): + data = read_json(path) + return str(data.get("accessToken") or "").strip() + + +def jwt_claims(token): + parts = token.split(".") + if len(parts) != 3: + return {} + try: + claims = json.loads(base64.urlsafe_b64decode(parts[1] + "=" * (-len(parts[1]) % 4))) + except (ValueError, UnicodeDecodeError): + return {} + return claims if isinstance(claims, dict) else {} + + +# A sign-in token whose expiry has passed. The CLI and the editor each renew +# their own the next time they run; an opaque API key never lapses here. +def token_lapsed(token): + expires = finite(jwt_claims(token).get("exp")) + return 0 < expires <= datetime.now(timezone.utc).timestamp() + + +# The same account the Cursor collector reads, found the same way: the key +# the user exported, the credential the Cursor CLI stores, or the editor's +# sign-in. A lapsed sign-in gives way to the next one that hasn't lapsed; +# with all of them lapsed, the first one decides. def read_token(): - token = str(os.environ.get("CURSOR_API_KEY") or "").strip() - token = token or str(read_json(config_home() / "cursor" / "auth.json").get("accessToken") or "").strip() - return token or read_ide_token(config_home() / "Cursor" / "User" / "globalStorage" / "state.vscdb") + key = str(os.environ.get("CURSOR_API_KEY") or "").strip() + if key: + return key + cli = read_cli_token(config_home() / "cursor" / "auth.json") + if cli and not token_lapsed(cli): + return cli + ide = read_ide_token(config_home() / "Cursor" / "User" / "globalStorage" / "state.vscdb") + if ide and not token_lapsed(ide): + return ide + return cli or ide # Who the token signs in as (the JWT's subject), so a machine signed in to # someone else never shows the last one's limit, and a refreshed token keeps -# the same cache. +# the same cache. An opaque API key keys by its own hash. def account_key(token): - parts = token.split(".") - subject = "" - if len(parts) == 3: - try: - subject = str(json.loads(base64.urlsafe_b64decode(parts[1] + "=" * (-len(parts[1]) % 4))).get("sub") or "") - except (ValueError, AttributeError): - subject = "" + subject = str(jwt_claims(token).get("sub") or "") return hashlib.sha256((subject or token).encode("utf-8")).hexdigest()[:16] @@ -127,7 +172,10 @@ def fetch_status(token): method="POST", ) with urllib.request.urlopen(request, timeout=15) as response: - return json.loads(response.read() or b"{}") + status = json.loads(response.read() or b"{}") + if not isinstance(status, dict): + raise ValueError("not an object") + return status # The answer is protobuf JSON, which leaves out a field holding zero or @@ -139,13 +187,27 @@ def weekly_limits(status): resets = parse_time(status.get("nextResetTimestampUtc")) return [{ "label": "Weekly", - "percent": max(0.0, min(1.0, float(status.get("usagePercent") or 0) / 100)), + "percent": max(0.0, min(1.0, finite(status.get("usagePercent")) / 100)), "resetsAt": resets.isoformat() if resets else "", }] -def collect(): - record = { +# A kept week that has ended is a fresh allowance, so it comes back at zero +# with no reset time until a check says when the next one ends. One with no +# known end is kept as it is. +def open_limits(limits): + now = datetime.now(timezone.utc) + kept = [] + for limit in limits if isinstance(limits, list) else []: + if not isinstance(limit, dict): + continue + resets = parse_time(limit.get("resetsAt")) if limit.get("resetsAt") else None + kept.append(limit if resets is None or resets > now else dict(limit, percent=0.0, resetsAt="")) + return kept + + +def empty_record(): + return { "schemaVersion": 1, "id": AGENT_ID, "name": AGENT_NAME, @@ -155,63 +217,103 @@ def collect(): "hasPromptStats": False, "tierLabel": "", "usageStatusText": "", - "authHelpText": AUTH_HELP, + "authHelpText": "", "limits": [], "limitsStale": False, "limitsFetchedAt": 0, + "todayPrompts": 0, + "todaySessions": 0, + "todayTotalTokens": 0, + "todayTokensByModel": {}, + "recentDays": [], + "totalPrompts": 0, + "totalSessions": 0, + "activeDays": 0, + "activeDates": [], + "modelUsage": {}, } + + +def collect(force=False): + record = empty_record() + # Without a sign-in there is no telling whether this account has Grok Bot + # at all, so the machine reads as not using it; the Cursor collector is the + # one that asks for a sign-in. token = read_token() if not token: return record cache = cache_root() / f"grokbot-limits-{account_key(token)}.json" - try: - status = fetch_status(token) - limits = weekly_limits(status) - except urllib.error.HTTPError as error: - auth = error.code in (401, 403) - # Unreachable, timed out, or reset mid-read are all OSError; a body cut - # short is an HTTPException. - except (OSError, http.client.HTTPException, ValueError, TypeError, AttributeError): - auth = False + cached = read_json(cache) + kept = open_limits(cached.get("limits")) + fetched_at = max(0, round(finite(cached.get("fetchedAtMs")))) + kept_tier = str(cached.get("tierLabel") or "") + + # A panel opened and shut repeatedly must not turn into a request per + # flick, so an answer this recent is the answer. + age = now_ms() - fetched_at + if not force and fetched_at > 0 and isinstance(cached.get("limits"), list) and 0 <= age < PROBE_MIN_INTERVAL_SECONDS * 1000: + record.update(ready=bool(kept), tierLabel=kept_tier, limits=kept, limitsFetchedAt=fetched_at) + return record + + retry = False + if token_lapsed(token): + status = "Limits paused" + detail = "Cursor's access token has lapsed. Open the Cursor editor, or run `cursor-agent`, to renew it; signing in again is not needed." else: - fetched_at = round(datetime.now(timezone.utc).timestamp() * 1000) - tier = str(status.get("grokPlanLabel") or "") - record.update(ready=bool(limits), tierLabel=tier, authHelpText="", limits=limits, limitsFetchedAt=fetched_at) - # The cache only backs a later failed check; failing to keep it never - # costs this run's live answer. try: - write_json(cache, {"fetchedAtMs": fetched_at, "tierLabel": tier, "limits": limits}) + answer = fetch_status(token) + limits = weekly_limits(answer) + except urllib.error.HTTPError as error: + if error.code in (401, 403): + status, detail = "Sign-in expired", "Cursor refused the saved sign-in. " + AUTH_HELP + else: + status, detail = "", f"Cursor's dashboard returned status {error.code}." + # A reply cut short or one that isn't the expected JSON came from a + # server that answered. + except (http.client.HTTPException, ValueError, TypeError, AttributeError): + status, detail = "", "Cursor's dashboard sent a reply that couldn't be read." + # No route, no DNS, refused, reset or timed out: nothing answered, so the + # panel tries again soon. except OSError: - pass - return record + status, detail, retry = "", "Couldn't reach Cursor's dashboard. Retrying shortly.", True + else: + tier = str(answer.get("grokPlanLabel") or "") + fetched_at = now_ms() + record.update(ready=bool(limits), tierLabel=tier, limits=limits, limitsFetchedAt=fetched_at) + # The cache only backs later runs; failing to keep it never costs this + # run's live answer. + try: + write_json(cache, {"fetchedAtMs": fetched_at, "tierLabel": tier, "limits": limits}) + except OSError: + pass + return record - # A failed check keeps the last good limit until its week ends, marked - # stale so the panel dims it; one with no known end isn't kept at all. - # With nothing kept, the record stays empty, so a Cursor account that - # never had Grok Bot doesn't appear on a blip. - cached = read_json(cache) - now = datetime.now(timezone.utc) - limits = [limit for limit in cached.get("limits") or [] if (resets := parse_time(limit.get("resetsAt"))) and resets > now] + # A failed check keeps the last good limit, marked stale so the panel + # gives its age on hover. With nothing kept the record stays unused, so a + # Cursor account that never had Grok Bot doesn't appear on a blip. record.update( - ready=bool(limits), - tierLabel=str(cached.get("tierLabel") or ""), - usageStatusText="Waiting for auth" if auth else ("" if limits else "Grok Bot limits unavailable"), - limits=limits, + ready=bool(kept), + tierLabel=kept_tier if kept else "", + usageStatusText=status or ("" if kept else "Grok Bot limits unavailable"), + authHelpText=detail + (" Showing the last known limit." if kept else ""), + limits=kept, limitsStale=True, - limitsFetchedAt=int(cached.get("fetchedAtMs") or 0) if limits else 0, + limitsFetchedAt=fetched_at if kept else 0, ) + if retry: + record["retryAdvised"] = True return record def main(): parser = argparse.ArgumentParser(description="Print the Grok Bot usage record as JSON") - # Every run is one live call, so there is no cache to force past and no - # scan to skip; the flags exist so every collector takes the same invocation. parser.add_argument("--force", action="store_true") + # Grok Bot has no local scan to reuse, so a near-limit refresh is a normal + # run; the flag keeps the invocation every collector takes. parser.add_argument("--limits-only", action="store_true") - parser.parse_args() - print(json.dumps(collect(), separators=(",", ":"))) + args = parser.parse_args() + print(json.dumps(collect(force=args.force), separators=(",", ":"))) if __name__ == "__main__": diff --git a/test/shell.d/agent-usage-grokbot-test.sh b/test/shell.d/agent-usage-grokbot-test.sh index 38798eab184..a0d830c0df9 100755 --- a/test/shell.d/agent-usage-grokbot-test.sh +++ b/test/shell.d/agent-usage-grokbot-test.sh @@ -17,19 +17,23 @@ unset CURSOR_API_KEY reset_at=$(python3 -c 'import datetime as dt; print((dt.datetime.now(dt.timezone.utc) + dt.timedelta(days=3)).strftime("%Y-%m-%dT%H:%M:%S.%fZ"))') -# A JWT whose subject names the account, as Cursor's sign-in tokens do. +# A JWT whose subject names the account, as Cursor's sign-in tokens do, with +# an expiry an hour out unless one is given. jwt() { - printf 'h.%s.s' "$(printf '{"sub":"%s"}' "$1" | base64 -w0 | tr '+/' '-_' | tr -d '=')" + local exp=${2:-$(( $(date +%s) + 3600 ))} + printf 'h.%s.s' "$(printf '{"sub":"%s","exp":%s}' "$1" "$exp" | base64 -w0 | tr '+/' '-_' | tr -d '=')" } cli_signed_in() { mkdir -p "$XDG_CONFIG_HOME/cursor" - jq -n --arg token "$(jwt "$1")" '{accessToken: $token}' >"$XDG_CONFIG_HOME/cursor/auth.json" + jq -n --arg token "$(jwt "$@")" '{accessToken: $token}' >"$XDG_CONFIG_HOME/cursor/auth.json" } -# ANSWER is the dashboard's JSON reply, an HTTP status code to fail with, or -# cut/reset for a connection that breaks mid-reply. With EXPECT set, any other -# token is refused, as Cursor would refuse it. +# ANSWER is the dashboard's JSON reply, an HTTP status code to fail with, +# cut/reset for a connection that breaks mid-reply, unreachable for one that +# never connects, or unasked to fail the test if anything is sent. With EXPECT +# set, any other token is refused, as Cursor would refuse it. Every run is +# --force unless ARGS says otherwise, so each one really asks. collect() { COLLECTOR="$ROOT/bin/omarchy-agent-usage-grokbot" python3 - <<'PY' import http.client, importlib.machinery, importlib.util, io, os, sys @@ -42,6 +46,8 @@ loader.exec_module(collector) def urlopen(request, timeout=None): assert request.full_url == collector.USAGE_URL answer = os.environ["ANSWER"] + if answer == "unasked": + raise AssertionError("the dashboard should not have been asked") expected = os.environ.get("EXPECT") if expected and request.get_header("Authorization") != f"Bearer {expected}": answer = "401" @@ -49,37 +55,53 @@ def urlopen(request, timeout=None): raise http.client.IncompleteRead(b"{") if answer == "reset": raise ConnectionResetError() + if answer == "unreachable": + raise collector.urllib.error.URLError("no route to host") if answer.isdigit(): raise collector.urllib.error.HTTPError(request.full_url, int(answer), "", {}, None) return io.BytesIO(answer.encode()) collector.urllib.request.urlopen = urlopen -sys.argv = ["omarchy-agent-usage-grokbot", "--limits-only"] +sys.argv = ["omarchy-agent-usage-grokbot", *os.environ.get("ARGS", "--limits-only --force").split()] collector.main() PY } record=$(ANSWER='{}' collect) -[[ $(jq -c '{id, ready, limits}' <<<"$record") == '{"id":"grokbot","ready":false,"limits":[]}' ]] || - fail "a machine signed in to Cursor nowhere gives an empty record" "$record" -pass "a machine signed in to Cursor nowhere gives an empty record" +[[ $(jq -c '{id, ready, limits, usageStatusText, authHelpText, hasLocalStats, hasPromptStats, todayPrompts, todayTotalTokens, recentDays, modelUsage, activeDates}' <<<"$record") == '{"id":"grokbot","ready":false,"limits":[],"usageStatusText":"","authHelpText":"","hasLocalStats":false,"hasPromptStats":false,"todayPrompts":0,"todayTotalTokens":0,"recentDays":[],"modelUsage":{},"activeDates":[]}' ]] || + fail "a machine signed in to Cursor nowhere gives an unused record" "$record" +pass "a machine signed in to Cursor nowhere gives an unused record" cli_signed_in u-main answer=$(jq -nc --arg at "$reset_at" '{usagePercent: 42.5, nextResetTimestampUtc: $at, hasNonZeroIncludedLimit: true, grokPlanLabel: "Pro"}') record=$(ANSWER="$answer" collect) -[[ $(jq -c '{ready, tierLabel, stale: .limitsStale, label: .limits[0].label, percent: .limits[0].percent, help: .authHelpText}' <<<"$record") == '{"ready":true,"tierLabel":"Pro","stale":false,"label":"Weekly","percent":0.425,"help":""}' ]] || +[[ $(jq -c '{ready, tierLabel, stale: .limitsStale, label: .limits[0].label, percent: .limits[0].percent, help: .authHelpText, status: .usageStatusText, retry: (.retryAdvised // false)}' <<<"$record") == '{"ready":true,"tierLabel":"Pro","stale":false,"label":"Weekly","percent":0.425,"help":"","status":"","retry":false}' ]] || fail "Grok Bot's weekly allowance comes from the Cursor sign-in" "$record" python3 -c 'import datetime as dt, sys; d = dt.datetime.fromisoformat(sys.argv[1]) - dt.datetime.now(dt.timezone.utc); sys.exit(0 if dt.timedelta(days=2) < d < dt.timedelta(days=3) else 1)' "$(jq -r '.limits[0].resetsAt' <<<"$record")" || fail "the weekly window says when it resets" "$record" pass "Grok Bot's weekly allowance comes from the Cursor sign-in" -# Failed checks keep the last good limit, dimmed, and say why when it's auth. +# An answer this recent is reused rather than asked again, unless forced. +record=$(ARGS="" ANSWER=unasked collect) +[[ $(jq -c '{ready, tierLabel, stale: .limitsStale, percent: .limits[0].percent, help: .authHelpText}' <<<"$record") == '{"ready":true,"tierLabel":"Pro","stale":false,"percent":0.425,"help":""}' ]] || + fail "a recent answer is reused" "$record" +record=$(ARGS="--limits-only" ANSWER=unasked collect) +[[ $(jq -r '.limits[0].percent' <<<"$record") == "0.425" ]] || + fail "a near-limit refresh reuses a recent answer too" "$record" +answer_later=$(jq -nc --arg at "$reset_at" '{usagePercent: 50, nextResetTimestampUtc: $at, hasNonZeroIncludedLimit: true, grokPlanLabel: "Pro"}') +record=$(ARGS="--force" ANSWER="$answer_later" collect) +[[ $(jq -r '.limits[0].percent' <<<"$record") == "0.5" ]] || + fail "--force asks again" "$record" +record=$(ANSWER="$answer" collect) +pass "a recent answer is reused, and --force asks again" + +# Failed checks keep the last good limit, stale, and say why when it's auth. record=$(ANSWER=500 collect) -[[ $(jq -c '{ready, tierLabel, stale: .limitsStale, percent: .limits[0].percent, status: .usageStatusText}' <<<"$record") == '{"ready":true,"tierLabel":"Pro","stale":true,"percent":0.425,"status":""}' ]] || +[[ $(jq -c '{ready, tierLabel, stale: .limitsStale, percent: .limits[0].percent, status: .usageStatusText, retry: (.retryAdvised // false), help: (.authHelpText | test("500"))}' <<<"$record") == '{"ready":true,"tierLabel":"Pro","stale":true,"percent":0.425,"status":"","retry":false,"help":true}' ]] || fail "a failed check keeps the last limit, marked stale" "$record" record=$(ANSWER=401 collect) -[[ $(jq -c '{stale: .limitsStale, percent: .limits[0].percent, status: .usageStatusText}' <<<"$record") == '{"stale":true,"percent":0.425,"status":"Waiting for auth"}' ]] || - fail "a refused sign-in keeps the last limit and says so" "$record" +[[ $(jq -c '{stale: .limitsStale, percent: .limits[0].percent, status: .usageStatusText}' <<<"$record") == '{"stale":true,"percent":0.425,"status":"Sign-in expired"}' ]] || + fail "a refused sign-in keeps the last limit and says it expired" "$record" for broken in cut reset; do record=$(ANSWER=$broken collect) [[ $(jq -c '{stale: .limitsStale, percent: .limits[0].percent}' <<<"$record") == '{"stale":true,"percent":0.425}' ]] || @@ -87,6 +109,25 @@ for broken in cut reset; do done pass "failed checks keep the last limit, marked stale" +# Only a dashboard that couldn't be reached at all asks for a sooner retry. +record=$(ANSWER=unreachable collect) +[[ $(jq -c '{retry: .retryAdvised, stale: .limitsStale, percent: .limits[0].percent}' <<<"$record") == '{"retry":true,"stale":true,"percent":0.425}' ]] || + fail "an unreachable dashboard advises a retry" "$record" +record=$(ANSWER=cut collect) +[[ $(jq -r '.retryAdvised // false' <<<"$record") == "false" ]] || + fail "a reply cut short is an answer, not a reason to retry" "$record" +pass "only an unreachable dashboard advises a retry" + +# A kept week that has ended comes back at zero, never dropped. +cache_file=$(ls "$XDG_CACHE_HOME"/omarchy/agent-usage/grokbot-limits-*.json) +kept=$(<"$cache_file") +jq '.limits[0].resetsAt = "2020-01-01T00:00:00+00:00"' <<<"$kept" >"$cache_file" +record=$(ANSWER=500 collect) +[[ $(jq -c '{ready, stale: .limitsStale, limits: [.limits[] | {label, percent, resetsAt}]}' <<<"$record") == '{"ready":true,"stale":true,"limits":[{"label":"Weekly","percent":0.0,"resetsAt":""}]}' ]] || + fail "an elapsed kept week comes back at zero" "$record" +printf '%s' "$kept" >"$cache_file" +pass "an elapsed kept week comes back at zero" + # A cache that can't be written never costs the live answer. chmod a-w "$XDG_CACHE_HOME/omarchy/agent-usage" record=$(ANSWER="$answer" collect) @@ -95,10 +136,18 @@ chmod u+w "$XDG_CACHE_HOME/omarchy/agent-usage" fail "an unwritable cache still gives the live limit" "$record" pass "an unwritable cache still gives the live limit" +# A lapsed sign-in isn't sent: the CLI or editor renews it the next time it +# runs, so the last limit is kept with the limits paused. +cli_signed_in u-main "$(( $(date +%s) - 60 ))" +record=$(ANSWER=unasked collect) +[[ $(jq -c '{ready, stale: .limitsStale, percent: .limits[0].percent, status: .usageStatusText}' <<<"$record") == '{"ready":true,"stale":true,"percent":0.425,"status":"Limits paused"}' ]] || + fail "a lapsed sign-in pauses the limit without asking" "$record" +pass "a lapsed sign-in pauses the limit without asking" + # Another account's limit never stands in for this one's. cli_signed_in u-other record=$(ANSWER=500 collect) -[[ $(jq -c '{ready, limits}' <<<"$record") == '{"ready":false,"limits":[]}' ]] || +[[ $(jq -c '{ready, limits, tierLabel}' <<<"$record") == '{"ready":false,"limits":[],"tierLabel":""}' ]] || fail "another account's cached limit is never shown" "$record" pass "another account's cached limit is never shown" @@ -110,23 +159,25 @@ record=$(ANSWER="$answer" collect) fail "an untouched week reads as 0%" "$record" record=$(ANSWER='{"usagePercent": 10}' collect) [[ $(jq -c '{ready, limits}' <<<"$record") == '{"ready":false,"limits":[]}' ]] || - fail "a Cursor account without Grok Bot gives an empty record" "$record" + fail "a Cursor account without Grok Bot gives an unused record" "$record" +record=$(ANSWER='{"usagePercent": "NaN", "hasNonZeroIncludedLimit": true}' collect) +[[ $(jq -r '.limits[0].percent' <<<"$record") == "0.0" ]] || + fail "an unreadable percentage reads as 0%" "$record" pass "an untouched week reads 0%, and no allowance means no Grok Bot" -# A limit with no known end shows while live, but is never kept through a -# failed check, where it could stand for a week long gone. +# A limit with no known end is kept as it is through a failed check. cli_signed_in u-no-reset record=$(ANSWER='{"usagePercent": 5, "hasNonZeroIncludedLimit": true}' collect) [[ $(jq -c '{ready, percent: .limits[0].percent, resetsAt: .limits[0].resetsAt}' <<<"$record") == '{"ready":true,"percent":0.05,"resetsAt":""}' ]] || fail "a limit with no known reset still shows while live" "$record" record=$(ANSWER=500 collect) -[[ $(jq -c '{ready, limits}' <<<"$record") == '{"ready":false,"limits":[]}' ]] || - fail "a limit with no known reset is not kept through a failed check" "$record" -pass "a limit with no known reset is never kept through a failed check" +[[ $(jq -c '{ready, stale: .limitsStale, percent: .limits[0].percent, resetsAt: .limits[0].resetsAt}' <<<"$record") == '{"ready":true,"stale":true,"percent":0.05,"resetsAt":""}' ]] || + fail "a limit with no known reset is kept through a failed check" "$record" +pass "a limit with no known reset is kept as it is through a failed check" # The same account the Cursor collector reads: an exported key, then the CLI's # sign-in, then the editor's, with all three present and each asked with its -# own token. +# own token. A lapsed CLI sign-in gives way to a live editor one. cli_signed_in p-cli state_db="$XDG_CONFIG_HOME/Cursor/User/globalStorage/state.vscdb" mkdir -p "$(dirname "$state_db")" @@ -145,8 +196,32 @@ record=$(CURSOR_API_KEY="$(jwt p-key)" EXPECT="$(jwt p-key)" ANSWER="$answer" co record=$(EXPECT="$(jwt p-cli)" ANSWER="$answer" collect) [[ $(jq -c '{ready, stale: .limitsStale, percent: .limits[0].percent}' <<<"$record") == "$live" ]] || fail "the Cursor CLI's sign-in wins over the editor's" "$record" +cli_signed_in p-cli "$(( $(date +%s) - 60 ))" +record=$(EXPECT="$(jwt p-editor)" ANSWER="$answer" collect) +[[ $(jq -c '{ready, stale: .limitsStale, percent: .limits[0].percent}' <<<"$record") == "$live" ]] || + fail "a live editor sign-in stands in for a lapsed CLI one" "$record" rm -f "$XDG_CONFIG_HOME/cursor/auth.json" record=$(EXPECT="$(jwt p-editor)" ANSWER="$answer" collect) [[ $(jq -c '{ready, stale: .limitsStale, percent: .limits[0].percent}' <<<"$record") == "$live" ]] || fail "the Cursor editor's sign-in stands in for the CLI's" "$record" pass "the sign-in is the one the Cursor collector reads, in the same order" + +# Both collectors pick the same sign-in from the same files. +same=$(COLLECTORS="$ROOT/bin" python3 - <<'PY' +import importlib.machinery, importlib.util, os +from pathlib import Path + +def load(name): + loader = importlib.machinery.SourceFileLoader(name.replace("-", "_"), os.path.join(os.environ["COLLECTORS"], name)) + spec = importlib.util.spec_from_loader(loader.name, loader) + module = importlib.util.module_from_spec(spec) + loader.exec_module(module) + return module + +cursor = load("omarchy-agent-usage-cursor") +grokbot = load("omarchy-agent-usage-grokbot") +print(cursor.read_token(cursor.default_auth_path()) == grokbot.read_token() != "") +PY +) +[[ $same == "True" ]] || fail "Grok Bot and Cursor read the same sign-in" "$same" +pass "Grok Bot and Cursor read the same sign-in" From 320284355eb687c2fbdc17be24fe011c05400821 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:09:02 -0400 Subject: [PATCH 052/546] Bring the Copilot collector in line with the usage collector contract A quota answer is reused for 15 seconds unless forced, and gh is only run when GitHub is actually asked: its account is keyed by the user hosts.yml names. The cache is now copilot-limits-*, adopting an answer under the old copilot-quota-* name once, and is read whatever its file time. A kept window that reset has no reset time instead of a guessed one, an unreachable endpoint advises a retry, and the help text says what failed instead of the sign-in hint. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-copilot | 224 +++++++++++++++++------ test/shell.d/agent-usage-copilot-test.sh | 158 +++++++++++++++- 2 files changed, 329 insertions(+), 53 deletions(-) diff --git a/bin/omarchy-agent-usage-copilot b/bin/omarchy-agent-usage-copilot index 81ad101e374..7bb02f07c80 100755 --- a/bin/omarchy-agent-usage-copilot +++ b/bin/omarchy-agent-usage-copilot @@ -19,9 +19,11 @@ The monthly meter is an allowance of AI credits: GitHub's own account-wide figure from the internal quota endpoint (the same one the editor plugins use), asked with a token from COPILOT_QUOTA_TOKEN, GH_TOKEN, GITHUB_TOKEN, or `gh auth token`. It names the plan, the allowance, and when it resets, and a -failed check keeps the last answer, marked stale. Only when it can't be asked -at all does a monthlyCredits budget from ~/.config/omarchy/agents/copilot.json -give a local estimate, labeled as one, of this machine's CLI spend: +failed check keeps the last answer, marked stale. An answer under 15 seconds +old is reused rather than asked again (--force always asks), and `gh` is only +run when GitHub is actually asked. Only when it can't be asked at all does a +monthlyCredits budget from ~/.config/omarchy/agents/copilot.json give a local +estimate, labeled as one, of this machine's CLI spend: { "monthlyCredits": 1500, @@ -41,6 +43,7 @@ import calendar import datetime as dt import fcntl import hashlib +import http.client import json import os import shutil @@ -57,7 +60,7 @@ from typing import Any AGENT_ID = "copilot" AGENT_NAME = "GitHub Copilot" AUTH_HELP = ( - "Run `copilot` in a terminal and sign in to GitHub to start recording sessions." + "Run `gh auth login`, or set COPILOT_QUOTA_TOKEN, to show Copilot's monthly allowance." ) ULTRA_SECONDS = 1_000_000_000 @@ -81,6 +84,9 @@ MAX_DB_BYTES = 512 * 1024 * 1024 SCAN_REUSE_SECONDS = 20 LIMITS_ONLY_REUSE_SECONDS = 900 +# Repeated panel opens reuse a quota answer this recent instead of asking +# GitHub again; --force (a person pressing refresh) always asks. +PROBE_MIN_INTERVAL_SECONDS = 15 def number(value: Any) -> int: @@ -717,17 +723,19 @@ def direct_scan(home_dir: Path) -> tuple[dict[str, Any], bool]: # ------------------------------------------------------------ remote quota -def github_token() -> str: - """Find a GitHub token without ever storing or logging one. +def short_hash(value: str) -> str: + return hashlib.sha256(value.encode("utf-8")).hexdigest()[:16] - Environment first, so a user can scope exactly what this sees, then the - `gh` CLI's own credential store, which is where a machine that has signed - in already keeps it. - """ + +def env_token() -> str: for var in ("COPILOT_QUOTA_TOKEN", "GH_TOKEN", "GITHUB_TOKEN"): value = os.environ.get(var) if value and value.strip(): return value.strip() + return "" + + +def gh_token() -> str: exe = shutil.which("gh") if not exe: return "" @@ -744,6 +752,55 @@ def github_token() -> str: return result.stdout.strip() if result.returncode == 0 else "" +def gh_user() -> str: + """Who `gh` is signed in to github.com as, from its hosts.yml. + + Read as plain lines rather than YAML: only the host's own `user:` key + matters, and reading it costs no process, so a reused answer never has to + run `gh` to know whose it is. + """ + config = os.environ.get("GH_CONFIG_DIR") or str(Path(os.environ.get("XDG_CONFIG_HOME") or (Path.home() / ".config")) / "gh") + try: + lines = (Path(config) / "hosts.yml").read_text(encoding="utf-8").splitlines() + except (OSError, UnicodeDecodeError): + return "" + in_host, child_indent = False, None + for line in lines: + stripped = line.strip() + if not stripped or stripped.startswith("#"): + continue + indent = len(line) - len(line.lstrip(" ")) + if indent == 0: + in_host, child_indent = stripped == "github.com:", None + continue + if not in_host: + continue + child_indent = indent if child_indent is None else child_indent + if indent == child_indent and stripped.startswith("user:"): + return stripped.split(":", 1)[1].strip().strip("'\"") + return "" + + +def github_credential() -> tuple[str, str]: + """Who the quota is asked as, and the token when it is already in hand. + + Environment first, so a user can scope exactly what this sees, then the + `gh` CLI's own credential store, which is where a machine that has signed + in already keeps it. An exported token keys by its own hash. gh keys by the + account its hosts.yml names, and its token is only fetched when GitHub is + asked; a hosts.yml without a user falls back to fetching it now. Nothing + here is ever stored or logged. + """ + token = env_token() + if token: + return short_hash(token), token + user = gh_user() if shutil.which("gh") else "" + if user: + return short_hash(f"gh:github.com:{user}"), "" + token = gh_token() + return (short_hash(token), token) if token else ("", "") + + def ask_quota(token: str) -> dict[str, Any]: """Ask GitHub for the account-wide credit balance.""" request = urllib.request.Request( @@ -759,7 +816,11 @@ def ask_quota(token: str) -> dict[str, Any]: payload = json.loads(response.read(2_000_000).decode("utf-8", errors="replace")) except urllib.error.HTTPError as exc: return {"available": False, "reason": f"http-{exc.code}"} - except (urllib.error.URLError, OSError, ValueError, TimeoutError): + # A reply that came back unreadable is still an answer from GitHub. + except (http.client.HTTPException, ValueError): + return {"available": False, "reason": "bad-payload"} + # No route, no DNS, refused, reset or timed out: nothing answered. + except OSError: return {"available": False, "reason": "unreachable"} if not isinstance(payload, dict): @@ -782,20 +843,58 @@ def ask_quota(token: str) -> dict[str, Any]: } -def fetch_remote_quota() -> dict[str, Any]: +def read_json(path: Path) -> dict[str, Any]: + try: + data = json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError): + return {} + return data if isinstance(data, dict) else {} + + +# Before the limits caches were named copilot-limits-*, each token kept its +# answer as copilot-quota-. Read once, when this account has no +# answer under the new name, so the last allowance survives the rename. +def adopt_legacy_cache(cache: Path, token: str) -> dict[str, Any]: + legacy = cache_root() / f"copilot-quota-{short_hash(token)}.json" + cached = read_json(legacy) + if not cached.get("available"): + return {} + try: + write_json(cache, cached) + legacy.unlink(missing_ok=True) + except OSError: + pass + return cached + + +def fetch_remote_quota(force: bool = False) -> dict[str, Any]: """The live quota, or the last good one marked stale. Every failure is soft: a missing token, an offline machine, or a changed - payload keeps the last answer this token got, marked stale, or leaves the record - on its local stats rather than failing the whole collection. The cache is - keyed by the token, so a machine signed in to another account never shows - the last one's allowance. + payload keeps the last answer this account got, marked stale, or leaves + the record on its local stats rather than failing the whole collection. + Each account keeps its own cache, so a machine signed in to another + account never shows the last one's allowance. The cache is read whatever + its file time, so a clock that went back doesn't lose it. """ - token = github_token() + key, token = github_credential() + if not key: + return {"available": False, "reason": "no-token"} + + cache = cache_root() / f"copilot-limits-{key}.json" + cached = read_json(cache) + # A panel opened and shut repeatedly must not turn into a request (and a + # gh process) per flick, so an answer this recent is the answer. + age = time.time() - number(cached.get("fetchedAtMs")) / 1000 + if not force and cached.get("available") and number(cached.get("fetchedAtMs")) > 0 and 0 <= age < PROBE_MIN_INTERVAL_SECONDS: + return open_quota(dict(cached, stale=False)) + + token = token or gh_token() if not token: return {"available": False, "reason": "no-token"} + if not cached: + cached = adopt_legacy_cache(cache, token) - cache = cache_root() / f"copilot-quota-{hashlib.sha256(token.encode('utf-8')).hexdigest()[:16]}.json" quota = ask_quota(token) if quota.get("available"): quota.update(stale=False, fetchedAtMs=round(time.time() * 1000)) @@ -805,25 +904,31 @@ def fetch_remote_quota() -> dict[str, Any]: pass return quota - cached = read_fresh_json(cache, float("inf")) or {} if cached.get("available"): - if not quota_window_open(cached): - cached.update(used=0.0, resetsAt="") - cached.update(stale=True, failure=quota.get("reason", "")) - return cached + return open_quota(dict(cached, stale=True, failure=quota.get("reason", ""))) return quota -# A kept answer describes its own allowance window and nothing after it: once -# that window resets, its usage is back at zero. Without a reset date, the window is -# the calendar month (UTC) it was fetched in. -def quota_window_open(cached: dict[str, Any]) -> bool: - now = dt.datetime.now(dt.timezone.utc) +# A kept answer describes its own allowance window and nothing after it: +# once GitHub's reset date passes, its usage is back at zero and the next +# reset is unknown until GitHub is asked again. One without a reset date is +# kept as it is. +def open_quota(cached: dict[str, Any]) -> dict[str, Any]: resets_at = parse_event_time(str(cached.get("resetsAt") or "")) - if resets_at is not None: - return resets_at > now - fetched = number(cached.get("fetchedAtMs")) / 1000 - return fetched > 0 and utc_month_key(fetched) == now.strftime("%Y-%m") + if resets_at is not None and resets_at <= dt.datetime.now(dt.timezone.utc): + cached.update(used=0.0, resetsAt="") + return cached + + +# What went wrong asking GitHub, and what to do about it. +def quota_trouble(reason: str) -> str: + if reason == "unreachable": + return "Couldn't reach GitHub's Copilot quota endpoint. Retrying shortly." + if reason in ("http-401", "http-403"): + return "GitHub refused the token for Copilot's quota. " + AUTH_HELP + if reason.startswith("http-"): + return f"GitHub's Copilot quota endpoint returned status {reason[5:]}." + return "GitHub's Copilot quota endpoint sent no allowance." # ---------------------------------------------------------------- allowance @@ -847,40 +952,53 @@ def allowance_limits( "limits": [], "tierLabel": "", "usageStatusText": "", - "authHelpText": AUTH_HELP, + "authHelpText": "", "limitsStale": False, "limitsFetchedAt": 0, + "retryAdvised": False, } - resets_at = next_month_start(now).isoformat() - - live = remote if isinstance(remote, dict) and remote.get("available") else None - if live: - result["tierLabel"] = str(live.get("planLabel") or "") - if live.get("unlimited") or float(live.get("entitlement") or 0) <= 0: + remote = remote if isinstance(remote, dict) else {} + reason = str(remote.get("failure") or remote.get("reason") or "") + result["retryAdvised"] = reason == "unreachable" + + if remote.get("available"): + result["tierLabel"] = str(remote.get("planLabel") or "") + if remote.get("stale"): + result["authHelpText"] = quota_trouble(reason) + " Showing the last known allowance." + if remote.get("unlimited") or float(remote.get("entitlement") or 0) <= 0: return result - parsed = parse_event_time(str(live.get("resetsAt") or "")) - if parsed is not None and parsed > now: - resets_at = parsed.isoformat() + # GitHub's own reset date, or none: a kept window that has reset has no + # known next one. + parsed = parse_event_time(str(remote.get("resetsAt") or "")) result["limits"] = [{ "label": "Monthly allowance", - "percent": min(1.0, float(live.get("used") or 0.0) / float(live["entitlement"])), - "resetsAt": resets_at, + "percent": min(1.0, float(remote.get("used") or 0.0) / float(remote["entitlement"])), + "resetsAt": parsed.isoformat() if parsed is not None and parsed > now else "", }] - result["limitsStale"] = bool(live.get("stale")) - result["limitsFetchedAt"] = number(live.get("fetchedAtMs")) + result["limitsStale"] = bool(remote.get("stale")) + result["limitsFetchedAt"] = number(remote.get("fetchedAtMs")) return result - if remote is not None and remote.get("reason") not in ("no-token", "disabled"): + if reason == "no-token": + result["usageStatusText"] = "Waiting for auth" + result["authHelpText"] = AUTH_HELP + elif reason in ("http-401", "http-403"): + result["usageStatusText"] = "Sign-in expired" + result["authHelpText"] = quota_trouble(reason) + elif reason and reason != "disabled": result["usageStatusText"] = "Copilot limits unavailable" + result["authHelpText"] = quota_trouble(reason) # A budget must not advertise a meter on a machine the CLI has never - # touched: the estimate needs local numbers to mean anything. + # touched: the estimate needs local numbers to mean anything. Its window is + # the calendar month in UTC, the boundary GitHub documents for Copilot's + # allowance. budget = config.get("monthlyCredits") if isinstance(budget, (int, float)) and budget > 0 and has_stats: result["limits"] = [{ "label": "Monthly allowance (est.)", "percent": min(1.0, month_credits / float(budget)), - "resetsAt": resets_at, + "resetsAt": next_month_start(now).isoformat(), }] return result @@ -912,17 +1030,18 @@ def main() -> int: # machine with gh signed in but no Copilot stays an empty record. remote = {"available": False, "reason": "disabled"} if config.get("remote") is True or (config.get("remote", True) is not False and home_dir.is_dir()): - remote = fetch_remote_quota() + remote = fetch_remote_quota(args.force) has_stats = number(stats.get("totalPrompts")) > 0 or float(stats.get("monthCredits") or 0.0) > 0 allowance = allowance_limits(float(stats.get("monthCredits") or 0.0), config, remote, now, has_stats) + ready = number(stats.get("totalPrompts")) > 0 or len(allowance["limits"]) > 0 record: dict[str, Any] = { "schemaVersion": 1, "id": AGENT_ID, "name": AGENT_NAME, "updatedAt": now.isoformat(), - "ready": number(stats.get("totalPrompts")) > 0 or len(allowance["limits"]) > 0, + "ready": ready, "hasLocalStats": True, "hasPromptStats": stats.get("hasPromptStats") is not False, "tierLabel": allowance["tierLabel"], @@ -932,6 +1051,9 @@ def main() -> int: "limitsStale": allowance["limitsStale"], "limitsFetchedAt": allowance["limitsFetchedAt"], } + # The first probe after login often fires before the network is up. + if allowance["retryAdvised"]: + record["retryAdvised"] = True for key in ( "todayPrompts", "todaySessions", diff --git a/test/shell.d/agent-usage-copilot-test.sh b/test/shell.d/agent-usage-copilot-test.sh index 746969d950d..ce75c1648d0 100755 --- a/test/shell.d/agent-usage-copilot-test.sh +++ b/test/shell.d/agent-usage-copilot-test.sh @@ -187,7 +187,7 @@ record = run() collector.urllib.request.urlopen = offline stale = run() # Once the kept answer's window has reset, its usage is back at zero. -for cache in collector.cache_root().glob("copilot-quota-*.json"): +for cache in collector.cache_root().glob("copilot-limits-*.json"): kept = json.loads(cache.read_text()) cache.write_text(json.dumps(dict(kept, resetsAt="2020-01-01T00:00:00.000Z"))) lapsed = run() @@ -206,7 +206,7 @@ pass "Copilot collector stamps the live allowance" fail "Copilot collector keeps the last live allowance, stale, after a failed check" "$result" pass "Copilot collector keeps the last live allowance, stale, after a failed check" -[[ $(jq -c '.lapsed | {label: .limits[0].label, zero: (.limits[0].percent == 0), stale: .limitsStale, tierLabel}' <<<"$result") == '{"label":"Monthly allowance","zero":true,"stale":true,"tierLabel":"Business"}' ]] || +[[ $(jq -c '.lapsed | {label: .limits[0].label, zero: (.limits[0].percent == 0), resetsAt: .limits[0].resetsAt, stale: .limitsStale, tierLabel}' <<<"$result") == '{"label":"Monthly allowance","zero":true,"resetsAt":"","stale":true,"tierLabel":"Business"}' ]] || fail "Copilot collector resets a kept allowance once its window ends" "$result" pass "Copilot collector resets a kept allowance once its window ends" @@ -348,3 +348,157 @@ result=$(run_collector_in "$EMPTY_HOME") [[ $(jq -r '.ready' <<<"$result") == "false" && $(jq -c '.limits' <<<"$result") == "[]" ]] || fail "Copilot collector stays hidden on a clean machine" "$result" pass "Copilot collector stays hidden on a clean machine" + +# The quota as gh's account sees it: a fake gh logs every run, and hosts.yml +# names who it is signed in as, so a reused answer needs no gh at all. +mkdir -p "$TEST_HOME/fakebin" "$TEST_HOME/nogh" "$TEST_HOME/gh" +cat >"$TEST_HOME/fakebin/gh" <<'EOF' +#!/bin/bash +printf '%s\n' "$*" >>"$GH_LOG" +[[ $1 == "auth" && $2 == "token" ]] && echo "gh-token" +EOF +chmod +x "$TEST_HOME/fakebin/gh" +printf 'github.com:\n users:\n tester:\n git_protocol: ssh\n user: tester\n' >"$TEST_HOME/gh/hosts.yml" +printf '{"monthlyCredits": 3000}' >"$TEST_HOME/.config/omarchy/agents/copilot.json" + +result=$(HOME="$TEST_HOME" COPILOT_HOME="$TEST_HOME/.copilot" XDG_CONFIG_HOME="$TEST_HOME/.config" \ + XDG_CACHE_HOME="$TEST_HOME/.cache-limits" GH_CONFIG_DIR="$TEST_HOME/gh" GH_LOG="$TEST_HOME/gh.log" \ + FAKEBIN="$TEST_HOME/fakebin" NOGH="$TEST_HOME/nogh" \ + python3 - "$ROOT/bin/omarchy-agent-usage-copilot" <<'PY' +import datetime as dt +import hashlib +import importlib.machinery +import importlib.util +import io +import json +import os +import sys +import time +from pathlib import Path + +for var in ("COPILOT_QUOTA_TOKEN", "GH_TOKEN", "GITHUB_TOKEN"): + os.environ.pop(var, None) +os.environ["PATH"] = os.environ["FAKEBIN"] + ":" + os.environ["PATH"] + +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) + +reset = (dt.datetime.now(dt.timezone.utc) + dt.timedelta(days=10)).replace(microsecond=0) +payload = { + "copilot_plan": "individual", + "quota_reset_date_utc": reset.strftime("%Y-%m-%dT%H:%M:%S.000Z"), + "quota_snapshots": {"premium_interactions": {"unlimited": False, "entitlement": 300, "credits_used": 30}}, +} +gh_log = Path(os.environ["GH_LOG"]) +asked = [] + + +def answer(request, timeout=None): + asked.append(request.get_header("Authorization")) + return io.BytesIO(json.dumps(payload).encode()) + + +def offline(request, timeout=None): + raise collector.urllib.error.URLError("offline") + + +def status(code): + def urlopen(request, timeout=None): + raise collector.urllib.error.HTTPError(request.full_url, code, "", {}, None) + return urlopen + + +def unasked(request, timeout=None): + raise AssertionError("the quota should not have been asked") + + +def run(urlopen, *args): + collector.urllib.request.urlopen = urlopen + out = io.StringIO() + sys.argv = ["omarchy-agent-usage-copilot", *args] + sys.stdout, real_stdout = out, sys.stdout + collector.main() + sys.stdout = real_stdout + record = json.loads(out.getvalue()) + return { + "limits": [[limit["label"], limit["percent"], limit["resetsAt"] != ""] for limit in record["limits"]], + "stale": record["limitsStale"], + "status": record["usageStatusText"], + "help": record["authHelpText"], + "retry": record.get("retryAdvised", False), + "tier": record["tierLabel"], + "gh": len(gh_log.read_text().splitlines()) if gh_log.exists() else 0, + } + + +out = {} +out["live"] = run(answer, "--force") +out["reused"] = run(unasked) +out["limitsOnly"] = run(unasked, "--limits-only") +out["forced"] = run(answer, "--force") +out["asked"] = sorted(set(asked)) +out["offline"] = run(offline, "--force") + +# A cache stamped in the future (the clock went back) is still the answer +# kept past a failed check. +cache = next(collector.cache_root().glob("copilot-limits-*.json")) +future = time.time() + 86400 +os.utime(cache, (future, future)) +out["future"] = run(offline, "--force") +out["status500"] = run(status(500), "--force") + +# Another gh account has nothing kept; how it failed decides what it says. +(Path(os.environ["GH_CONFIG_DIR"]) / "hosts.yml").write_text("github.com:\n user: someone-else\n") +out["otherOffline"] = run(offline, "--force") +out["other500"] = run(status(500), "--force") +out["other401"] = run(status(401), "--force") + +# No token anywhere, with the CLI in use: waiting for a sign-in. +os.environ["PATH"] = os.environ["NOGH"] +out["noToken"] = run(unasked, "--force") + +# An answer under the old copilot-quota-* name is adopted once. +os.environ["COPILOT_QUOTA_TOKEN"] = "legacy-token" +legacy = collector.cache_root() / f"copilot-quota-{hashlib.sha256(b'legacy-token').hexdigest()[:16]}.json" +legacy.write_text(json.dumps({"available": True, "entitlement": 100, "used": 40, "resetsAt": reset.isoformat(), "planLabel": "Pro", "fetchedAtMs": 1000})) +out["legacy"] = run(offline, "--force") +out["legacyMoved"] = [legacy.exists(), (collector.cache_root() / f"copilot-limits-{hashlib.sha256(b'legacy-token').hexdigest()[:16]}.json").exists()] +print(json.dumps(out)) +PY +) + +[[ $(jq -c '.live' <<<"$result") == '{"limits":[["Monthly allowance",0.1,true]],"stale":false,"status":"","help":"","retry":false,"tier":"Pro","gh":1}' ]] || + fail "Copilot collector asks the quota as gh's account" "$result" +pass "Copilot collector asks the quota as gh's account" + +[[ $(jq -c '[.reused, .limitsOnly] | map({limits, stale, help, gh})' <<<"$result") == '[{"limits":[["Monthly allowance",0.1,true]],"stale":false,"help":"","gh":1},{"limits":[["Monthly allowance",0.1,true]],"stale":false,"help":"","gh":1}]' ]] || + fail "Copilot collector reuses a recent quota answer without running gh" "$result" +[[ $(jq -c '{gh: .forced.gh, asked}' <<<"$result") == '{"gh":2,"asked":["Bearer gh-token"]}' ]] || + fail "Copilot collector asks again on --force" "$result" +pass "Copilot collector reuses a recent quota answer without running gh, and --force asks again" + +[[ $(jq -c '.offline | {stale, status, retry, kept: (.limits | length), help: (.help | test("Retrying"))}' <<<"$result") == '{"stale":true,"status":"","retry":true,"kept":1,"help":true}' ]] || + fail "Copilot collector advises a retry when GitHub can't be reached" "$result" +[[ $(jq -c '.status500 | {stale, retry, help: (.help | test("500"))}' <<<"$result") == '{"stale":true,"retry":false,"help":true}' ]] || + fail "Copilot collector doesn't advise a retry for an error status" "$result" +pass "Copilot collector advises a retry only when GitHub can't be reached" + +[[ $(jq -c '.future | {stale, kept: (.limits | length)}' <<<"$result") == '{"stale":true,"kept":1}' ]] || + fail "Copilot collector keeps a cache whose file time is in the future" "$result" +pass "Copilot collector keeps a cache whose file time is in the future" + +[[ $(jq -c '[.otherOffline, .other500, .other401] | map({status, retry, help: (.help != "")})' <<<"$result") == '[{"status":"Copilot limits unavailable","retry":true,"help":true},{"status":"Copilot limits unavailable","retry":false,"help":true},{"status":"Sign-in expired","retry":false,"help":true}]' ]] || + fail "Copilot collector says what failed when nothing is kept" "$result" +[[ $(jq -r '.other500.help' <<<"$result") == *"status 500"* && $(jq -r '.other500.help' <<<"$result") != *"gh auth login"* ]] || + fail "Copilot collector gives an error status as the detail, not a sign-in hint" "$result" +pass "Copilot collector says what failed when nothing is kept" + +[[ $(jq -c '.noToken | {status, help: (.help | test("gh auth login"))}' <<<"$result") == '{"status":"Waiting for auth","help":true}' ]] || + fail "Copilot collector waits for a GitHub sign-in" "$result" +pass "Copilot collector waits for a GitHub sign-in" + +[[ $(jq -c '{limits: .legacy.limits, stale: .legacy.stale, moved: .legacyMoved}' <<<"$result") == '{"limits":[["Monthly allowance",0.4,true]],"stale":true,"moved":[false,true]}' ]] || + fail "Copilot collector adopts an answer cached under the old name" "$result" +pass "Copilot collector adopts an answer cached under the old name" From 4105f1701f4c7245b5b8a3db332249390630a89c Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 22:57:57 -0400 Subject: [PATCH 053/546] Keep Antigravity's undated and out-of-range times from crashing its scan Co-authored-by: steph4n <284007938+steph4n-gh@users.noreply.github.com> Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-agy | 31 ++- .../agent-usage-agy-timestamps-test.sh | 211 ++++++++++++++++++ 2 files changed, 234 insertions(+), 8 deletions(-) create mode 100755 test/shell.d/agent-usage-agy-timestamps-test.sh diff --git a/bin/omarchy-agent-usage-agy b/bin/omarchy-agent-usage-agy index 32432dc7b94..969416d04e2 100755 --- a/bin/omarchy-agent-usage-agy +++ b/bin/omarchy-agent-usage-agy @@ -67,22 +67,37 @@ def write_json(path, payload): raise +# Times on disk are untrusted: a boolean, an epoch at or below zero, one too +# large for a date, NaN, or text that isn't ISO is no time at all, and the +# row it dates stays undated rather than taking the record down. def parse_time(value): - if value is None or value == "": + if value is None or value == "" or isinstance(value, bool): 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: + if isinstance(value, (int, float)) or str(value).isdigit(): + stamp = float(value) + if not stamp > 0: + return None + parsed = datetime.fromtimestamp(stamp / 1000 if stamp > 1e12 else stamp, timezone.utc) + else: + parsed = datetime.fromisoformat(str(value).replace("Z", "+00:00")) + if not parsed.tzinfo: + parsed = parsed.replace(tzinfo=timezone.utc) + except (ValueError, OverflowError, OSError): return None - return parsed if parsed.tzinfo else parsed.replace(tzinfo=timezone.utc) + # Go's zero time marks an uninitialized summary, not dated activity. + return None if parsed == datetime.min.replace(tzinfo=timezone.utc) else parsed def local_day(value): parsed = parse_time(value) - return parsed.astimezone().date().isoformat() if parsed else "" + if parsed is None: + return "" + try: + return parsed.astimezone().date().isoformat() + except (ValueError, OverflowError, OSError): + # Even a parsed date can cross years 1/9999 during local conversion. + return "" # The sign-in agy keeps in the keyring: its access token, when that lapses, diff --git a/test/shell.d/agent-usage-agy-timestamps-test.sh b/test/shell.d/agent-usage-agy-timestamps-test.sh new file mode 100755 index 00000000000..71e335a886c --- /dev/null +++ b/test/shell.d/agent-usage-agy-timestamps-test.sh @@ -0,0 +1,211 @@ +#!/bin/bash + +set -euo pipefail + +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" + +require_command python3 + +python3 - "$ROOT/bin/omarchy-agent-usage-agy" <<'PY' +import contextlib +import io +import json +import os +import runpy +import sqlite3 +import subprocess +import sys +import tempfile +import time +import types +import unittest +from datetime import datetime, timezone +from pathlib import Path +from unittest.mock import Mock, patch + +collector_path = Path(sys.argv.pop()) +collector = types.SimpleNamespace(**runpy.run_path(str(collector_path))) +# Functions loaded by runpy share this namespace, rather than the wrapper. +namespace = collector.local_stats.__globals__ +NOW = datetime(2026, 10, 2, 12, tzinfo=timezone.utc) +ZONES = ('UTC0', 'EST5', 'JST-9') # Fixed offsets; no host zoneinfo or DST rules. +INVALID = (None, '', 'not a date', 0, -1, '0', -1.5, '2026-02-30T12:00:00Z', [], {}, True, + '0001-01-01 00:00:00+00:00', '0001-01-01T00:00:00Z', + '0001-01-01T00:00:00', '0001-01-01T01:00:00+01:00', + 10**30, -(10**30), 10**400, '9' * 400, float('inf'), float('nan')) + + +@contextlib.contextmanager +def local_zone(zone): + old = os.environ.get('TZ') + try: + os.environ['TZ'] = zone + time.tzset() + yield + finally: + if old is None: + os.environ.pop('TZ', None) + else: + os.environ['TZ'] = old + time.tzset() + + +class FrozenDatetime(datetime): + @classmethod + def now(cls, tz=None): + return NOW.astimezone(tz) if tz else NOW.astimezone().replace(tzinfo=None) + + +def write_fixture(root, invalid): + history = [dict(timestamp='2026-10-02T02:00:00Z', conversation_id='valid')] + history += [dict(timestamp=value, conversation_id='bad-history') for value in invalid] + history += [{}, dict(timestamp='1969-07-20T12:00:00Z', conversation_id='old'), + dict(timestamp='2026-10-02T12:00:00Z', conversation_id='valid')] + (root / 'history.jsonl').write_text('\n'.join(map(json.dumps, history))) + with contextlib.closing(sqlite3.connect(root / 'conversation_summaries.db')) as db, db: + db.execute('CREATE TABLE conversation_summaries (conversation_id TEXT, last_modified_time TEXT)') + db.executemany('INSERT INTO conversation_summaries VALUES (?, ?)', [ + ('valid', '2026-10-02T12:00:00Z'), ('undated', '0001-01-01 00:00:00+00:00'), + ('undated', None), ('undated', 'not a date'), + ('summary-only', '2026-10-02T12:00:00Z')]) + transcript = root / 'brain/valid/.system_generated/logs/transcript.jsonl' + transcript.parent.mkdir(parents=True) + steps = [dict(created_at='2026-10-02T12:00:00Z', model='dated', + usage=dict(input_tokens=10, output_tokens=5))] + steps += [dict(created_at=value, model='undated', usage=dict(input_tokens=2, output_tokens=1)) + for value in invalid] + steps += [dict(model='undated', usage=dict(input_tokens=2, output_tokens=1)), + dict(created_at='2026-10-02T02:00:00Z', model='dated', + usage=dict(input_tokens=7, output_tokens=3))] + transcript.write_text('\n'.join(map(json.dumps, steps))) + + +class TimestampTests(unittest.TestCase): + def test_zero_and_invalid_times_are_undated_in_every_zone(self): + for zone in ZONES: + with local_zone(zone): + for value in INVALID: + with self.subTest(zone=zone, value=value): + self.assertEqual(collector.local_day(value), '') + self.assertIsNone(collector.parse_time('0001-01-01 00:00:00+00:00')) + + def test_valid_times_keep_existing_formats_and_local_dates(self): + for zone, early_day in zip(ZONES, ('2026-10-02', '2026-10-01', '2026-10-02')): + with local_zone(zone), self.subTest(zone=zone): + for value in ('2026-10-02T02:00:00Z', '2026-10-02 02:00:00', + '2026-10-02T04:00:00+02:00', 1790906400, 1790906400000, + '1790906400000'): + self.assertEqual(collector.local_day(value), early_day) + self.assertEqual(collector.local_day('1969-07-20T12:00:00Z'), '1969-07-20') + self.assertEqual(collector.local_day(1), '1969-12-31' if zone == 'EST5' else '1970-01-01') + self.assertEqual(collector.local_day('2026-10-02T20:00:00Z'), + '2026-10-03' if zone == 'JST-9' else '2026-10-02') + + def test_conversion_boundaries(self): + cases = ( + ('0001-01-01T00:00:01Z', ('0001-01-01', '', '0001-01-01')), + ('0001-01-01T12:00:00Z', ('0001-01-01', '0001-01-01', '0001-01-01')), + ('9999-12-31T23:59:59Z', ('9999-12-31', '9999-12-31', '')), + ('9999-12-31T12:00:00Z', ('9999-12-31', '9999-12-31', '9999-12-31')), + ('0001-01-01T00:00:01+14:00', ('', '', '')), + ('9999-12-31T23:59:59-12:00', ('', '', '')), + ) + for value, expected in cases: + for zone, day in zip(ZONES, expected): + with local_zone(zone), self.subTest(zone=zone, value=value): + self.assertEqual(collector.local_day(value), day) + + def test_mixed_activity_survives_without_inventing_dates(self): + for zone in ZONES: + # These offset timestamps overflow while normalizing to UTC in all zones. + invalid = INVALID + ('0001-01-01T00:00:01+14:00', '9999-12-31T23:59:59-12:00') + if zone == 'EST5': + invalid += ('0001-01-01T00:00:01Z',) + elif zone == 'JST-9': + invalid += ('9999-12-31T23:59:59Z',) + with local_zone(zone), tempfile.TemporaryDirectory() as scratch, self.subTest(zone=zone): + root = Path(scratch) + write_fixture(root, invalid) + with patch.dict(namespace, datetime=FrozenDatetime): + stats = collector.local_stats(root) + self.assertTrue(stats['hasLocalStats']) + self.assertEqual(stats['totalPrompts'], 3) + self.assertEqual(stats['todayPrompts'], 1 if zone == 'EST5' else 2) + # Summary identities remain countable without a date; bad history isn't. + self.assertEqual(stats['totalSessions'], 4) + self.assertEqual(stats['todaySessions'], 2) + dates = ['1969-07-20', '2026-10-02'] + if zone == 'EST5': + dates.insert(1, '2026-10-01') + self.assertEqual(stats['activeDates'], dates) + self.assertEqual(stats['activeDays'], len(dates)) + self.assertEqual(stats['todayTotalTokens'], 15 if zone == 'EST5' else 25) + self.assertEqual(stats['todayTokensByModel'], {'dated': stats['todayTotalTokens']}) + recent = {row['date']: row['messageCount'] for row in stats['recentDays']} + self.assertEqual(recent['2026-10-02'], stats['todayTotalTokens']) + self.assertEqual(sum(recent.values()), 25) + self.assertEqual(stats['modelUsage']['dated']['inputTokens'], 17) + self.assertEqual(stats['modelUsage']['dated']['outputTokens'], 8) + # Undated usage retains lifetime totals, just like a missing timestamp. + self.assertEqual(stats['modelUsage']['undated']['inputTokens'], 2 * (len(invalid) + 1)) + self.assertEqual(stats['modelUsage']['undated']['outputTokens'], len(invalid) + 1) + + def test_quota_reset_display_keeps_unknown_and_valid_windows(self): + # A usable quota remains useful without a trustworthy reset countdown. + cases = [({}, '')] + [({'resetTime': value}, '') for value in INVALID] + cases += [ + ({'resetTime': '2026-10-02T11:00:00Z'}, '2026-10-02T11:00:00+00:00'), + ({'resetTime': '2026-10-02T15:00:00+02:00'}, '2026-10-02T15:00:00+02:00'), + ({'resetTime': '2026-10-02 13:00:00'}, '2026-10-02T13:00:00+00:00'), + ({'resetTime': (NOW.timestamp() + 3600) * 1000}, '2026-10-02T13:00:00+00:00'), + ] + for zone in ZONES: + for reset, expected in cases: + with local_zone(zone), self.subTest(zone=zone, reset=reset): + bucket = dict(reset, window='5h', remainingFraction=0.75) + self.assertEqual(collector.reset_of(bucket), expected) + limits = collector.extract_limits({'groups': [ + {'displayName': 'Gemini', 'buckets': [bucket]}]}) + self.assertEqual(len(limits), 1) + self.assertEqual(limits[0]['resetsAt'], expected) + self.assertEqual(limits[0]['percent'], 0.25) + + def test_unrelated_errors_propagate(self): + with patch.dict(namespace, datetime=types.SimpleNamespace( + fromisoformat=lambda value: (_ for _ in ()).throw(RuntimeError('unexpected parser failure')))): + with self.assertRaisesRegex(RuntimeError, 'unexpected parser failure'): + collector.local_day('2026-10-02T12:00:00Z') + with patch.dict(namespace, parse_time=lambda value: types.SimpleNamespace( + astimezone=lambda: (_ for _ in ()).throw(RuntimeError('unexpected conversion failure')))): + with self.assertRaisesRegex(RuntimeError, 'unexpected conversion failure'): + collector.local_day('2026-10-02T12:00:00Z') + + def test_isolated_cli_smoke_and_cached_record(self): + for zone in ZONES: + with tempfile.TemporaryDirectory() as scratch, self.subTest(zone=zone): + root = Path(scratch) + write_fixture(root, ('0001-01-01 00:00:00+00:00', None, 'not a date')) + stub_dir = root / 'bin' + stub_dir.mkdir() + secret_tool = stub_dir / 'secret-tool' + secret_tool.write_text('#!/bin/bash\nexit 1\n') + secret_tool.chmod(0o755) + env = dict(os.environ, HOME=scratch, AGY_DIR=scratch, XDG_CACHE_HOME=str(root / 'cache'), + PATH=str(stub_dir) + os.pathsep + os.defpath, TZ=zone) + for flag in ('--force', '--limits-only'): + result = subprocess.run([str(collector_path), flag], env=env, capture_output=True, + text=True, check=True, timeout=10) + record = json.loads(result.stdout) + self.assertEqual(result.stderr, '') + self.assertEqual(record['id'], 'agy') + self.assertFalse(record['ready']) + self.assertEqual(record['limits'], []) + self.assertEqual(record['totalPrompts'], 3) + self.assertEqual(record['totalSessions'], 4) + self.assertEqual(record['modelUsage']['dated']['inputTokens'], 17) + + +unittest.main(verbosity=2) +PY + +pass "Antigravity timestamp boundaries, mixed records, and isolated CLI smoke" From eca61b5effa6f2949a1a0b8bd139a72b9a43da1b Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:00:21 -0400 Subject: [PATCH 054/546] Read Antigravity's transcript tokens and models in every shape agy writes Count the flat input/output/cache counts newer transcripts carry beside Gemini's nested ones, file tokens under the model a step or a Model Selection names, skip unreadable counts, count transcript days as active, and fold subagent conversations into the one that spawned them. Co-authored-by: Slawomir S Co-authored-by: Ashish Daga Co-authored-by: Cedric Darne Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-agy | 153 +++++++++++++++---- test/shell.d/agent-usage-agy-scanner-test.sh | 56 +++++++ 2 files changed, 180 insertions(+), 29 deletions(-) diff --git a/bin/omarchy-agent-usage-agy b/bin/omarchy-agent-usage-agy index 969416d04e2..970bf6c3987 100755 --- a/bin/omarchy-agent-usage-agy +++ b/bin/omarchy-agent-usage-agy @@ -18,6 +18,7 @@ import argparse import base64 import json import os +import re import sqlite3 import subprocess import tempfile @@ -35,6 +36,8 @@ EXPIRED_HELP = "Antigravity's saved sign-in expired. Start `agy` to refresh it." PROBE_REUSE_SECONDS = 60 # A limits-only refresh needs fresh limits, not a fresh history scan. LIMITS_ONLY_SCAN_REUSE_SECONDS = 900 +# More tokens than any one step could hold is a damaged field, not usage. +MAX_TOKENS = 1e12 def app_root(): @@ -269,6 +272,67 @@ def model_bucket(models, model): return models.setdefault(model, {"inputTokens": 0, "outputTokens": 0, "cacheReadInputTokens": 0, "cacheCreationInputTokens": 0}) +# A token count from disk is untrusted: a boolean, a negative, NaN, an +# absurdly large or a non-numeric value counts as nothing, so one bad field +# costs that field rather than the scan. +def number(value): + if isinstance(value, bool): + return 0 + try: + count = float(value or 0) + except (TypeError, ValueError, OverflowError): + return 0 + return round(count) if 0 < count < MAX_TOKENS else 0 + + +# A transcript step's tokens as (input, output, cache read, cache write). +# Steps nest Gemini's counts under usage, where the prompt count includes +# the cached part and thinking sits beside the candidates; newer builds put +# input_tokens, output_tokens, cache_read_tokens and cache_write_tokens on the +# step itself. Every reader of those adds the four up, so input here excludes +# the cache reads, the way Anthropic's counts do. None when the step reports +# no usage. +def step_tokens(step): + usage = step.get("usage") + for source in ([usage] if isinstance(usage, dict) else []) + [step]: + if "prompt_token_count" in source or "candidates_token_count" in source: + prompt = number(source.get("prompt_token_count")) + cached = min(prompt, number(source.get("cached_content_token_count"))) + output = number(source.get("candidates_token_count")) + number(source.get("thoughts_token_count")) + return prompt - cached, output, cached, 0 + if any(key in source for key in ("input_tokens", "output_tokens", "cache_read_tokens", "cache_write_tokens")): + return (number(source.get("input_tokens")), number(source.get("output_tokens")), + number(source.get("cache_read_tokens")), number(source.get("cache_write_tokens"))) + return None + + +# A model step names its model by id; a "Model Selection" step names it the +# way the picker shows it ("Gemini 3.8 Flash (High)"). Both fold to the same +# key, so a conversation's tokens don't split across two spellings. +MODEL_SELECTION = re.compile(r"Model Selection` (?:from .+? )?to (.+?)\.(?:\s|$|<)") + + +def model_key(name): + key = re.sub(r"[\s_]+", "-", re.sub(r"[()]", " ", name).strip().lower()) + return re.sub(r"-+", "-", key).strip("-") + + +def step_model(step): + for field in ("model", "model_name", "model_id"): + value = step.get(field) + if isinstance(value, str) and model_key(value): + return model_key(value) + return "" + + +def selected_model(step): + content = step.get("content") + if not isinstance(content, str) or "Model Selection" not in content: + return "" + match = MODEL_SELECTION.search(content) + return model_key(match.group(1)) if match else "" + + def history_entries(root): try: lines = (root / "history.jsonl").read_text(encoding="utf-8", errors="ignore").splitlines() @@ -283,9 +347,28 @@ def history_entries(root): yield entry +# Each conversation in the summaries with when it last changed and, where +# agy records it, the conversation it was spawned from. +def conversation_summaries(root): + summaries = root / "conversation_summaries.db" + if not summaries.is_file(): + return [] + try: + conn = sqlite3.connect(summaries.resolve().as_uri() + "?mode=ro", uri=True, timeout=2) + try: + columns = {row[1] for row in conn.execute("PRAGMA table_info(conversation_summaries)")} + parent = "parent_conversation_id" if "parent_conversation_id" in columns else "NULL" + return conn.execute(f"SELECT conversation_id, last_modified_time, {parent} FROM conversation_summaries").fetchall() + finally: + conn.close() + except sqlite3.Error: + return [] + + # 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. +# reports. Steps that report no usage add no tokens; steps without a usable +# date still count toward the lifetime totals. def local_stats(root): today = datetime.now().astimezone().date() week = [(today - timedelta(days=offset)).isoformat() for offset in range(6, -1, -1)] @@ -295,60 +378,72 @@ def local_stats(root): days, sessions, today_sessions = set(), set(), set() today_by_model, models = {}, {} + # A subagent's conversation is part of the one that spawned it, so it + # counts toward that conversation rather than as a session of its own. + summaries = conversation_summaries(root) + parents = {str(conversation): str(parent) for conversation, _, parent in summaries if conversation and parent} + + def session_of(conversation): + conversation, seen = str(conversation), set() + while conversation in parents and conversation not in seen: + seen.add(conversation) + conversation = parents[conversation] + return conversation + 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") + conversation = entry.get("conversationId") or entry.get("conversation_id") or entry.get("sessionId") if conversation: - sessions.add(conversation) + sessions.add(session_of(conversation)) if day == today: today_prompts += 1 if conversation: - today_sessions.add(conversation) + today_sessions.add(session_of(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 conversation, modified, _ in summaries: + if conversation: + sessions.add(session_of(conversation)) + if local_day(modified) == today: + today_sessions.add(session_of(conversation)) 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 + # The model a conversation runs on carries from step to step until a + # step names another one. + current = "" 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): + if not isinstance(step, dict): + continue + current = step_model(step) or selected_model(step) or current + tokens = step_tokens(step) + if not tokens or not any(tokens): 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") + model = current or "unknown" bucket = model_bucket(models, model) - bucket["inputTokens"] += uncached - bucket["outputTokens"] += output - day = local_day(step.get("created_at")) + bucket["inputTokens"] += tokens[0] + bucket["outputTokens"] += tokens[1] + bucket["cacheReadInputTokens"] += tokens[2] + bucket["cacheCreationInputTokens"] += tokens[3] + total = sum(tokens) + day = local_day(step.get("created_at") or step.get("timestamp")) + if day: + days.add(day) if day in by_day: - by_day[day] += uncached + output + by_day[day] += total if day == today: - today_tokens += uncached + output - today_by_model[model] = today_by_model.get(model, 0) + uncached + output + today_tokens += total + today_by_model[model] = today_by_model.get(model, 0) + total return { "hasLocalStats": prompts > 0 or bool(sessions) or bool(models), diff --git a/test/shell.d/agent-usage-agy-scanner-test.sh b/test/shell.d/agent-usage-agy-scanner-test.sh index a7e6d7539d5..55a18ab3415 100755 --- a/test/shell.d/agent-usage-agy-scanner-test.sh +++ b/test/shell.d/agent-usage-agy-scanner-test.sh @@ -70,6 +70,62 @@ pass "Antigravity collector parses input and output tokens" fail "Antigravity collector makes up no limits from prompt counts" "$result" pass "Antigravity collector makes up no limits from prompt counts" +# Transcripts in both shapes agy writes: Gemini's counts nested under usage, +# with the cached part inside the prompt count, and the four flat counts on +# the step. A "Model Selection" step names the model the following steps run +# on, and a step that names its own model id wins. +mkdir -p "$AGY_DIR/brain/flat-conv/.system_generated/logs" +cat >"$AGY_DIR/brain/test-conv/.system_generated/logs/transcript.jsonl" <"$AGY_DIR/brain/flat-conv/.system_generated/logs/transcript.jsonl" <"$AGY_DIR/history.jsonl" <"$TEST_HOME/bin/secret-tool" <<'EOF' #!/bin/bash From c25be492e4ce31b84f3d0b1ac422ad88f791c779 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:04:57 -0400 Subject: [PATCH 055/546] Bring the Antigravity collector in line with the usage collector contract Cache limits per account, zero elapsed windows instead of dropping them, pause rather than expire limits while agy's access token lapses, advise a retry only when Google couldn't be reached, clear the help text when healthy, reuse probes for 15 s and scans for 20 s, mark the machine in use from local history alone, read snake_case and ungrouped quota buckets, and send the token only to an https CLOUD_CODE_URL. Co-authored-by: ivanvan08 Co-authored-by: Cedric Darne Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-agy | 298 +++++++++++++----- test/shell.d/agent-usage-agy-limits-test.sh | 168 +++++++++- .../agent-usage-agy-timestamps-test.sh | 4 +- 3 files changed, 382 insertions(+), 88 deletions(-) diff --git a/bin/omarchy-agent-usage-agy b/bin/omarchy-agent-usage-agy index 970bf6c3987..eb5345e2484 100755 --- a/bin/omarchy-agent-usage-agy +++ b/bin/omarchy-agent-usage-agy @@ -16,6 +16,9 @@ Secret Service keyring. The agents panel only ever reads the JSON this prints. import argparse import base64 +import fcntl +import hashlib +import http.client import json import os import re @@ -31,11 +34,13 @@ 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 +# How each source of a sign-in is named when it needs attention. +SOURCE_NAMES = {"agy": "`agy`"} +PROBE_MIN_INTERVAL_SECONDS = 15 +SCAN_REUSE_SECONDS = 20 # A limits-only refresh needs fresh limits, not a fresh history scan. -LIMITS_ONLY_SCAN_REUSE_SECONDS = 900 +LIMITS_ONLY_REUSE_SECONDS = 900 +DEFAULT_CLOUD_CODE_URL = "https://daily-cloudcode-pa.googleapis.com" # More tokens than any one step could hold is a damaged field, not usage. MAX_TOKENS = 1e12 @@ -104,7 +109,7 @@ def local_day(value): # The sign-in agy keeps in the keyring: its access token, when that lapses, -# and which kind of account it is. +# whether a refresh token stands behind it, and which kind of account it is. def login(): try: raw = subprocess.check_output( @@ -120,13 +125,20 @@ def login(): token = data.get("token") if isinstance(data.get("token"), dict) else {} expires = parse_time(token.get("expiry") or data.get("expiry")) return { + "source": "agy", "accessToken": str(token.get("access_token") or data.get("access_token") or ""), "expiresAt": expires.timestamp() if expires else 0, + "refresh": str(token.get("refresh_token") or data.get("refresh_token") or ""), "authMethod": str(data.get("auth_method") or "consumer"), "email": email_of(data.get("id_token")), } +def logins(): + entry = login() + return [entry] if entry else [] + + def email_of(id_token): try: claims = str(id_token).split(".")[1] @@ -135,6 +147,21 @@ def email_of(id_token): return "" +# Which Google account a sign-in belongs to, so the limits kept for it are +# its own: the email when the sign-in names one, otherwise its refresh token, +# which outlives the access tokens it mints. Only a digest of either names +# the cache file. +def limits_cache(entry): + identity = entry.get("email") or entry.get("refresh") or entry.get("accessToken") + if not identity: + return None + return cache_root() / f"agy-limits-{hashlib.sha256(identity.encode()).hexdigest()[:16]}.json" + + +def token_fresh(entry, now): + return bool(entry["accessToken"]) and (not entry["expiresAt"] or entry["expiresAt"] > now) + + def normalize_tier_label(name, tier_id=""): for word in ("ultra", "pro", "free"): if word in tier_id.lower() or word in name.lower(): @@ -147,7 +174,7 @@ def normalize_tier_label(name, tier_id=""): # 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) + tier = data.get(key) if isinstance(data, dict) else None if isinstance(tier, dict): name = str(tier.get("name") or "").strip() tier_id = str(tier.get("id") or "").strip() @@ -156,20 +183,35 @@ def extract_tier(data): return "" +# How much of a bucket is spent, or None when its fraction isn't a number. def used(bucket): - return round(max(0.0, min(1.0, 1.0 - float(bucket.get("remainingFraction", 1.0)))), 4) + remaining = bucket.get("remainingFraction", bucket.get("remaining_fraction", 1.0)) + if isinstance(remaining, bool): + return None + try: + remaining = float(remaining) + except (TypeError, ValueError): + return None + if remaining != remaining: + return None + return round(max(0.0, min(1.0, 1.0 - remaining)), 4) def reset_of(bucket): - parsed = parse_time(bucket.get("resetTime")) + parsed = parse_time(bucket.get("resetTime") or bucket.get("reset_time")) 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. +# they have been touched. An answer that lists its buckets without groups +# lists the Gemini ones. def extract_limits(data): + if not isinstance(data, dict): + return [] groups = [g for g in data.get("groups") or [] if isinstance(g, dict)] + if not groups and isinstance(data.get("buckets"), list): + groups = [{"displayName": "Gemini", "buckets": data["buckets"]}] if not groups: return [] gemini = next((g for g in groups if "gemini" in str(g.get("displayName") or "").lower()), groups[0]) @@ -177,7 +219,7 @@ def extract_limits(data): limits = [] session = weekly = None for bucket in gemini.get("buckets") or []: - if not isinstance(bucket, dict): + if not isinstance(bucket, dict) or used(bucket) is None: continue window = str(bucket.get("window") or "").lower() bucket_id = str(bucket.get("bucketId") or "").lower() @@ -197,14 +239,18 @@ def extract_limits(data): 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: + if isinstance(bucket, dict) and (used(bucket) or 0) > 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 +# The bearer token goes wherever CLOUD_CODE_URL points, so a destination +# named there is taken only over https, and anything else is Google's own. def post(endpoint, token): - base = (os.environ.get("CLOUD_CODE_URL") or "https://daily-cloudcode-pa.googleapis.com").rstrip("/") + base = (os.environ.get("CLOUD_CODE_URL") or DEFAULT_CLOUD_CODE_URL).rstrip("/") + if not base.startswith("https://"): + base = DEFAULT_CLOUD_CODE_URL request = urllib.request.Request( f"{base}/v1internal:{endpoint}", data=b"{}", headers={"Authorization": f"Bearer {token}", "Content-Type": "application/json", "User-Agent": "antigravity"}, @@ -213,58 +259,124 @@ def post(endpoint, token): 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): +# A window kept from an earlier check whose reset has passed is back to +# untouched, not unknown, so it reads zero until Google names the next reset. +# A damaged cache entry is skipped rather than taking the collector down. +def usable_cached_limits(cached): + limits = cached.get("limits") if isinstance(cached, dict) else None 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] + kept = [] + for limit in limits if isinstance(limits, list) else []: + if not isinstance(limit, dict): + continue + resets = parse_time(limit.get("resetsAt")) + kept.append(dict(limit, percent=0.0, resetsAt="") if resets and resets <= now else limit) + return kept -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 +def cached_answer(entry): + path = limits_cache(entry) + cached = (read_json(path) if path else None) or {} + return cached if isinstance(cached, dict) else {} + + +def fetched_at_ms(cached): + try: + return max(0, int(cached.get("fetchedAtMs") or 0)) + except (TypeError, ValueError, OverflowError): + return 0 - 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") +def lapsed_help(entry, kept): + name = SOURCE_NAMES[entry["source"]] + shown = " — showing the last known limits." if kept else "." + if entry["refresh"]: + return f"The Antigravity access token {name} keeps has lapsed{shown} Start {name} to refresh it; signing in again is not needed." + return f"The Antigravity sign-in {name} keeps expired{shown} Start {name} to sign in again." + + +def refused_help(entry): + name = SOURCE_NAMES[entry["source"]] + return f"Google refused the Antigravity sign-in {name} keeps. Start {name} to sign in again." + + +# The plan and limits, asked of Google with one sign-in after another until +# one answers. Each account keeps its own last good answer, which stands in, +# marked stale, when its check fails. +def collect_limits(entries, force): + result = {"limits": [], "tierLabel": "", "usageStatusText": "", "authHelpText": "", "live": False, "fetchedAtMs": 0, "retryAdvised": False} + now = time.time() + # Each harness refreshes its own token while it runs, so the one lapsing + # last was refreshed most recently. A token of unknown expiry comes after + # those, since only Google can say whether it still holds, and a lapsed + # one last. + ordered = sorted(entries, key=lambda e: (not token_fresh(e, now), -e["expiresAt"])) + fresh = [e for e in ordered if token_fresh(e, now)] + lead = next(iter(fresh or [e for e in ordered if e["accessToken"]] or ordered), None) + + def keep(entry): + cached = cached_answer(entry) if entry else {} + result.update(limits=usable_cached_limits(cached), tierLabel=str(cached.get("tierLabel") or ""), fetchedAtMs=fetched_at_ms(cached)) + + keep(lead) + if not lead or not lead["accessToken"]: + result.update(usageStatusText="Waiting for auth", authHelpText=AUTH_HELP) 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) + # until agy starts again, which needs no new sign-in while a refresh token + # stands behind it. + if not fresh: + result.update(usageStatusText="Limits paused" if lead["refresh"] else "Sign-in expired", authHelpText=lapsed_help(lead, result["limits"])) return result - if kept and not force and time.time() - fetched_at / 1000 < PROBE_REUSE_SECONDS: - result.update(limits=kept, live=True) + # A panel that is opened and shut repeatedly must not turn into a request + # per flick; --force is a person asking for fresh numbers. + if not force: + for entry in fresh: + cached = cached_answer(entry) + if cached and 0 <= now - fetched_at_ms(cached) / 1000 < PROBE_MIN_INTERVAL_SECONDS: + keep(entry) + result["live"] = True + return result + + refused, failure = [], None + for entry in fresh: + try: + tier = extract_tier(post("loadCodeAssist", entry["accessToken"])) + limits = extract_limits(post("retrieveUserQuotaSummary", entry["accessToken"])) + except urllib.error.HTTPError as error: + # A refused sign-in gives way to the next one; any other status is + # Google answering for every sign-in alike. + if error.code in (401, 403): + refused.append(entry) + continue + failure = entry, f"Google's quota service answered with status {error.code}.", False + break + except (urllib.error.URLError, OSError, http.client.HTTPException): + failure = entry, "Couldn't reach Google's quota service. Retrying shortly.", True + break + except ValueError: + failure = entry, "Google's quota service sent an answer that couldn't be read.", False + break + # A plan with no quota windows to report is healthy, with no meters. + fetched = round(time.time() * 1000) + result.update(limits=limits, tierLabel=tier, live=True, fetchedAtMs=fetched) + path = limits_cache(entry) + if path: + write_json(path, {"fetchedAtMs": fetched, "limits": limits, "tierLabel": tier}) 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: + if failure: + entry, help_text, transport = failure + keep(entry) + result.update(authHelpText=help_text, retryAdvised=transport) + if not result["limits"]: result["usageStatusText"] = "Antigravity limits unavailable" + return result + + keep(refused[0]) + result.update(usageStatusText="Sign-in expired", authHelpText=refused_help(refused[0])) return result @@ -333,10 +445,13 @@ def selected_model(step): return model_key(match.group(1)) if match else "" -def history_entries(root): +def history_entries(root, problems): try: lines = (root / "history.jsonl").read_text(encoding="utf-8", errors="ignore").splitlines() + except FileNotFoundError: + return except OSError: + problems.append("history") return for line in lines: try: @@ -349,7 +464,7 @@ def history_entries(root): # Each conversation in the summaries with when it last changed and, where # agy records it, the conversation it was spawned from. -def conversation_summaries(root): +def conversation_summaries(root, problems): summaries = root / "conversation_summaries.db" if not summaries.is_file(): return [] @@ -362,6 +477,7 @@ def conversation_summaries(root): finally: conn.close() except sqlite3.Error: + problems.append("summaries") return [] @@ -369,7 +485,8 @@ def conversation_summaries(root): # conversation summaries, and the tokens each model step in a transcript # reports. Steps that report no usage add no tokens; steps without a usable # date still count toward the lifetime totals. -def local_stats(root): +def local_stats(root, problems=None): + problems = [] if problems is None else problems today = datetime.now().astimezone().date() week = [(today - timedelta(days=offset)).isoformat() for offset in range(6, -1, -1)] today = today.isoformat() @@ -380,7 +497,7 @@ def local_stats(root): # A subagent's conversation is part of the one that spawned it, so it # counts toward that conversation rather than as a session of its own. - summaries = conversation_summaries(root) + summaries = conversation_summaries(root, problems) parents = {str(conversation): str(parent) for conversation, _, parent in summaries if conversation and parent} def session_of(conversation): @@ -390,7 +507,7 @@ def local_stats(root): conversation = parents[conversation] return conversation - for entry in history_entries(root): + for entry in history_entries(root, problems): day = local_day(entry.get("timestamp")) if not day: continue @@ -414,6 +531,7 @@ def local_stats(root): try: lines = transcript.read_text(encoding="utf-8", errors="ignore").splitlines() except OSError: + problems.append(str(transcript)) continue # The model a conversation runs on carries from step to step until a # step names another one. @@ -460,16 +578,39 @@ def local_stats(root): } -# A scan from before midnight would count yesterday's prompts as today's. +# A scan is reused for max_age seconds, and only on the day it was made: one +# from before midnight would count yesterday's prompts as today's. The scan +# is rebuilt under a lock so overlapping runs don't each walk the history, +# and one that couldn't read every source isn't kept. def cached_local_stats(root, max_age): - cache = cache_root() / "agy-stats.json" - cached = read_json(cache) or {} + digest = hashlib.sha256(str(root).encode()).hexdigest()[:16] + cache = cache_root() / f"agy-scan-{digest}.json" 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 reusable(): + cached = read_json(cache) + if not isinstance(cached, dict) or cached.get("scanDate") != today or not isinstance(cached.get("stats"), dict): + return None + try: + age = time.time() - float(cached.get("at")) + except (TypeError, ValueError): + return None + return cached["stats"] if 0 <= age < max_age else None + + stats = reusable() if max_age > 0 else None + if stats is not None: + return stats + cache.parent.mkdir(parents=True, exist_ok=True) + with open(cache.with_suffix(".lock"), "w") as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + stats = reusable() if max_age > 0 else None + if stats is not None: + return stats + problems = [] + stats = local_stats(root, problems) + if not problems: + write_json(cache, {"scanDate": today, "at": time.time(), "stats": stats}) + return stats def main(): @@ -478,14 +619,26 @@ def main(): 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) + if args.force: + scan_age = 0 + elif args.limits_only: + scan_age = LIMITS_ONLY_REUSE_SECONDS + else: + scan_age = SCAN_REUSE_SECONDS + stats = cached_local_stats(app_root(), scan_age) + + # Nobody signed in to Antigravity here and no local history: no limits and + # no plan, and the panel skips the record. + entries = logins() + ready = bool(entries) or bool(stats.get("hasLocalStats")) + limits = collect_limits(entries, args.force) if ready else { + "limits": [], "tierLabel": "", "usageStatusText": "", "authHelpText": "", "live": False, "fetchedAtMs": 0, "retryAdvised": False, + } tier = limits["tierLabel"] - if not tier and entry and entry["authMethod"] == "enterprise": + agy = next((e for e in entries if e["source"] == "agy"), None) + if not tier and agy and agy["authMethod"] == "enterprise": tier = "Enterprise" - elif not tier and entry and entry["authMethod"] in ("cloud", "vertex"): + elif not tier and agy and agy["authMethod"] in ("cloud", "vertex"): tier = "Cloud" record = { @@ -493,15 +646,16 @@ def main(): "id": AGENT_ID, "name": AGENT_NAME, "updatedAt": datetime.now(timezone.utc).isoformat(), - "ready": entry is not None, + "ready": ready, "tierLabel": tier, "usageStatusText": limits["usageStatusText"], "authHelpText": limits["authHelpText"], "limits": limits["limits"], "limitsStale": not limits["live"], "limitsFetchedAt": limits["fetchedAtMs"], + "retryAdvised": limits["retryAdvised"], } - record.update(cached_local_stats(app_root(), LIMITS_ONLY_SCAN_REUSE_SECONDS if args.limits_only and not args.force else 0)) + record.update(stats) print(json.dumps(record, separators=(",", ":"))) diff --git a/test/shell.d/agent-usage-agy-limits-test.sh b/test/shell.d/agent-usage-agy-limits-test.sh index 1198a4356e8..8f4c8f518d4 100755 --- a/test/shell.d/agent-usage-agy-limits-test.sh +++ b/test/shell.d/agent-usage-agy-limits-test.sh @@ -143,21 +143,45 @@ reordered=$(test_limits '{"groups": [ 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. +snake=$(test_limits '{"buckets": [ + {"window": "5h", "remaining_fraction": 0.4, "reset_time": "2026-09-15T19:31:46Z"}, + {"window": "weekly", "remainingFraction": "plenty"} +]}') +[[ $(jq -c 'map({label, percent, resetsAt})' <<<"$snake") == '[{"label":"Session (5-hour)","percent":0.6,"resetsAt":"2026-09-15T19:31:46+00:00"}]' ]] || + fail "Antigravity collector reads snake_case buckets listed without groups and skips unreadable fractions" "$snake" +pass "Antigravity collector reads snake_case buckets listed without groups and skips unreadable fractions" + +# The whole record with Google's endpoints and the keyring standing in. The +# keyring answers with keyring.json, Google answers as STATE says, and each +# request's endpoint, token, and host is logged so a test can tell what was +# asked. 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"}}' +cat "$HOME/keyring.json" 2>/dev/null EOF chmod +x "$test_tmp/bin/secret-tool" +in_an_hour=$(python3 -c 'import datetime as dt; print((dt.datetime.now(dt.timezone.utc) + dt.timedelta(hours=1)).isoformat())') +an_hour_ago=$(python3 -c 'import datetime as dt; print((dt.datetime.now(dt.timezone.utc) - dt.timedelta(hours=1)).isoformat())') reset_at=$(python3 -c 'import datetime as dt; print((dt.datetime.now(dt.timezone.utc) + dt.timedelta(hours=3)).isoformat())') +# The id token names the account: header.{"email":"me@example.com"}.signature +me_id_token="header.eyJlbWFpbCI6Im1lQGV4YW1wbGUuY29tIn0.signature" +other_id_token="header.eyJlbWFpbCI6Im90aGVyQGV4YW1wbGUuY29tIn0.signature" + +sign_in() { + local refresh=${4-refresh-main} + printf '{"auth_method":"consumer","id_token":"%s","token":{"access_token":"%s","refresh_token":"%s","expiry":"%s"}}\n' \ + "${3:-$me_id_token}" "$1" "$refresh" "$2" >"$test_tmp/keyring.json" +} + collect() { - HOME="$test_tmp" XDG_CACHE_HOME="$test_tmp/cache" PATH="$test_tmp/bin:$PATH" COLLECTOR="$COLLECTOR" RESET_AT="$reset_at" python3 - "$@" <<'PY' + local state=$1 + shift + HOME="$test_tmp" AGY_DIR="$test_tmp/agy" XDG_CACHE_HOME="$test_tmp/cache" PATH="$test_tmp/bin:$PATH" \ + COLLECTOR="$COLLECTOR" RESET_AT="$reset_at" STATE="$state" ASKED="$test_tmp/asked" python3 - "$@" <<'PY' import importlib.machinery, importlib.util, io, json, os, sys, urllib.error loader = importlib.machinery.SourceFileLoader("collector", os.environ["COLLECTOR"]) @@ -171,35 +195,149 @@ answers = { {"window": "5h", "remainingFraction": 0.75, "resetTime": os.environ["RESET_AT"]} ]}]}, } -state = sys.argv[1] +state = os.environ["STATE"] def urlopen(request, timeout=None): + endpoint = request.full_url.rsplit(":", 1)[1] + with open(os.environ["ASKED"], "a") as asked: + asked.write(f"{endpoint} {request.get_header('Authorization')} {request.full_url.split('/')[2]}\n") 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()) + if state in ("401", "500"): + raise urllib.error.HTTPError(request.full_url, int(state), "No", {}, io.BytesIO()) + if state == "noquota" and endpoint == "retrieveUserQuotaSummary": + return io.BytesIO(b'{"groups": []}') + return io.BytesIO(json.dumps(answers[endpoint]).encode()) collector.urllib.request.urlopen = urlopen -sys.argv = ["omarchy-agent-usage-agy", "--force"] +sys.argv = ["omarchy-agent-usage-agy"] + sys.argv[1:] 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}' ]] || +asked() { + if [[ -e $test_tmp/asked ]]; then + cat "$test_tmp/asked" + rm -f "$test_tmp/asked" + fi +} + +sign_in token-main "$in_an_hour" +record=$(collect up --force) +[[ $(jq -c '{id, ready, tierLabel, usageStatusText, authHelpText, stale: .limitsStale, retryAdvised, percent: .limits[0].percent}' <<<"$record") == '{"id":"agy","ready":true,"tierLabel":"Pro","usageStatusText":"","authHelpText":"","stale":false,"retryAdvised":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" +[[ $(asked | awk '{print $3}' | sort -u) == "token-main" ]] || fail "Antigravity collector asks with agy's token" "$record" pass "Antigravity collector reports its plan and live limits" +caches=$(cd "$test_tmp/cache/omarchy/agent-usage" && ls agy-limits-*.json) +[[ $caches =~ ^agy-limits-[0-9a-f]{16}\.json$ && ! -e $test_tmp/cache/omarchy/agent-usage/agy-limits.json ]] || + fail "Antigravity collector caches limits per account" "$caches" +grep -q "me@example.com" <<<"$caches" && fail "Antigravity collector names its cache by a digest, not the email" "$caches" +pass "Antigravity collector caches limits per account" + +# Within the probe interval a normal run reuses the answer it just got; +# --force asks again. +record=$(collect up) +[[ -z $(asked) && $(jq -c '{stale: .limitsStale, percent: .limits[0].percent}' <<<"$record") == '{"stale":false,"percent":0.25}' ]] || + fail "Antigravity collector reuses a fresh answer without asking Google again" "$record" +record=$(collect up --force) +[[ $(asked | wc -l) == "2" ]] || fail "Antigravity collector asks again when forced" "$record" +pass "Antigravity collector reuses a fresh answer, and asks again when forced" + fetched=$(jq '.limitsFetchedAt' <<<"$record") -record=$(collect down) -[[ $(jq -c '{tierLabel, stale: .limitsStale, percent: .limits[0].percent}' <<<"$record") == '{"tierLabel":"Pro","stale":true,"percent":0.25}' ]] || +record=$(collect down --force) +[[ $(jq -c '{tierLabel, usageStatusText, stale: .limitsStale, retryAdvised, percent: .limits[0].percent}' <<<"$record") == '{"tierLabel":"Pro","usageStatusText":"","stale":true,"retryAdvised":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" +[[ $(jq -r '.authHelpText' <<<"$record") == *"Couldn't reach"* ]] || fail "Antigravity collector says the check couldn't reach Google" "$record" pass "Antigravity collector keeps the last limits, marked stale, after a failed check" +record=$(collect 500 --force) +[[ $(jq -c '{usageStatusText, stale: .limitsStale, retryAdvised}' <<<"$record") == '{"usageStatusText":"","stale":true,"retryAdvised":false}' ]] || + fail "Antigravity collector advises no quick retry when Google answered with an error" "$record" +pass "Antigravity collector advises no quick retry when Google answered with an error" + +record=$(collect 401 --force) +[[ $(jq -c '{usageStatusText, stale: .limitsStale, percent: .limits[0].percent}' <<<"$record") == '{"usageStatusText":"Sign-in expired","stale":true,"percent":0.25}' ]] || + fail "Antigravity collector reports a refused sign-in as expired" "$record" +asked >/dev/null +pass "Antigravity collector reports a refused sign-in as expired" + +# A kept window whose reset has passed is back at zero, not dropped. +cache_file=$(ls "$test_tmp/cache/omarchy/agent-usage"/agy-limits-*.json) +python3 - "$cache_file" <<'PY' +import json, sys +path = sys.argv[1] +data = json.load(open(path)) +data["limits"][0]["resetsAt"] = "2020-01-01T00:00:00+00:00" +data["limits"].append("damaged") +json.dump(data, open(path, "w")) +PY +record=$(collect down --force) +[[ $(jq -c '.limits' <<<"$record") == '[{"label":"Session (5-hour)","title":"Session","percent":0.0,"resetsAt":""}]' ]] || + fail "Antigravity collector shows an elapsed kept window at zero" "$record" +asked >/dev/null +pass "Antigravity collector shows an elapsed kept window at zero" + +# A lapsed access token is renewed the next time agy runs, so the kept limits +# stay up without a new sign-in; without a refresh token behind it, a new +# sign-in is needed. Neither is sent to Google. +sign_in token-main "$an_hour_ago" +record=$(collect up --force) +[[ -z $(asked) && $(jq -c '{usageStatusText, stale: .limitsStale, n: (.limits | length)}' <<<"$record") == '{"usageStatusText":"Limits paused","stale":true,"n":1}' ]] || + fail "Antigravity collector pauses limits while agy's access token has lapsed" "$record" +[[ $(jq -r '.authHelpText' <<<"$record") == *"signing in again is not needed"* ]] || fail "Antigravity collector says no new sign-in is needed" "$record" +pass "Antigravity collector pauses limits while agy's access token has lapsed" + +sign_in token-main "$an_hour_ago" "$me_id_token" "" +record=$(collect up --force) +[[ -z $(asked) && $(jq -r '.usageStatusText' <<<"$record") == "Sign-in expired" ]] || + fail "Antigravity collector reports a lapsed sign-in without a refresh token as expired" "$record" +pass "Antigravity collector reports a lapsed sign-in without a refresh token as expired" + +# Another account signed in here never wears the last one's limits. +sign_in token-other "$an_hour_ago" "$other_id_token" +record=$(collect down --force) +[[ $(jq -c '{limits, tierLabel, limitsFetchedAt}' <<<"$record") == '{"limits":[],"tierLabel":"","limitsFetchedAt":0}' ]] || + fail "Antigravity collector shows no other account's kept limits" "$record" +pass "Antigravity collector shows no other account's kept limits" + +# A plan with no quota windows to report is healthy, with no meters. +sign_in token-other "$in_an_hour" "$other_id_token" +record=$(collect noquota --force) +[[ $(jq -c '{limits, tierLabel, usageStatusText, authHelpText, stale: .limitsStale}' <<<"$record") == '{"limits":[],"tierLabel":"Pro","usageStatusText":"","authHelpText":"","stale":false}' ]] || + fail "Antigravity collector treats a plan without quota windows as healthy" "$record" +asked >/dev/null +pass "Antigravity collector treats a plan without quota windows as healthy" + +# The bearer token goes only to an https destination. +record=$(CLOUD_CODE_URL="http://evil.example" collect up --force) +[[ $(asked | awk '{print $4}' | sort -u) == "daily-cloudcode-pa.googleapis.com" ]] || + fail "Antigravity collector ignores a plain-http destination" "$record" +record=$(CLOUD_CODE_URL="https://mirror.example" collect up --force) +[[ $(asked | awk '{print $4}' | sort -u) == "mirror.example" ]] || + fail "Antigravity collector honors an https destination" "$record" +pass "Antigravity collector asks only an https destination, defaulting to Google's" + rm -rf "$test_tmp/cache" -record=$(collect down) -[[ $(jq -c '{limits, usageStatusText}' <<<"$record") == '{"limits":[],"usageStatusText":"Antigravity limits unavailable"}' ]] || +record=$(collect down --force) +[[ $(jq -c '{limits, usageStatusText, retryAdvised}' <<<"$record") == '{"limits":[],"usageStatusText":"Antigravity limits unavailable","retryAdvised":true}' ]] || fail "Antigravity collector invents no limits when it has never reached Google" "$record" +asked >/dev/null pass "Antigravity collector invents no limits when it has never reached Google" + +# A normal run reuses a scan for a few seconds; --force rescans. +mkdir -p "$test_tmp/agy" +now=$(date -u +%Y-%m-%dT%H:%M:%SZ) +echo '{"timestamp":"'"$now"'","conversation_id":"c1"}' >"$test_tmp/agy/history.jsonl" +record=$(collect up --force) +echo '{"timestamp":"'"$now"'","conversation_id":"c2"}' >>"$test_tmp/agy/history.jsonl" +record=$(collect up) +[[ $(jq '.totalPrompts' <<<"$record") == "1" ]] || fail "Antigravity collector reuses a recent scan" "$record" +record=$(collect up --limits-only) +[[ $(jq '.totalPrompts' <<<"$record") == "1" ]] || fail "Antigravity collector reuses a scan on a limits-only run" "$record" +record=$(collect up --force) +[[ $(jq '.totalPrompts' <<<"$record") == "2" ]] || fail "Antigravity collector rescans when forced" "$record" +asked >/dev/null +pass "Antigravity collector reuses a recent scan, and rescans when forced" diff --git a/test/shell.d/agent-usage-agy-timestamps-test.sh b/test/shell.d/agent-usage-agy-timestamps-test.sh index 71e335a886c..0f68b9ea5c0 100755 --- a/test/shell.d/agent-usage-agy-timestamps-test.sh +++ b/test/shell.d/agent-usage-agy-timestamps-test.sh @@ -198,7 +198,9 @@ class TimestampTests(unittest.TestCase): record = json.loads(result.stdout) self.assertEqual(result.stderr, '') self.assertEqual(record['id'], 'agy') - self.assertFalse(record['ready']) + # Local history alone puts Antigravity in use here, waiting for a sign-in. + self.assertTrue(record['ready']) + self.assertEqual(record['usageStatusText'], 'Waiting for auth') self.assertEqual(record['limits'], []) self.assertEqual(record['totalPrompts'], 3) self.assertEqual(record['totalSessions'], 4) From f9bb0b51ce2e87645c0b677a56b8b21f9aba3d88 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:07:45 -0400 Subject: [PATCH 056/546] Read the Antigravity sign-in omp, pi, or opencode keeps when agy has none Try the latest-expiring token first and move on from one Google refuses, fall back to agy's own token file without a keyring, and ask a lapsed omp sign-in through `omp usage`, which renews it. Every store is only read. Co-authored-by: ivanvan08 Co-authored-by: steph4n <284007938+steph4n-gh@users.noreply.github.com> Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-agy | 178 +++++++++++++-- test/shell.d/agent-usage-agy-harness-test.sh | 214 +++++++++++++++++++ 2 files changed, 371 insertions(+), 21 deletions(-) create mode 100755 test/shell.d/agent-usage-agy-harness-test.sh diff --git a/bin/omarchy-agent-usage-agy b/bin/omarchy-agent-usage-agy index eb5345e2484..239e42bfcfb 100755 --- a/bin/omarchy-agent-usage-agy +++ b/bin/omarchy-agent-usage-agy @@ -11,7 +11,10 @@ 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. +Secret Service keyring (or its antigravity-oauth-token file), or with the +Antigravity sign-in omp, pi, or opencode's Antigravity plugin keeps; a lapsed +omp sign-in is asked through `omp usage`, which renews it. Those stores are +only ever read. The agents panel only ever reads the JSON this prints. """ import argparse @@ -22,6 +25,7 @@ import http.client import json import os import re +import shutil import sqlite3 import subprocess import tempfile @@ -33,9 +37,9 @@ from pathlib import Path AGENT_ID = "agy" AGENT_NAME = "Antigravity" -AUTH_HELP = "Run `agy` to sign in." +AUTH_HELP = "Run `agy` to sign in, or sign in to Antigravity in omp, pi, or opencode." # How each source of a sign-in is named when it needs attention. -SOURCE_NAMES = {"agy": "`agy`"} +SOURCE_NAMES = {"agy": "`agy`", "omp": "omp", "pi": "pi", "opencode": "opencode"} PROBE_MIN_INTERVAL_SECONDS = 15 SCAN_REUSE_SECONDS = 20 # A limits-only refresh needs fresh limits, not a fresh history scan. @@ -110,6 +114,7 @@ def local_day(value): # The sign-in agy keeps in the keyring: its access token, when that lapses, # whether a refresh token stands behind it, and which kind of account it is. +# Without a Secret Service, agy keeps the same sign-in in a file of its own. def login(): try: raw = subprocess.check_output( @@ -119,6 +124,8 @@ def login(): data = json.loads(raw) if raw.strip() else None except Exception: data = None + if not isinstance(data, dict): + data = read_json(app_root() / "antigravity-oauth-token") if not isinstance(data, dict): return None @@ -134,9 +141,76 @@ def login(): } +# omp, pi, and opencode's Antigravity plugin sign in with the same Google +# client agy uses, so the access tokens they keep reach the same plan and +# quota. Each refreshes its own token while it runs; these are only ever +# read, never refreshed or written back. +def harness_login(source, data): + if not isinstance(data, dict) or data.get("type", "oauth") != "oauth" or not isinstance(data.get("access"), str) or not data["access"]: + return None + expires = parse_time(data.get("expires")) + return { + "source": source, + "accessToken": data["access"], + "expiresAt": expires.timestamp() if expires else 0, + "refresh": str(data.get("refresh") or ""), + "authMethod": "consumer", + "email": str(data.get("email") or ""), + } + + +def omp_logins(): + database = Path.home() / ".omp" / "agent" / "agent.db" + if not database.is_file(): + return [] + try: + conn = sqlite3.connect(database.resolve().as_uri() + "?mode=ro", uri=True, timeout=2) + try: + rows = conn.execute( + "SELECT data FROM auth_credentials WHERE provider = 'google-antigravity' " + "AND credential_type = 'oauth' AND disabled_cause IS NULL" + ).fetchall() + finally: + conn.close() + except sqlite3.Error: + return [] + logins = [] + for (data,) in rows: + try: + entry = harness_login("omp", json.loads(data)) + except (TypeError, ValueError): + continue + if entry: + logins.append(entry) + return logins + + +def pi_agent_dir(): + return Path(os.environ.get("PI_CODING_AGENT_DIR") or (Path.home() / ".pi" / "agent")) + + +def pi_logins(): + auth = read_json(pi_agent_dir() / "auth.json") + entry = harness_login("pi", auth.get("google-antigravity")) if isinstance(auth, dict) else None + return [entry] if entry else [] + + +# opencode files the Antigravity plugin's sign-in under its "google" +# provider, which other Google sign-ins share, so it only counts alongside +# the plugin's own accounts file. +def opencode_logins(): + config_home = Path(os.environ.get("XDG_CONFIG_HOME") or (Path.home() / ".config")) + if not (config_home / "opencode" / "antigravity-accounts.json").is_file(): + return [] + data_home = Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) + auth = read_json(data_home / "opencode" / "auth.json") + entry = harness_login("opencode", auth.get("google")) if isinstance(auth, dict) else None + return [entry] if entry else [] + + def logins(): entry = login() - return [entry] if entry else [] + return ([entry] if entry else []) + omp_logins() + pi_logins() + opencode_logins() def email_of(id_token): @@ -313,27 +387,31 @@ def collect_limits(entries, force): ordered = sorted(entries, key=lambda e: (not token_fresh(e, now), -e["expiresAt"])) fresh = [e for e in ordered if token_fresh(e, now)] lead = next(iter(fresh or [e for e in ordered if e["accessToken"]] or ordered), None) + omp = shutil.which("omp") if any(e["source"] == "omp" for e in entries) else None + omp_entries = [e for e in ordered if e["source"] == "omp"] if omp else [] def keep(entry): cached = cached_answer(entry) if entry else {} result.update(limits=usable_cached_limits(cached), tierLabel=str(cached.get("tierLabel") or ""), fetchedAtMs=fetched_at_ms(cached)) + def answered(entry, tier, limits): + # A plan with no quota windows to report is healthy, with no meters. + fetched = round(time.time() * 1000) + result.update(limits=limits, tierLabel=tier, live=True, fetchedAtMs=fetched) + path = limits_cache(entry) + if path: + write_json(path, {"fetchedAtMs": fetched, "limits": limits, "tierLabel": tier}) + return result + keep(lead) if not lead or not lead["accessToken"]: result.update(usageStatusText="Waiting for auth", authHelpText=AUTH_HELP) return result - # agy refreshes its token while it runs; one left to lapse can't be used - # until agy starts again, which needs no new sign-in while a refresh token - # stands behind it. - if not fresh: - result.update(usageStatusText="Limits paused" if lead["refresh"] else "Sign-in expired", authHelpText=lapsed_help(lead, result["limits"])) - return result - # A panel that is opened and shut repeatedly must not turn into a request # per flick; --force is a person asking for fresh numbers. if not force: - for entry in fresh: + for entry in fresh + [e for e in omp_entries if e not in fresh]: cached = cached_answer(entry) if cached and 0 <= now - fetched_at_ms(cached) / 1000 < PROBE_MIN_INTERVAL_SECONDS: keep(entry) @@ -359,13 +437,7 @@ def collect_limits(entries, force): except ValueError: failure = entry, "Google's quota service sent an answer that couldn't be read.", False break - # A plan with no quota windows to report is healthy, with no meters. - fetched = round(time.time() * 1000) - result.update(limits=limits, tierLabel=tier, live=True, fetchedAtMs=fetched) - path = limits_cache(entry) - if path: - write_json(path, {"fetchedAtMs": fetched, "limits": limits, "tierLabel": tier}) - return result + return answered(entry, tier, limits) if failure: entry, help_text, transport = failure @@ -375,11 +447,75 @@ def collect_limits(entries, force): result["usageStatusText"] = "Antigravity limits unavailable" return result - keep(refused[0]) - result.update(usageStatusText="Sign-in expired", authHelpText=refused_help(refused[0])) + # No token in hand answered, but omp renews its own sign-in when asked + # for usage. + answer = omp_answer(omp, omp_entries) if omp_entries else None + if answer: + entry, tier, limits = answer + return answered(entry, tier or str(cached_answer(entry).get("tierLabel") or ""), limits) + + if refused: + keep(refused[0]) + result.update(usageStatusText="Sign-in expired", authHelpText=refused_help(refused[0])) + return result + + # Every harness refreshes its token while it runs; one left to lapse can't + # be used until that harness starts again, which needs no new sign-in + # while a refresh token stands behind it. + result.update(usageStatusText="Limits paused" if lead["refresh"] else "Sign-in expired", authHelpText=lapsed_help(lead, result["limits"])) return result +# omp refreshes its own Antigravity sign-in when asked for usage, so its +# limits come from `omp usage` even after the token it keeps has lapsed. The +# report is folded back into the shape Google's quota summary has. +def omp_usage(omp): + try: + output = subprocess.run( + [omp, "usage", "--json", "--provider", "google-antigravity"], + capture_output=True, text=True, timeout=20, check=True, + ).stdout + reports = [r for r in json.loads(output).get("reports") or [] if isinstance(r, dict) and r.get("provider") == "google-antigravity"] + except (OSError, ValueError, AttributeError, subprocess.SubprocessError): + return None + if not reports: + return None + groups, seen = {}, set() + for limit in reports[0].get("limits") or []: + if not isinstance(limit, dict): + continue + window = limit.get("window") if isinstance(limit.get("window"), dict) else {} + amount = limit.get("amount") if isinstance(limit.get("amount"), dict) else {} + label = str(limit.get("label") or "") + # A shared allowance is listed once per model family that draws on it. + if (label, window.get("id")) in seen: + continue + seen.add((label, window.get("id"))) + fraction = amount.get("usedFraction") + if isinstance(fraction, bool) or not isinstance(fraction, (int, float)): + continue + groups.setdefault(label, []).append({"window": str(window.get("id") or ""), "remainingFraction": 1.0 - fraction, "resetTime": window.get("resetsAt")}) + metadata = reports[0].get("metadata") if isinstance(reports[0].get("metadata"), dict) else {} + quota = {"groups": [{"displayName": name, "buckets": buckets} for name, buckets in groups.items()]} + return quota, str(metadata.get("email") or "") + + +# The sign-in `omp usage` answered for, its plan when the token omp just +# renewed can still ask Google for it, and its limits. +def omp_answer(omp, entries): + answer = omp_usage(omp) + if not answer: + return None + quota, email = answer + renewed = [e for e in omp_logins() if not email or e["email"] == email] + entry = (renewed or [e for e in entries if not email or e["email"] == email] or [dict(entries[0], email=email)])[0] + tier = "" + if token_fresh(entry, time.time()): + try: + tier = extract_tier(post("loadCodeAssist", entry["accessToken"])) + except (urllib.error.URLError, OSError, http.client.HTTPException, ValueError): + pass + return entry, tier, extract_limits(quota) def model_bucket(models, model): return models.setdefault(model, {"inputTokens": 0, "outputTokens": 0, "cacheReadInputTokens": 0, "cacheCreationInputTokens": 0}) diff --git a/test/shell.d/agent-usage-agy-harness-test.sh b/test/shell.d/agent-usage-agy-harness-test.sh new file mode 100755 index 00000000000..4af6be048ed --- /dev/null +++ b/test/shell.d/agent-usage-agy-harness-test.sh @@ -0,0 +1,214 @@ +#!/bin/bash + +source "$(dirname "$0")/base-test.sh" + +require_command jq +require_command python3 + +COLLECTOR="$ROOT/bin/omarchy-agent-usage-agy" + +# Antigravity signed in through another harness, with agy itself absent: the +# record still carries the plan and limits, asked with that harness' token. +test_tmp=$(mktemp -d) +trap 'rm -rf "$test_tmp"' EXIT +mkdir -p "$test_tmp/bin" +cat >"$test_tmp/bin/secret-tool" <<'EOF' +#!/bin/bash +exit 1 +EOF +# omp answers `omp usage` with the report in omp-report.json, when there is +# one, and fails otherwise. +cat >"$test_tmp/bin/omp" <<'EOF' +#!/bin/bash +[[ "$*" == "usage --json --provider google-antigravity" ]] || exit 2 +cat "$HOME/omp-report.json" 2>/dev/null +EOF +chmod +x "$test_tmp/bin/secret-tool" "$test_tmp/bin/omp" + +in_ms() { + python3 -c "import sys, time; print(round((time.time() + float(sys.argv[1])) * 1000))" "$1" +} + +sqlite() { + python3 - "$@" <<'PY' +import sqlite3, sys +db = sqlite3.connect(sys.argv[1]) +db.executescript(sys.argv[2]) +db.commit() +PY +} + +omp_login() { + mkdir -p "$test_tmp/.omp/agent" + rm -f "$test_tmp/.omp/agent/agent.db" + sqlite "$test_tmp/.omp/agent/agent.db" \ + "CREATE TABLE auth_credentials (id INTEGER PRIMARY KEY, provider TEXT, credential_type TEXT, data TEXT, disabled_cause TEXT); + INSERT INTO auth_credentials (provider, credential_type, data) VALUES + ('google-antigravity', 'oauth', '{\"access\":\"$1\",\"refresh\":\"r\",\"expires\":$2,\"email\":\"$3\"}'); + INSERT INTO auth_credentials (provider, credential_type, data, disabled_cause) VALUES + ('google-antigravity', 'oauth', '{\"access\":\"disabled-token\",\"refresh\":\"r\",\"expires\":$(in_ms 99999),\"email\":\"$3\"}', 'revoked');" +} + +pi_login() { + mkdir -p "${PI_DIR:-$test_tmp/.pi/agent}" + printf '{"google-antigravity":{"type":"oauth","access":"%s","refresh":"r","expires":%s,"email":"%s"}}\n' "$1" "$2" "$3" \ + >"${PI_DIR:-$test_tmp/.pi/agent}/auth.json" +} + +# Google answers every token in ACCEPTED and refuses the rest; each request's +# token is logged so a test can tell which sign-in was asked. +collect() { + env -u OPENCLAW_STATE_DIR -u XDG_CONFIG_HOME -u XDG_DATA_HOME -u CLOUD_CODE_URL \ + HOME="$test_tmp" XDG_CACHE_HOME="$test_tmp/cache" PATH="$test_tmp/bin:$PATH" \ + PI_CODING_AGENT_DIR="${PI_DIR:-}" COLLECTOR="$COLLECTOR" ACCEPTED="$1" ASKED="$test_tmp/asked" 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": "2999-01-01T00:00:00Z"} + ]}]}, +} + +def urlopen(request, timeout=None): + token = request.get_header("Authorization").removeprefix("Bearer ") + with open(os.environ["ASKED"], "a") as asked: + asked.write(token + "\n") + if token not in os.environ["ACCEPTED"].split(): + raise urllib.error.HTTPError(request.full_url, 403, "Forbidden", {}, io.BytesIO()) + 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 +} + +asked() { + if [[ -e $test_tmp/asked ]]; then + sort -u "$test_tmp/asked" | paste -sd' ' + rm -f "$test_tmp/asked" + fi +} + +# Every store the collector reads, so a test can tell none was written. +stores() { + find "$test_tmp/.omp" "$test_tmp/.pi" "$test_tmp/.local/share/opencode" "$test_tmp/.gemini" "$test_tmp/pi-elsewhere" \ + -type f -exec md5sum {} + 2>/dev/null | sort +} + +omp_login omp-token "$(in_ms 3600)" me@example.com +record=$(collect omp-token) +[[ $(jq -c '{ready, tierLabel, usageStatusText, authHelpText, stale: .limitsStale, percent: .limits[0].percent}' <<<"$record") == '{"ready":true,"tierLabel":"Pro","usageStatusText":"","authHelpText":"","stale":false,"percent":0.25}' ]] || + fail "Antigravity collector reports limits from an omp sign-in without agy" "$record" +[[ $(asked) == "omp-token" ]] || fail "Antigravity collector asks Google with omp's token, never a disabled one" "$record" +pass "Antigravity collector reports limits from an omp sign-in without agy" + +pi_login pi-token "$(in_ms 7200)" me@example.com +record=$(collect "omp-token pi-token") +[[ $(asked) == "pi-token" ]] || + fail "Antigravity collector asks with the most recently refreshed sign-in" "$record" +pass "Antigravity collector asks with the most recently refreshed sign-in" + +record=$(collect omp-token) +[[ $(asked) == "omp-token pi-token" && $(jq -c '{stale: .limitsStale, usageStatusText}' <<<"$record") == '{"stale":false,"usageStatusText":""}' ]] || + fail "Antigravity collector moves on from a sign-in Google refuses" "$record" +pass "Antigravity collector moves on from a sign-in Google refuses" + +record=$(collect "") +[[ $(jq -c '{usageStatusText, stale: .limitsStale, percent: .limits[0].percent}' <<<"$record") == '{"usageStatusText":"Sign-in expired","stale":true,"percent":0.25}' ]] || + fail "Antigravity collector reports every sign-in refused as expired, keeping the account's limits" "$record" +[[ $(jq -r '.authHelpText' <<<"$record") == "Google refused the Antigravity sign-in pi keeps. Start pi to sign in again." ]] || + fail "Antigravity collector names the harness whose sign-in was refused" "$record" +asked >/dev/null +pass "Antigravity collector reports every sign-in refused as expired, keeping the account's limits" + +# A lapsed token is renewed the next time its harness runs, so the limits are +# paused, not signed out, and the token goes nowhere. +rm -rf "$test_tmp/.pi" +omp_login omp-token "$(in_ms -60)" me@example.com +before=$(stores) +record=$(collect omp-token) +[[ -z $(asked) ]] || fail "Antigravity collector sends no lapsed token to Google" "$record" +[[ $(jq -c '{ready, stale: .limitsStale, usageStatusText, percent: .limits[0].percent}' <<<"$record") == '{"ready":true,"stale":true,"usageStatusText":"Limits paused","percent":0.25}' ]] || + fail "Antigravity collector keeps the last limits while omp's token has lapsed" "$record" +[[ $(jq -r '.authHelpText' <<<"$record") == *"Start omp to refresh it; signing in again is not needed." ]] || + fail "Antigravity collector names the harness whose token lapsed" "$record" +[[ $(stores) == "$before" ]] || fail "Antigravity collector never writes omp's store" +pass "Antigravity collector keeps the last limits while omp's token has lapsed" + +omp_login other-token "$(in_ms -60)" other@example.com +record=$(collect "") +[[ $(jq -c '{limits, tierLabel}' <<<"$record") == '{"limits":[],"tierLabel":""}' ]] || + fail "Antigravity collector shows no other Google account's kept limits" "$record" +pass "Antigravity collector shows no other Google account's kept limits" + +# omp renews its own sign-in when asked for usage, so a lapsed token in its +# database still gets live limits through `omp usage`. +cat >"$test_tmp/omp-report.json" <<'EOF' +{"reports":[{"provider":"google-antigravity","metadata":{"email":"other@example.com"},"limits":[ + {"label":"Gemini","amount":{"usedFraction":0.6},"window":{"id":"5h","resetsAt":32503680000000}}, + {"label":"Gemini","amount":{"usedFraction":0.3},"window":{"id":"weekly","resetsAt":32503680000000}}, + {"label":"Claude & GPT (shared)","amount":{"usedFraction":0.1},"window":{"id":"weekly","resetsAt":32503680000000}}, + {"label":"Claude & GPT (shared)","amount":{"usedFraction":0.1},"window":{"id":"weekly","resetsAt":32503680000000}}, + {"label":"Gemini","amount":{"usedFraction":"lots"},"window":{"id":"daily"}} +]}]} +EOF +record=$(collect "") +[[ -z $(asked) ]] || fail "Antigravity collector sends omp's lapsed token nowhere" "$record" +[[ $(jq -c '{stale: .limitsStale, usageStatusText, authHelpText, limits: [.limits[] | {title, percent}]}' <<<"$record") == '{"stale":false,"usageStatusText":"","authHelpText":"","limits":[{"title":"Session","percent":0.6},{"title":"Weekly","percent":0.3},{"title":"Claude/GPT Weekly","percent":0.1}]}' ]] || + fail "Antigravity collector takes omp's limits from omp usage" "$record" +rm "$test_tmp/omp-report.json" +pass "Antigravity collector takes omp's limits from omp usage" +rm -rf "$test_tmp/.omp" "$test_tmp/cache" + +# opencode files any Google sign-in under "google"; only the Antigravity +# plugin's accounts file makes it an Antigravity one. +mkdir -p "$test_tmp/.local/share/opencode" +printf '{"google":{"type":"oauth","access":"opencode-token","refresh":"r|p","expires":%s}}\n' "$(in_ms 3600)" \ + >"$test_tmp/.local/share/opencode/auth.json" +record=$(collect opencode-token) +[[ $(jq -c '{ready, limits}' <<<"$record") == '{"ready":false,"limits":[]}' && -z $(asked) ]] || + fail "Antigravity collector ignores opencode's Google sign-in without the Antigravity plugin" "$record" +mkdir -p "$test_tmp/.config/opencode" +echo '{"version":3,"accounts":[]}' >"$test_tmp/.config/opencode/antigravity-accounts.json" +record=$(collect opencode-token) +[[ $(asked) == "opencode-token" && $(jq '.limits[0].percent' <<<"$record") == "0.25" ]] || + fail "Antigravity collector reads opencode's Antigravity plugin sign-in" "$record" +pass "Antigravity collector reads opencode's Antigravity plugin sign-in only with the plugin" + +# opencode records no email for its sign-in, so a replaced sign-in is told +# apart by its refresh token: the kept limits go rather than pass for the +# new account's. +printf '{"google":{"type":"oauth","access":"switched-token","refresh":"switched|p","expires":%s}}\n' "$(in_ms 3600)" \ + >"$test_tmp/.local/share/opencode/auth.json" +record=$(collect "") +[[ $(jq -c '{limits, tierLabel, usageStatusText}' <<<"$record") == '{"limits":[],"tierLabel":"","usageStatusText":"Sign-in expired"}' ]] || + fail "Antigravity collector drops kept limits when the account cannot be matched" "$record" +asked >/dev/null +pass "Antigravity collector drops kept limits when the account cannot be matched" +rm -rf "$test_tmp/.local/share/opencode" "$test_tmp/.config/opencode" + +# Without a keyring, agy keeps its sign-in in a file under its home. +mkdir -p "$test_tmp/.gemini/antigravity-cli" +printf '{"token":{"access_token":"agy-file-token","refresh_token":"agy-refresh","expiry":"2999-01-01T00:00:00Z"}}\n' \ + >"$test_tmp/.gemini/antigravity-cli/antigravity-oauth-token" +record=$(collect agy-file-token) +[[ $(asked) == "agy-file-token" && $(jq '.limits[0].percent' <<<"$record") == "0.25" ]] || + fail "Antigravity collector reads agy's sign-in file when there is no keyring" "$record" +pass "Antigravity collector reads agy's sign-in file when there is no keyring" +rm -rf "$test_tmp/.gemini" + +# PI_CODING_AGENT_DIR moves pi's sign-in. +PI_DIR="$test_tmp/pi-elsewhere" pi_login moved-token "$(in_ms 3600)" me@example.com +before=$(stores) +record=$(PI_DIR="$test_tmp/pi-elsewhere" collect moved-token) +[[ $(asked) == "moved-token" && $(jq '.limits[0].percent' <<<"$record") == "0.25" ]] || + fail "Antigravity collector reads pi's sign-in where PI_CODING_AGENT_DIR puts it" "$record" +[[ $(stores) == "$before" ]] || fail "Antigravity collector never writes pi's store" +pass "Antigravity collector reads pi's sign-in where PI_CODING_AGENT_DIR puts it" From 0b05b75bb04f94a05447b9db5749018345051100 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:09:30 -0400 Subject: [PATCH 057/546] Count Antigravity used through pi, omp, OpenClaw, and opencode in its local stats Find google-antigravity answers in pi and omp sessions (with profiles and PI_CODING_AGENT_DIR) and OpenClaw's, counting a forked copy once, and opencode's antigravity- models under a google provider in both of its message tables. No other collector counts these. Co-authored-by: ivanvan08 Co-authored-by: steph4n <284007938+steph4n-gh@users.noreply.github.com> Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-agy | 164 ++++++++++++++++++- test/shell.d/agent-usage-agy-harness-test.sh | 46 ++++++ 2 files changed, 207 insertions(+), 3 deletions(-) diff --git a/bin/omarchy-agent-usage-agy b/bin/omarchy-agent-usage-agy index 239e42bfcfb..9592b5061b2 100755 --- a/bin/omarchy-agent-usage-agy +++ b/bin/omarchy-agent-usage-agy @@ -8,7 +8,9 @@ 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 +conversation transcripts under ~/.gemini/antigravity-cli/, and from the +Antigravity answers pi, omp, OpenClaw, and opencode's Antigravity plugin +recorded; no other collector counts those. 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 (or its antigravity-oauth-token file), or with the @@ -617,10 +619,146 @@ def conversation_summaries(root, problems): return [] +# pi and omp keep sessions under /agent/sessions, and `--profile=` +# moves the whole agent tree under /profiles//; PI_CODING_AGENT_DIR +# moves pi's. OpenClaw keeps the same transcript format, one directory per +# agent. Sorted, so the scan order doesn't depend on directory order. +def session_roots(): + roots = [pi_agent_dir() / "sessions", Path.home() / ".omp" / "agent" / "sessions"] + for base in (Path.home() / ".pi", Path.home() / ".omp"): + try: + roots.extend(sorted(child / "agent" / "sessions" for child in (base / "profiles").iterdir() if child.is_dir())) + except OSError: + pass + openclaw = Path(os.environ.get("OPENCLAW_STATE_DIR") or (Path.home() / ".openclaw")) + try: + roots.extend(sorted(child / "sessions" for child in (openclaw / "agents").iterdir() if child.is_dir())) + except OSError: + pass + return list(dict.fromkeys(roots)) + + +def opencode_database(): + return Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) / "opencode" / "opencode.db" + + +# Each answer pi, omp, or OpenClaw got through its google-antigravity +# provider: the session it belongs to, its day, its model, and its input, +# output, cache read, and cache write tokens. +def session_messages(problems): + roots = [root for root in session_roots() if root.is_dir()] + rg = shutil.which("rg") if roots else None + if roots and not rg: + problems.append("rg") + return + seen = set() + for root in roots: + try: + proc = subprocess.Popen( + # --no-ignore: these are data files, not a source tree. A home that + # is itself a git checkout would otherwise hide every session file + # behind its ignore rules. + [rg, "--json", "--no-ignore", "-e", r'"provider"\s*:\s*"google-antigravity"', str(root)], + stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, text=True, errors="replace", + ) + except OSError: + problems.append(str(root)) + continue + for raw in proc.stdout: + try: + event = json.loads(raw) + if event.get("type") != "match": + continue + path = event["data"]["path"]["text"] + entry = json.loads(event["data"]["lines"]["text"]) + except (ValueError, KeyError, TypeError, AttributeError): + continue + if not isinstance(entry, dict) or entry.get("type") != "message": + continue + message = entry.get("message") if isinstance(entry.get("message"), dict) else {} + if message.get("role") != "assistant" or message.get("provider") != "google-antigravity": + continue + # A fork copies its parent's entries, and they keep their id and time, + # so each answer counts once. An entry without an id always counts. + if entry.get("id"): + key = (str(entry["id"]), str(entry["timestamp"])) if entry.get("timestamp") else (path, str(entry["id"])) + if key in seen: + continue + seen.add(key) + usage = message.get("usage") if isinstance(message.get("usage"), dict) else {} + tokens = (number(usage.get("input")), number(usage.get("output")), number(usage.get("cacheRead")), number(usage.get("cacheWrite"))) + if not any(tokens): + tokens = (number(usage.get("totalTokens")), 0, 0, 0) + model = message.get("model") if isinstance(message.get("model"), str) else "" + if any(tokens): + yield path, local_day(entry.get("timestamp") or message.get("timestamp")), model_key(model) or "unknown", tokens + try: + proc.wait(timeout=1) + except subprocess.TimeoutExpired: + proc.kill() + + +# opencode's Antigravity plugin runs under a "google" provider, with model +# ids it prefixes with antigravity- to keep them apart from Gemini API ones. +# opencode v1 keeps messages in `message`, with the role, provider and model +# in each row's data; v2 writes `session_message`, with the role in its type +# column and the provider and model under data.model. Read-only: opencode +# may be writing right now. +def opencode_messages(problems): + database = opencode_database() + if not database.is_file(): + return + try: + conn = sqlite3.connect(database.resolve().as_uri() + "?mode=ro", uri=True, timeout=2) + except sqlite3.Error: + problems.append("opencode") + return + try: + tables = {name for (name,) in conn.execute("SELECT name FROM sqlite_master WHERE type = 'table'")} + rows = [] + if "message" in tables: + rows += [("opencode:", session, raw, None) for session, raw in conn.execute("SELECT session_id, data FROM message WHERE data LIKE '%antigravity-%'")] + if "session_message" in tables: + rows += [("opencode-v2:", session, raw, kind) for session, raw, kind in conn.execute("SELECT session_id, data, type FROM session_message WHERE data LIKE '%antigravity-%'")] + except sqlite3.Error: + problems.append("opencode") + return + finally: + conn.close() + for prefix, session, raw, kind in rows: + try: + entry = json.loads(raw) + if kind is None: + role, provider, model = entry.get("role"), entry.get("providerID"), entry.get("modelID") + else: + info = entry.get("model") if isinstance(entry.get("model"), dict) else {} + role, provider, model = kind, info.get("providerID"), info.get("id") + tokens = entry.get("tokens") if isinstance(entry.get("tokens"), dict) else {} + cache = tokens.get("cache") if isinstance(tokens.get("cache"), dict) else {} + created = (entry.get("time") or {}).get("created") + except (ValueError, TypeError, AttributeError): + continue + if role != "assistant" or not str(provider or "").startswith("google") or not str(model or "").startswith("antigravity-"): + continue + # opencode keeps thinking tokens out of output; both are generated. + counts = (number(tokens.get("input")), number(tokens.get("output")) + number(tokens.get("reasoning")), number(cache.get("read")), number(cache.get("write"))) + if any(counts): + yield prefix + str(session), local_day(created), model_key(str(model)), counts + + +# Hermes reaches Antigravity only through plugins that run agy itself, so its +# sign-in and transcripts are agy's own, which local_stats reads already. +def harness_messages(problems): + yield from session_messages(problems) + yield from opencode_messages(problems) + + # 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; steps without a usable -# date still count toward the lifetime totals. +# date still count toward the lifetime totals. Antigravity driven from +# another harness adds each answer it recorded as a prompt, the way the +# Codex collector counts them, with every token it reports. def local_stats(root, problems=None): problems = [] if problems is None else problems today = datetime.now().astimezone().date() @@ -699,6 +837,25 @@ def local_stats(root, problems=None): today_tokens += total today_by_model[model] = today_by_model.get(model, 0) + total + for session, day, model, counts in harness_messages(problems): + total = sum(counts) + prompts += 1 + sessions.add(session) + bucket = model_bucket(models, model) + bucket["inputTokens"] += counts[0] + bucket["outputTokens"] += counts[1] + bucket["cacheReadInputTokens"] += counts[2] + bucket["cacheCreationInputTokens"] += counts[3] + if day: + days.add(day) + if day in by_day: + by_day[day] += total + if day == today: + today_prompts += 1 + today_sessions.add(session) + today_tokens += total + today_by_model[model] = today_by_model.get(model, 0) + total + return { "hasLocalStats": prompts > 0 or bool(sessions) or bool(models), "todayPrompts": today_prompts, @@ -719,7 +876,8 @@ def local_stats(root, problems=None): # is rebuilt under a lock so overlapping runs don't each walk the history, # and one that couldn't read every source isn't kept. def cached_local_stats(root, max_age): - digest = hashlib.sha256(str(root).encode()).hexdigest()[:16] + sources = [str(root), str(opencode_database())] + [str(path) for path in session_roots()] + digest = hashlib.sha256("\n".join(sources).encode()).hexdigest()[:16] cache = cache_root() / f"agy-scan-{digest}.json" today = datetime.now().astimezone().date().isoformat() diff --git a/test/shell.d/agent-usage-agy-harness-test.sh b/test/shell.d/agent-usage-agy-harness-test.sh index 4af6be048ed..26f83b9dfe5 100755 --- a/test/shell.d/agent-usage-agy-harness-test.sh +++ b/test/shell.d/agent-usage-agy-harness-test.sh @@ -212,3 +212,49 @@ record=$(PI_DIR="$test_tmp/pi-elsewhere" collect moved-token) fail "Antigravity collector reads pi's sign-in where PI_CODING_AGENT_DIR puts it" "$record" [[ $(stores) == "$before" ]] || fail "Antigravity collector never writes pi's store" pass "Antigravity collector reads pi's sign-in where PI_CODING_AGENT_DIR puts it" + +# Antigravity used through pi, omp (and its profiles), OpenClaw, and +# opencode's plugin counts in the local stats; Gemini through any other +# provider does not, and a forked session's copy of an answer counts once. +if command -v rg >/dev/null; then + rm -rf "$test_tmp/.pi" "$test_tmp/.omp" "$test_tmp/pi-elsewhere" "$test_tmp/cache" + now=$(date -u +%Y-%m-%dT%H:%M:%S.000Z) + # The home is a git checkout that ignores everything, as dotfiles setups do. + printf '*\n' >"$test_tmp/.gitignore" + mkdir -p "$test_tmp/.git" "$test_tmp/.omp/agent/sessions/project" "$test_tmp/.omp/profiles/work/agent/sessions" \ + "$test_tmp/.openclaw/agents/main/sessions" "$test_tmp/pi-elsewhere/sessions" "$test_tmp/.local/share/opencode" + answer='{"type":"message","id":"a1","timestamp":"'$now'","message":{"role":"assistant","provider":"google-antigravity","model":"gemini-3-pro","usage":{"input":100,"output":20,"cacheRead":5,"cacheWrite":0}}}' + printf '%s\n%s\n' "$answer" \ + '{"type":"message","id":"a2","timestamp":"'$now'","message":{"role":"assistant","provider":"google-gemini-cli","model":"gemini-3-pro","usage":{"input":999,"output":999}}}' \ + >"$test_tmp/.omp/agent/sessions/project/one.jsonl" + printf '%s\n' "$answer" >"$test_tmp/.omp/agent/sessions/project/fork.jsonl" + # A profile's sessions count; an answer with an unreadable time and a + # total alone still counts toward the lifetime totals. + printf '%s\n' '{"type":"message","id":"p1","timestamp":"0001-01-01T00:00:00Z","message":{"role":"assistant","provider":"google-antigravity","model":"gemini-3-pro","usage":{"input":"many","totalTokens":30}}}' \ + >"$test_tmp/.omp/profiles/work/agent/sessions/s.jsonl" + printf '%s\n' '{"type":"message","id":"m1","timestamp":"'$now'","message":{"role":"assistant","provider":"google-antigravity","model":"gemini-3-pro","usage":{"input":3,"output":1}}}' \ + >"$test_tmp/pi-elsewhere/sessions/s.jsonl" + # OpenClaw's entries carry no id; two alike answers are two answers. + openclaw='{"type":"message","timestamp":"'$now'","message":{"role":"assistant","provider":"google-antigravity","model":"gemini-3-pro","usage":{"input":1,"output":1}}}' + printf '%s\n%s\n' "$openclaw" "$openclaw" >"$test_tmp/.openclaw/agents/main/sessions/s.jsonl" + sqlite "$test_tmp/.local/share/opencode/opencode.db" \ + "CREATE TABLE message (session_id TEXT, data TEXT); + CREATE TABLE session_message (session_id TEXT, data TEXT, type TEXT); + INSERT INTO message VALUES + ('s1', '{\"role\":\"assistant\",\"providerID\":\"google\",\"modelID\":\"antigravity-gemini-3-pro\",\"time\":{\"created\":$(in_ms 0)},\"tokens\":{\"input\":10,\"output\":5,\"reasoning\":1,\"cache\":{\"read\":2,\"write\":0}}}'), + ('s2', '{\"role\":\"assistant\",\"providerID\":\"google\",\"modelID\":\"gemini-3-pro\",\"time\":{\"created\":$(in_ms 0)},\"tokens\":{\"input\":999,\"output\":999}}'), + ('s3', '{\"role\":\"user\",\"providerID\":\"google\",\"modelID\":\"antigravity-gemini-3-pro\",\"tokens\":{\"input\":999}}'); + INSERT INTO session_message VALUES + ('s4', '{\"model\":{\"providerID\":\"google\",\"id\":\"antigravity-gemini-3-pro\"},\"time\":{\"created\":$(in_ms 0)},\"tokens\":{\"input\":4,\"output\":\"x\"}}', 'assistant');" + record=$(PI_DIR="$test_tmp/pi-elsewhere" collect "") + asked >/dev/null + [[ $(jq -c '{ready, hasLocalStats, totalPrompts, totalSessions, todayPrompts, todayTotalTokens, models: (.modelUsage | keys)}' <<<"$record") == '{"ready":true,"hasLocalStats":true,"totalPrompts":7,"totalSessions":6,"todayPrompts":6,"todayTotalTokens":155,"models":["antigravity-gemini-3-pro","gemini-3-pro"]}' ]] || + fail "Antigravity collector counts usage from pi, omp, OpenClaw, and opencode's plugin" "$record" + [[ $(jq -c '.modelUsage["gemini-3-pro"]' <<<"$record") == '{"inputTokens":135,"outputTokens":23,"cacheReadInputTokens":5,"cacheCreationInputTokens":0}' ]] || + fail "Antigravity collector counts each harness answer's tokens once" "$record" + [[ $(jq -r '.usageStatusText' <<<"$record") == "Waiting for auth" ]] || + fail "Antigravity collector in use only through local stats waits for a sign-in" "$record" + pass "Antigravity collector counts usage from pi, omp, OpenClaw, and opencode's plugin" +else + skip "Antigravity collector counts usage from pi, omp, OpenClaw, and opencode's plugin (no rg)" +fi From 0e104aa4d4370ea00d1e9dd1e7b851c0daa49c3f Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:10:05 -0400 Subject: [PATCH 058/546] Skip Antigravity quota buckets whose fraction isn't finite Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-agy | 3 ++- test/shell.d/agent-usage-agy-limits-test.sh | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/bin/omarchy-agent-usage-agy b/bin/omarchy-agent-usage-agy index 9592b5061b2..6847dfbde85 100755 --- a/bin/omarchy-agent-usage-agy +++ b/bin/omarchy-agent-usage-agy @@ -25,6 +25,7 @@ import fcntl import hashlib import http.client import json +import math import os import re import shutil @@ -268,7 +269,7 @@ def used(bucket): remaining = float(remaining) except (TypeError, ValueError): return None - if remaining != remaining: + if not math.isfinite(remaining): return None return round(max(0.0, min(1.0, 1.0 - remaining)), 4) diff --git a/test/shell.d/agent-usage-agy-limits-test.sh b/test/shell.d/agent-usage-agy-limits-test.sh index 8f4c8f518d4..7357219d9cf 100755 --- a/test/shell.d/agent-usage-agy-limits-test.sh +++ b/test/shell.d/agent-usage-agy-limits-test.sh @@ -145,7 +145,8 @@ pass "Antigravity collector finds the Gemini group wherever it is listed" snake=$(test_limits '{"buckets": [ {"window": "5h", "remaining_fraction": 0.4, "reset_time": "2026-09-15T19:31:46Z"}, - {"window": "weekly", "remainingFraction": "plenty"} + {"window": "weekly", "remainingFraction": "plenty"}, + {"bucketId": "gemini-weekly", "remainingFraction": Infinity} ]}') [[ $(jq -c 'map({label, percent, resetsAt})' <<<"$snake") == '[{"label":"Session (5-hour)","percent":0.6,"resetsAt":"2026-09-15T19:31:46+00:00"}]' ]] || fail "Antigravity collector reads snake_case buckets listed without groups and skips unreadable fractions" "$snake" From 8193d6cee2e6e1a5e349d7834d04999ff3d5c5fd Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:04:52 -0400 Subject: [PATCH 059/546] Bring the OpenRouter collector in line with the collector contract Each key keeps its own meter and balance, kept stale past a failed check and reused for 15 seconds unless --force. A refused key reads as an expired sign-in, retryAdvised marks only an unreachable API, an unused key prints an empty status, and the scan is rebuilt under its lock. pi and omp sessions on the openrouter provider now count, and broken timestamps or token counts skip their row. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-openrouter | 544 +++++++++++++----- .../agent-usage-openrouter-scanner-test.sh | 266 ++++++++- 2 files changed, 677 insertions(+), 133 deletions(-) diff --git a/bin/omarchy-agent-usage-openrouter b/bin/omarchy-agent-usage-openrouter index 836dc6cdea1..67a0768262f 100755 --- a/bin/omarchy-agent-usage-openrouter +++ b/bin/omarchy-agent-usage-openrouter @@ -6,7 +6,10 @@ 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. +(providerID == "openrouter", role == "assistant") and from pi/omp session +files whose assistant messages ran on the openrouter provider, read-only. +Each key keeps its own last good meter and balance, shown stale when a +check fails. 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 @@ -22,13 +25,17 @@ omacom/omarchy#10330 (dedicated pi collector proposal). from __future__ import annotations import argparse +import fcntl import hashlib import json +import math import os 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 @@ -45,33 +52,53 @@ API_BASE = "https://openrouter.ai/api/v1" SCAN_REUSE_SECONDS = 20 LIMITS_ONLY_REUSE_SECONDS = 900 +# Opening and shutting the panel must not turn into three API calls per flick, +# so a probe this recent is reused; --force, a person pressing refresh, skips +# it. +PROBE_MIN_INTERVAL_SECONDS = 15 + +# Token counts beyond this are a corrupt row, not usage. +MAX_TOKENS = 10**12 + class OpenRouterError(Exception): pass +# Counts from disk and from the API are untrusted: a boolean, a negative, a +# non-finite or an absurdly large value counts as nothing rather than taking +# the record down or inflating it. def number(value: Any) -> int: + if isinstance(value, bool): + return 0 try: - return max(0, int(value or 0)) - except (TypeError, ValueError): + parsed = float(value or 0) + except (TypeError, ValueError, OverflowError): + return 0 + if not math.isfinite(parsed) or parsed <= 0 or parsed > MAX_TOKENS: return 0 + return int(parsed) +# The local day a timestamp names, or "" when it names none: a row with a +# broken timestamp is skipped rather than filed under today or 1970. 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") + if isinstance(value, bool) or value is None: + return "" 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") + if isinstance(value, (int, float)): + if not math.isfinite(value) or value <= 0: + return "" + seconds = value / 1000 if value > 10_000_000_000 else value + if seconds > 32_503_680_000: + return "" + return datetime.fromtimestamp(seconds).astimezone().strftime("%Y-%m-%d") + parsed = datetime.fromisoformat(str(value).strip().replace("Z", "+00:00")) + if parsed.tzinfo is not None: + parsed = parsed.astimezone() + return parsed.strftime("%Y-%m-%d") + except (ValueError, OverflowError, OSError): + return "" def model_name(raw: Any) -> str: @@ -155,10 +182,21 @@ def management_key(config: dict[str, Any]) -> str: # ---------------------------------------------------------------- API +# No server answered at all (DNS, connection, timeout): the one failure worth +# a sooner retry. class OpenRouterTransportError(OpenRouterError): pass +# The key was refused. 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 what a refusal means. +class OpenRouterAuthError(OpenRouterError): + def __init__(self, code: int): + super().__init__(f"OpenRouter refused the key (status {code})") + self.code = code + + def api_get(path: str, key: str) -> dict: req = urllib.request.Request( API_BASE + path, headers={"Authorization": "Bearer " + key, "Accept": "application/json"} @@ -167,21 +205,28 @@ def api_get(path: str, key: str) -> dict: 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: + except urllib.error.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 + raise OpenRouterAuthError(error.code) + raise OpenRouterError(f"OpenRouter's API returned status {error.code}.") + except urllib.error.URLError as error: + raise OpenRouterTransportError("Couldn't reach the OpenRouter API. Retrying shortly.") 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 + raise OpenRouterTransportError("Couldn't reach the OpenRouter API. Retrying shortly.") from error except (json.JSONDecodeError, ValueError) as error: - raise OpenRouterError("OpenRouter returned an invalid response") from error + raise OpenRouterError("OpenRouter returned an invalid response.") from error + + +# /activity names its day as a date, sometimes with a time after it. Anything +# that isn't one files under no day rather than a made-up one. +def activity_day(value: Any) -> str: + text = str(value or "").strip()[:10] + try: + return datetime.strptime(text, "%Y-%m-%d").strftime("%Y-%m-%d") + except ValueError: + return "" def summarize_activity(items: list[Any], today: datetime) -> dict[str, Any]: @@ -202,7 +247,7 @@ def summarize_activity(items: list[Any], today: datetime) -> dict[str, Any]: for raw in items: if not isinstance(raw, dict): continue - day = str(raw.get("date") or "") + day = activity_day(raw.get("date")) 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")) @@ -230,24 +275,50 @@ def summarize_activity(items: list[Any], today: datetime) -> dict[str, Any]: } -def activity_stats(key: str, today: datetime) -> dict[str, Any] | None: +def credential_digest(secret: str) -> str: + return hashlib.sha256(secret.encode("utf-8")).hexdigest()[:16] + + +def activity_stats(key: str, today: datetime, max_age: float) -> tuple[dict[str, Any] | None, bool]: """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. + Returns (stats, transport_failed). Stats are None when there is no + management key (plain tier-1 run) or the key is refused -- the caller then + falls back to the local DB scan. The answer is a stats source like the + local scan, so it is reused on the same schedule; when /activity can't be + reached, today's last answer stands in for it. """ if not key: - return None + return None, False + scan_date = today.strftime("%Y-%m-%d") + try: + cache = cache_root() / f"openrouter-activity-{credential_digest(key)}.json" + cached = read_json(cache) or {} + except OSError: + cache, cached = None, {} + kept = cached.get("stats") if cached.get("scanDate") == scan_date and isinstance(cached.get("stats"), dict) else None + age = time.time() - epoch_ms(cached.get("fetchedAtMs")) / 1000 + if kept is not None and 0 <= age < max_age: + return kept, False + try: payload = api_get("/activity", key) - except OpenRouterError: - return None + except OpenRouterAuthError: + return None, False + except OpenRouterError as error: + print(f"omarchy-agent-usage-openrouter: activity lookup failed ({error})", file=sys.stderr) + return kept, isinstance(error, OpenRouterTransportError) items = payload.get("data") if not isinstance(items, list): - return None + return kept, False stats = empty_stats() stats.update(summarize_activity(items, today)) - return stats + if cache is not None: + try: + write_json(cache, {"schemaVersion": 1, "scanDate": scan_date, "fetchedAtMs": int(time.time() * 1000), "stats": stats}) + except OSError as exc: + print(f"omarchy-agent-usage-openrouter: could not write activity cache ({exc})", file=sys.stderr) + return stats, False # ---------------------------------------------------------------- local stats (opencode DB) @@ -346,6 +417,8 @@ def scan_opencode(scan: LocalScan, db: Path | None = None) -> bool: if not (in_tok or out_tok or cache_read or cache_write): continue day = local_day((entry.get("time") or {}).get("created")) + if not day: + continue model = model_name(str(entry.get("modelID") or "")) except Exception: continue @@ -363,21 +436,111 @@ def cache_root() -> Path: return root +# pi and omp keep sessions under their agent dir, and `--profile=` +# moves a whole agent tree under profiles//. PI_CODING_AGENT_DIR moves +# pi's default one, as pi itself reads it. +def pi_session_roots() -> list[Path]: + pi_agent = Path(os.environ.get("PI_CODING_AGENT_DIR") or (Path.home() / ".pi" / "agent")) + roots = [pi_agent / "sessions", Path.home() / ".omp" / "agent" / "sessions"] + for base in (Path.home() / ".pi", Path.home() / ".omp"): + try: + roots.extend(sorted(child / "agent" / "sessions" for child in (base / "profiles").iterdir() if child.is_dir())) + except OSError: + pass + return roots + + +def scan_pi(scan: LocalScan, roots: list[Path] | None = None) -> bool: + # Returns whether every session file could be read, like scan_opencode. + roots = [root for root in (roots if roots is not None else pi_session_roots()) if root.is_dir()] + if not roots: + return True + # --no-ignore: a $HOME that is a git checkout would otherwise hide the + # session dirs behind its .gitignore. Sorted, so which copy of a forked + # message counts, and so its session, doesn't change from run to run. + try: + proc = subprocess.run( + ["rg", "--json", "--no-ignore", "--hidden", "--sort", "path", "--glob", "*.jsonl", "-e", r'"provider"\s*:\s*"openrouter"', *map(str, roots)], + capture_output=True, text=True, errors="replace", timeout=60, + ) + except (OSError, subprocess.SubprocessError): + return False + + seen: set[tuple[str, str]] = set() + for raw in proc.stdout.splitlines(): + try: + event = json.loads(raw) + if event.get("type") != "match": + continue + data = event.get("data") or {} + path = str((data.get("path") or {}).get("text") or "") + entry = json.loads((data.get("lines") or {}).get("text") or "") + if not isinstance(entry, dict) or entry.get("type") != "message": + continue + message = entry.get("message") + if not isinstance(message, dict) or message.get("role") != "assistant": + continue + if str(message.get("provider") or "") != "openrouter": + continue + # A forked session repeats its parent's messages with their ids and + # timestamps, so those count once; a line without both always counts. + entry_id, stamp = entry.get("id"), entry.get("timestamp") + if entry_id and stamp: + identity = (str(entry_id), str(stamp)) + if identity in seen: + continue + seen.add(identity) + usage = message.get("usage") if isinstance(message.get("usage"), dict) else {} + in_tok = number(usage.get("input")) + out_tok = number(usage.get("output")) + cache_read = number(usage.get("cacheRead")) + cache_write = number(usage.get("cacheWrite")) + if not (in_tok or out_tok or cache_read or cache_write): + in_tok = number(usage.get("totalTokens")) + if not (in_tok or out_tok or cache_read or cache_write): + continue + day = local_day(stamp if stamp is not None else message.get("timestamp")) + if not day: + continue + model = model_name(str(message.get("model") or "")) + except Exception: + continue + scan.add_usage(day, "pi:" + path, model, in_tok, out_tok, cache_read, cache_write) + # rg exits 1 for no matches and 2 when something couldn't be read. + return proc.returncode in (0, 1) + + +def scan_local(scan: LocalScan) -> bool: + # Both sources always run, so one failing still leaves the other's numbers. + opencode_complete = scan_opencode(scan) + pi_complete = scan_pi(scan) + return opencode_complete and pi_complete + + def scan_cache_paths() -> tuple[Path, Path]: - db = opencode_db_path() - digest = hashlib.sha1(str(db).encode("utf-8")).hexdigest()[:16] + sources = [opencode_db_path(), *pi_session_roots()] + digest = hashlib.sha1("\n".join(map(str, sources)).encode("utf-8")).hexdigest()[:16] root = cache_root() return root / f"openrouter-scan-{digest}.json", root / f"openrouter-scan-{digest}.lock" +def read_json(path: Path) -> dict | None: + try: + parsed = json.loads(path.read_text(encoding="utf-8")) + except Exception: + return None + return parsed if isinstance(parsed, dict) else None + + 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 + # 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 read_json(path) except Exception: return None return None @@ -403,27 +566,40 @@ def cached_local_stats(today: datetime, max_age: float) -> tuple[dict, bool]: 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) + + def usable(cached: dict | None) -> dict | None: 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 + return cached["stats"] + return None + + try: + cache_file, lock_file = scan_cache_paths() + stats = usable(read_fresh_json(cache_file, max_age)) + if stats is not None: + return stats, True + # Concurrent runs (the update command and the panel's own refresh) wait + # for one scan instead of each reading the whole database. + with lock_file.open("w") as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + stats = usable(read_fresh_json(cache_file, max_age)) + if stats is not None: + return stats, True + complete = scan_local(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) + complete = scan_local(scan) return scan.snapshot(), complete @@ -437,9 +613,6 @@ def base_record(**overrides: Any) -> dict: "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": "", @@ -463,134 +636,241 @@ def key_limit_meter(key_info: dict) -> list[dict]: try: limit = key_info.get("limit") remaining = key_info.get("limit_remaining") - if limit is None or remaining is None or float(limit) <= 0: + if limit is None or remaining is None or isinstance(limit, bool) or isinstance(remaining, bool): return [] lim = float(limit) rem = max(0.0, float(remaining)) except (TypeError, ValueError): return [] + if not (math.isfinite(lim) and math.isfinite(rem)) or lim <= 0: + 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": ""}] + return [{"label": label, "title": label, "percent": min(1.0, (lim - min(rem, lim)) / lim), "resetsAt": ""}] + + +def dollars(value: Any) -> float | None: + if isinstance(value, bool): + return None + try: + parsed = float(value or 0) + except (TypeError, ValueError): + return None + return parsed if math.isfinite(parsed) else None -def limits_cache() -> Path: - return cache_root() / "openrouter-limits.json" +# Balance: total purchased minus total used, or None for an account that never +# bought credits. +def credits_balance(credits: dict) -> dict | None: + total_credits = dollars(credits.get("total_credits")) + total_usage = dollars(credits.get("total_usage")) + if not total_credits or total_credits <= 0 or total_usage is None: + return None + return { + "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, + } -# 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]: +def epoch_ms(value: Any) -> int: + if isinstance(value, bool): + return 0 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 + parsed = float(value or 0) + except (TypeError, ValueError): + return 0 + return int(parsed) if math.isfinite(parsed) and parsed > 0 else 0 + - limits = key_limit_meter(key_info if isinstance(key_info, dict) else {}) - fetched_at = round(time.time() * 1000) +# A kept window whose reset has passed is an untouched allowance, not an +# unknown one. The key meter only ever names a period, so this is a guard for +# a cache written by a meter that carries a reset time. +def limit_window_open(entry: dict, now: datetime) -> bool: + raw = str(entry.get("resetsAt") or "") + if raw == "": + return True 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, "", "" + 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 + + +def usable_cached_limits(cached: dict) -> list[dict]: + entries = cached.get("limits") + if not isinstance(entries, list): + return [] + now = datetime.now(timezone.utc) + return [ + entry if limit_window_open(entry, now) else dict(entry, percent=0.0, resetsAt="") + for entry in entries if isinstance(entry, dict) + ] + + +# Each credential keeps its own last good answer, so another key never shows +# this one's numbers. Refreshes can overlap (the panel's own and a sooner +# retry), so an answer only replaces one fetched before it. +def probe_cache(kind: str, secret: str) -> Path: + return cache_root() / f"openrouter-{kind}-{credential_digest(secret)}.json" + + +def remember(path: Path, part: dict) -> None: + current = read_json(path) or {} + if epoch_ms(current.get("fetchedAtMs")) <= part["fetchedAtMs"]: + write_json(path, part) + + +def probe(endpoint: str, key: str, cache: Path, parse, force: bool) -> tuple[dict | None, OpenRouterError | None]: + """One endpoint's answer as (part, error). + + A probe this recent is reused, a fresh answer replaces the cached one, and + a failed one falls back on it. part is None when nothing was ever fetched. + """ + cached = read_json(cache) or {} + fetched_at = epoch_ms(cached.get("fetchedAtMs")) + kept = dict(cached, fetchedAtMs=fetched_at) if fetched_at else None + age = time.time() - fetched_at / 1000 + if kept is not None and not force and 0 <= age < PROBE_MIN_INTERVAL_SECONDS: + return kept, None + try: + payload = api_get(endpoint, key) + except OpenRouterError as error: + return kept, error + data = payload.get("data") + # Floored, never rounded: a stamp half a millisecond ahead reads as a clock + # that went back and fails the 0 <= age guard on the very next run. + part = dict(parse(data if isinstance(data, dict) else {}), fetchedAtMs=int(time.time() * 1000)) + try: + remember(cache, part) + except OSError as exc: + print(f"omarchy-agent-usage-openrouter: could not write {cache.name} ({exc})", file=sys.stderr) + return part, None + + +def collect_billing(key: str, manager: str, force: bool) -> dict[str, Any]: + """The key-limit meter and the balance, each kept past a failed check. + + The meter describes an inference key's own credit limit, so with only a + management key there is none to show; the balance is the account's and + either key can read it. + """ + result: dict[str, Any] = { + "limits": [], "limitsStale": False, "limitsFetchedAt": 0, + "usageStatusText": "", "authHelpText": "", "balance": None, "retryAdvised": False, + } + parts: dict[str, dict | None] = {} + errors: dict[str, OpenRouterError] = {} + if key: + parts["key"], error = probe("/key", key, probe_cache("limits", key), lambda data: {"limits": key_limit_meter(data)}, force) + if error: + errors["key"] = error + credits_key = key or manager + parts["credits"], error = probe("/credits", credits_key, probe_cache("credits", credits_key), lambda data: {"balance": credits_balance(data)}, force) + if error: + errors["credits"] = error + + meter = parts.get("key") + if meter is not None: + result["limits"] = usable_cached_limits(meter) + credits = parts.get("credits") + if credits is not None and isinstance(credits.get("balance"), dict): + result["balance"] = credits["balance"] + kept = [part for part in parts.values() if part is not None] + if kept: + result["limitsFetchedAt"] = min(part["fetchedAtMs"] for part in kept) + # Anything that didn't come back fresh marks the record stale, kept or not; + # the panel shows the kept numbers' age on hover. + result["limitsStale"] = bool(errors) + result["retryAdvised"] = any(isinstance(error, OpenRouterTransportError) for error in errors.values()) + + refused = [error for error in errors.values() if isinstance(error, OpenRouterAuthError)] + if refused: + result["usageStatusText"] = "Sign-in expired" + result["authHelpText"] = "OpenRouter refused the API key. " + AUTH_HELP + elif "key" in errors and meter is None: + result["usageStatusText"] = "OpenRouter limits unavailable" + result["authHelpText"] = str(errors["key"]) + elif "credits" in errors and credits is None: + result["usageStatusText"] = "Balance unavailable" + result["authHelpText"] = str(errors["credits"]) + elif errors: + result["authHelpText"] = str(next(iter(errors.values()))) + return result def scan(args: Any) -> dict[str, Any]: config = read_config() key = api_key(config) manager = management_key(config) + # No key anywhere: OpenRouter isn't used here, and the panel skips it. if not key and not manager: - return base_record(usageStatusText="OpenRouter unavailable", authHelpText=AUTH_HELP) + return base_record() today = datetime.now().astimezone() + max_age = 0 if args.force else (LIMITS_ONLY_REUSE_SECONDS if args.limits_only else SCAN_REUSE_SECONDS) # 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(manager, today) + # every machine and client. Without one (or a refused one) fall back to + # the local opencode and pi scan below. + managed, activity_unreachable = activity_stats(manager, today, max_age) if managed is not None: - stats = managed + stats = dict(managed) complete = True has_prompt_stats = False - # Account-global numbers must merge by widest value, not by sum. + # Tier-2 numbers are account-global (billing API), like Fireworks: every + # synced device reports the same truth, so aggregation must merge them by + # widest value, not sum them. Tier 1 is machine-local instead. 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 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 - # one sooner retry instead of waiting out the full refresh interval. - balance = None - retry_advised = False - try: - 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: - 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) + billing = collect_billing(key, manager, args.force) 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, + limits=billing["limits"], + limitsStale=billing["limitsStale"], + limitsFetchedAt=billing["limitsFetchedAt"], + usageStatusText=billing["usageStatusText"], + authHelpText=billing["authHelpText"], ) - if retry_advised: + # Usually the seconds after login before the network is up; the panel + # tries once more soon instead of waiting out its refresh interval. + if billing["retryAdvised"] or activity_unreachable: 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: + if billing["balance"]: + record["balance"] = billing["balance"] + # A key nothing has been spent through, with no credits bought and no limit + # set, has nothing to show: the machine doesn't use OpenRouter yet. + if not has_data and not billing["balance"] and not record["limits"] and not record["usageStatusText"] and not record["limitsStale"]: record["ready"] = False - record["usageStatusText"] = "No OpenRouter usage found" - record["authHelpText"] = AUTH_HELP + record["tierLabel"] = "" 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") + # --force rescans and re-probes everything; --limits-only reuses a recent + # local scan or activity history for up to 15 minutes, while the key meter + # and balance are probed like a normal run. + ap.add_argument("--force", action="store_true", help="rescan sessions and re-probe the API, ignoring caches") + ap.add_argument("--limits-only", action="store_true", help="reuse any recent session 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") + record = base_record(usageStatusText="OpenRouter unavailable", authHelpText="OpenRouter usage scan failed") print(f"omarchy-agent-usage-openrouter: {type(error).__name__}", file=sys.stderr) print(json.dumps(record, separators=(",", ":"))) return 0 diff --git a/test/shell.d/agent-usage-openrouter-scanner-test.sh b/test/shell.d/agent-usage-openrouter-scanner-test.sh index a6b93b71f3d..e886166de3d 100755 --- a/test/shell.d/agent-usage-openrouter-scanner-test.sh +++ b/test/shell.d/agent-usage-openrouter-scanner-test.sh @@ -29,6 +29,12 @@ no_key=$("$COLLECTOR") [[ $(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" +[[ $(jq -c '{usageStatusText, authHelpText, limits, balance, limitsStale, limitsFetchedAt}' <<<"$no_key") == '{"usageStatusText":"","authHelpText":"","limits":[],"balance":null,"limitsStale":false,"limitsFetchedAt":0}' ]] || + fail "an unused machine reports no trouble" "$no_key" +pass "an unused machine reports no trouble" +[[ $(jq -c '[.todayPrompts, .todaySessions, .todayTotalTokens, .todayTokensByModel, .recentDays, .totalPrompts, .totalSessions, .activeDays, .activeDates, .modelUsage]' <<<"$no_key") == '[0,0,0,{},[],0,0,0,[],{}]' ]] || + fail "an unused machine prints the full empty stats block" "$no_key" +pass "an unused machine prints the full empty stats block" # 2. Fixture opencode DB: 2 assistant rows (openrouter) + 1 other provider (ignored) mkdir -p "$XDG_DATA_HOME/opencode" @@ -165,7 +171,7 @@ 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" +check "$(jq -r .status <<<"$stale_check")" "" "a failed lookup that kept the meter leaves the status to limitsStale" # 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. @@ -235,4 +241,262 @@ 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" +# The remaining checks drive the collector module with a stubbed API. Each +# starts from an empty cache so one check's answers never feed another's. +cat >"$TEST_HOME/prelude.py" <<'PY' +import importlib.machinery, importlib.util, json, os, shutil, sys, time +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) +shutil.rmtree(os.environ["XDG_CACHE_HOME"], ignore_errors=True) + +class Args: + def __init__(self, force=False, limits_only=False): + self.force = force + self.limits_only = limits_only + +calls = [] +answers = {} +def fake_api_get(path, key): + calls.append(f"{path}={key}") + answer = answers.get((path, key), answers.get(path)) + if isinstance(answer, Exception): + raise answer + if answer is None: + raise AssertionError(path) + return answer +mod.api_get = fake_api_get +os.environ["OPENROUTER_API_KEY"] = "sk-test" +os.environ["OPENROUTER_MANAGEMENT_KEY"] = "" +METER = {"data": {"limit": 10, "limit_remaining": 4, "limit_reset": "monthly"}} +CREDITS = {"data": {"total_credits": 10, "total_usage": 1}} +UNREACHABLE = mod.OpenRouterTransportError("Couldn't reach the OpenRouter API. Retrying shortly.") +pick = lambda r: {k: r.get(k) for k in ("ready", "usageStatusText", "authHelpText", "limits", "limitsStale", "limitsFetchedAt", "balance", "retryAdvised")} +PY +run_py() { + cat "$TEST_HOME/prelude.py" - | python3 - "$COLLECTOR" +} + +# 7. A key nothing was spent through, with no credits and no limit, is an +# unused machine: ready=false and no status from outside the vocabulary. +idle=$(run_py <<'PY' +answers.update({"/key": {"data": {"limit": None}}, "/credits": {"data": {"total_credits": 0, "total_usage": 0}}}) +os.environ["XDG_DATA_HOME"] = "/nonexistent" +print(json.dumps(pick(mod.scan(Args())))) +PY +) +check "$(jq -c '{ready, usageStatusText, authHelpText}' <<<"$idle")" '{"ready":false,"usageStatusText":"","authHelpText":""}' "a key with nothing to show is an unused machine" + +# 8. Probes are reused for PROBE_MIN_INTERVAL_SECONDS, and --force skips that. +reuse=$(run_py <<'PY' +answers.update({"/key": METER, "/credits": CREDITS}) +mod.scan(Args()) +first = len(calls) +mod.scan(Args()) +reused = len(calls) - first +mod.scan(Args(limits_only=True)) +limits_only = len(calls) - first - reused +mod.scan(Args(force=True)) +forced = len(calls) - first - reused - limits_only +print(json.dumps({"first": first, "reused": reused, "limitsOnly": limits_only, "forced": forced})) +PY +) +check "$(jq -c . <<<"$reuse")" '{"first":2,"reused":0,"limitsOnly":0,"forced":2}' "a recent probe is reused and --force probes again" + +# 9. Each key keeps its own meter: another key whose check fails never shows +# this one's numbers. +per_key=$(run_py <<'PY' +answers.update({"/key": METER, "/credits": CREDITS}) +mod.scan(Args(force=True)) +answers.update({"/key": mod.OpenRouterError("OpenRouter's API returned status 500."), "/credits": mod.OpenRouterError("OpenRouter's API returned status 500.")}) +os.environ["OPENROUTER_API_KEY"] = "sk-other" +other = mod.scan(Args(force=True)) +files = sorted(p.name for p in mod.cache_root().glob("openrouter-limits-*.json")) +print(json.dumps({"other": pick(other), "files": len(files)})) +PY +) +check "$(jq -c '.other | {limits, usageStatusText, retryAdvised}' <<<"$per_key")" '{"limits":[],"usageStatusText":"OpenRouter limits unavailable","retryAdvised":null}' "another key never shows this key's meter" +check "$(jq -r .other.authHelpText <<<"$per_key")" "OpenRouter's API returned status 500." "a failed check with nothing kept says what failed" +check "$(jq -r .files <<<"$per_key")" "1" "the meter is cached per key" + +# 10. A refused key is a sign-in to redo, never a raw status code, and is not +# worth a sooner retry. +for code in 401 403; do + refused=$(CODE=$code run_py <<'PY' +answers.update({"/key": METER, "/credits": CREDITS}) +mod.scan(Args(force=True)) +code = int(os.environ["CODE"]) +answers.update({"/key": mod.OpenRouterAuthError(code), "/credits": mod.OpenRouterAuthError(code)}) +print(json.dumps(pick(mod.scan(Args(force=True))))) +PY +) + check "$(jq -c '{usageStatusText, limitsStale, retryAdvised, meter: (.limits | length), ready}' <<<"$refused")" '{"usageStatusText":"Sign-in expired","limitsStale":true,"retryAdvised":null,"meter":1,"ready":true}' "a $code keeps the last meter and asks to sign in again" + [[ $(jq -r .authHelpText <<<"$refused") == "OpenRouter refused the API key. "* ]] || + fail "a $code explains itself without a raw status" "$refused" + pass "a $code explains itself without a raw status" +done + +# 11. A failed /credits keeps the last good balance, marked stale; with none +# kept, the balance says it is unavailable. +balance=$(run_py <<'PY' +answers.update({"/key": METER, "/credits": CREDITS}) +mod.scan(Args(force=True)) +answers["/credits"] = mod.OpenRouterError("OpenRouter's API returned status 502.") +kept = mod.scan(Args(force=True)) +for path in mod.cache_root().glob("openrouter-credits-*.json"): + path.unlink() +missing = mod.scan(Args(force=True)) +print(json.dumps({"kept": pick(kept), "missing": pick(missing)})) +PY +) +check "$(jq -c '.kept | {remaining: .balance.remaining, limitsStale, usageStatusText}' <<<"$balance")" '{"remaining":9.0,"limitsStale":true,"usageStatusText":""}' "a failed balance check keeps the last balance, stale" +check "$(jq -c '.missing | {balance, limitsStale, usageStatusText}' <<<"$balance")" '{"balance":null,"limitsStale":true,"usageStatusText":"Balance unavailable"}' "a failed balance check with nothing kept says so" + +# 12. retryAdvised marks only a service that couldn't be reached at all. +retry=$(run_py <<'PY' +answers.update({"/key": UNREACHABLE, "/credits": CREDITS}) +unreachable = mod.scan(Args(force=True)) +answers.update({"/key": mod.OpenRouterError("OpenRouter's API returned status 500.")}) +server = mod.scan(Args(force=True)) +answers.update({"/key": METER}) +healthy = mod.scan(Args(force=True)) +print(json.dumps({"unreachable": unreachable.get("retryAdvised"), "server": server.get("retryAdvised"), "healthy": pick(healthy)})) +PY +) +check "$(jq -c '[.unreachable, .server]' <<<"$retry")" '[true,null]' "retryAdvised marks only an unreachable service" +check "$(jq -c '.healthy | {usageStatusText, authHelpText, limitsStale}' <<<"$retry")" '{"usageStatusText":"","authHelpText":"","limitsStale":false}' "a healthy check clears the help text" + +# 13. A kept window whose reset has passed comes back at zero, never dropped. +elapsed=$(run_py <<'PY' +cache = mod.probe_cache("limits", "sk-test") +mod.write_json(cache, {"fetchedAtMs": round(time.time() * 1000) - 3_600_000, "limits": [ + {"label": "Daily key limit", "title": "Daily key limit", "percent": 0.7, "resetsAt": "2020-01-01T00:00:00Z"}, + {"label": "Key limit", "title": "Key limit", "percent": 0.4, "resetsAt": ""}, +]}) +answers.update({"/key": UNREACHABLE, "/credits": CREDITS}) +print(json.dumps(pick(mod.scan(Args())))) +PY +) +check "$(jq -c '[.limits[] | {percent, resetsAt}]' <<<"$elapsed")" '[{"percent":0.0,"resetsAt":""},{"percent":0.4,"resetsAt":""}]' "an elapsed kept window comes back at zero" + +# 14. A cache written under a clock that later went back is still the last +# good answer, but never counts as a recent probe. +clock=$(run_py <<'PY' +cache = mod.probe_cache("limits", "sk-test") +mod.write_json(cache, {"fetchedAtMs": round(time.time() * 1000) + 3_600_000, "limits": [{"label": "Key limit", "title": "Key limit", "percent": 0.4, "resetsAt": ""}]}) +answers.update({"/key": METER, "/credits": CREDITS}) +record = mod.scan(Args()) +print(json.dumps({"probed": any(c.startswith("/key=") for c in calls), "percent": record["limits"][0]["percent"]})) +PY +) +check "$(jq -c . <<<"$clock")" '{"probed":true,"percent":0.6}' "a future-dated probe cache is not reused" + +# 15. Rows with broken timestamps or tokens are skipped, never filed under +# today, and never take the record down. +sqlite3 "$XDG_DATA_HOME/opencode/opencode.db" <<'SQL' +INSERT INTO message VALUES + ('b1','ses_bad',0,0,'{"role":"assistant","providerID":"openrouter","modelID":"a/bad-time","tokens":{"input":50},"time":{"created":true}}'), + ('b2','ses_bad',0,0,'{"role":"assistant","providerID":"openrouter","modelID":"a/bad-time","tokens":{"input":50},"time":{"created":-5}}'), + ('b3','ses_bad',0,0,'{"role":"assistant","providerID":"openrouter","modelID":"a/bad-time","tokens":{"input":50},"time":{"created":"garbage"}}'), + ('b4','ses_bad',0,0,'{"role":"assistant","providerID":"openrouter","modelID":"a/bad-time","tokens":{"input":50},"time":{"created":1e300}}'), + ('b5','ses_bad',0,0,'{"role":"assistant","providerID":"openrouter","modelID":"a/bad-tokens","tokens":{"input":true,"output":-3,"reasoning":1e300,"cache":{"read":"NaN"}},"time":{"created":1787000000000}}'); +SQL +robust=$(run_py <<'PY' +from datetime import datetime +scan = mod.LocalScan(datetime.now().astimezone()) +complete = mod.scan_opencode(scan) +stats = scan.snapshot() +print(json.dumps({"complete": complete, "prompts": stats["totalPrompts"], "models": sorted(stats["modelUsage"])})) +PY +) +check "$(jq -c . <<<"$robust")" '{"complete":true,"prompts":1,"models":["gpt-4o-mini"]}' "broken timestamps and tokens skip their row" + +# 16. pi and omp sessions on the openrouter provider count, a forked +# session's repeated messages count once, a line without an id always counts, +# and a $HOME that is a git checkout doesn't hide them. +today_iso=$(date -u +%Y-%m-%dT12:00:00.000Z) +mkdir -p "$HOME/.pi/agent/sessions/--p--" "$HOME/.omp/profiles/work/agent/sessions/--q--" +printf '.pi/\n.omp/\n' >"$HOME/.gitignore" +git -C "$HOME" init -q +cat >"$HOME/.pi/agent/sessions/--p--/a-parent.jsonl" <"$HOME/.pi/agent/sessions/--p--/b-fork.jsonl" <"$HOME/.omp/profiles/work/agent/sessions/--q--/omp.jsonl" <= 1.0})) +PY +) +check "$(jq -r .waited <<<"$lock_wait")" "true" "the local scan is rebuilt under its lock" + +# 19. The 30-day history is reused like a scan, and when /activity can't be +# reached, today's last answer stands in for it. +activity=$(run_py <<'PY' +from datetime import datetime +today = datetime.now().astimezone().strftime("%Y-%m-%d") +os.environ["OPENROUTER_MANAGEMENT_KEY"] = "sk-mgmt" +answers.update({"/activity": {"data": [{"date": today, "model": "a/b", "prompt_tokens": 10, "completion_tokens": 5}]}, "/key": METER, "/credits": CREDITS}) +mod.scan(Args(force=True)) +calls.clear() +reused = mod.scan(Args()) +reused_calls = [c for c in calls if c.startswith("/activity")] +answers["/activity"] = UNREACHABLE +kept = mod.scan(Args(force=True)) +print(json.dumps({"reusedCalls": len(reused_calls), "reusedScope": reused.get("scope"), "keptScope": kept.get("scope"), "keptToday": kept["todayTotalTokens"], "retry": kept.get("retryAdvised")})) +PY +) +check "$(jq -c . <<<"$activity")" '{"reusedCalls":0,"reusedScope":"account","keptScope":"account","keptToday":15,"retry":true}' "the activity history is reused and kept past an unreachable check" + pass "ALL OPENROUTER SCANNER TESTS PASSED" From affbca57ea557cca5ca5510fe35c8215653ee1ed Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:09:54 -0400 Subject: [PATCH 060/546] Keep Fireworks' last good usage and balance past a failed check Each API key keeps its account's usage rows and balance, shown stale when a check fails instead of blanking the tab. Usage is reused like a scan and the balance for 15 seconds unless --force; --limits-only re-checks only the balance. An unused machine prints an empty status, a rejected key reads as an expired sign-in, and retryAdvised marks only an unreachable API. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-fireworks | 302 +++++++++++++++--- .../agent-usage-fireworks-scanner-test.sh | 194 ++++++++++- 2 files changed, 441 insertions(+), 55 deletions(-) diff --git a/bin/omarchy-agent-usage-fireworks b/bin/omarchy-agent-usage-fireworks index fe04bc1df38..0f79d2044f9 100755 --- a/bin/omarchy-agent-usage-fireworks +++ b/bin/omarchy-agent-usage-fireworks @@ -9,16 +9,23 @@ the last 30 days. Fireworks does not expose its prepaid ledger, so the record carries an estimated balance instead of rate limits: the credits configured in ~/.config/omarchy/agents/fireworks.json minus rated account costs since the funding date. The agents panel only ever reads the JSON this prints. + +Each API key keeps its account's last good usage rows and balance under the +agent-usage cache, so a failed check shows them stale rather than blank. """ from __future__ import annotations import argparse import configparser +import hashlib import json +import math import os import re import sys +import tempfile +import time as time_module import urllib.error import urllib.parse import urllib.request @@ -32,16 +39,44 @@ AGENT_NAME = "Fireworks" AUTH_HELP = "Set FIREWORKS_API_KEY, run `firectl set-api-key`, or sign in to Fireworks in pi or opencode." API_BASE_URL = "https://api.fireworks.ai" +# The usage rows are reused like a local scan: briefly to dedup concurrent +# runs, and for 15 minutes on the panel's near-limit refresh. The balance is +# re-checked at most this often unless --force. +SCAN_REUSE_SECONDS = 20 +LIMITS_ONLY_REUSE_SECONDS = 900 +PROBE_MIN_INTERVAL_SECONDS = 15 + +# Token counts beyond this are a corrupt row, not usage. +MAX_TOKENS = 10**12 + class FireworksError(Exception): pass +# No server answered at all (DNS, connection, timeout): the one failure worth +# a sooner retry. +class FireworksTransportError(FireworksError): + pass + + +# The key itself was refused, so only a new one helps. +class FireworksAuthError(FireworksError): + pass + + +# Counts from the API are untrusted: a boolean, a negative, a non-finite or +# an absurdly large value counts as nothing. def number(value: Any) -> int: + if isinstance(value, bool): + return 0 try: - return max(0, round(float(value or 0))) - except (TypeError, ValueError): + parsed = float(value or 0) + except (TypeError, ValueError, OverflowError): + return 0 + if not math.isfinite(parsed) or parsed <= 0 or parsed > MAX_TOKENS: return 0 + return round(parsed) def money_value(value: Any) -> Decimal: @@ -50,9 +85,11 @@ def money_value(value: Any) -> Decimal: try: units = Decimal(str(value.get("units", 0) or 0)) nanos = Decimal(str(value.get("nanos", 0) or 0)) / Decimal("1000000000") - return units + nanos + total = units + nanos except (InvalidOperation, TypeError, ValueError): return Decimal("0") + # A non-finite amount would print as invalid JSON. + return total if total.is_finite() else Decimal("0") def model_id(row: dict[str, Any]) -> str: @@ -73,7 +110,10 @@ def row_date(row: dict[str, Any]) -> str: try: parsed = datetime.fromisoformat(raw.replace("Z", "+00:00")) except ValueError: - return raw[:10] if len(raw) >= 10 else "" + try: + return date.fromisoformat(raw[:10]).isoformat() + except ValueError: + return "" if parsed.tzinfo is None: parsed = parsed.replace(tzinfo=timezone.utc) return parsed.astimezone().date().isoformat() @@ -121,6 +161,8 @@ def base_record(**overrides: Any) -> dict[str, Any]: "usageStatusText": "", "authHelpText": "", "limits": [], + "limitsStale": False, + "limitsFetchedAt": 0, } record.update(empty_stats()) record.update(overrides) @@ -362,15 +404,19 @@ class FireworksClient: return decoded if isinstance(decoded, dict) else {} except urllib.error.HTTPError as error: if error.code == 401: - raise FireworksError("Fireworks rejected the API key") + raise FireworksAuthError("Fireworks rejected the API key") if error.code == 403: raise FireworksError("The Fireworks API key cannot read billing data") if error.code == 404: raise FireworksError("Fireworks account not found") raise FireworksError(f"Fireworks API returned HTTP {error.code}") except urllib.error.URLError as error: - raise FireworksError("Could not reach the Fireworks API") from error - except (json.JSONDecodeError, TimeoutError) as error: + raise FireworksTransportError("Could not reach the Fireworks 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 FireworksTransportError("Could not reach the Fireworks API") from error + except (json.JSONDecodeError, ValueError) as error: raise FireworksError("Fireworks returned an invalid billing response") from error def discover_account(self) -> tuple[str, dict[str, Any]]: @@ -443,6 +489,10 @@ def live_balance(client: FireworksClient, account_id: str) -> Decimal | None: quoted = urllib.parse.quote(normalize_account_id(account_id), safe="") try: payload = client.request(f"/v1/accounts/{quoted}:getBalance") + except (FireworksTransportError, FireworksAuthError): + # Not the permission gate: the balance can't be checked at all, which a + # missing estimate must not paper over as "no balance". + raise except FireworksError: return None candidates = [payload] + [payload.get(field) for field in ("balance", "creditBalance", "prepaidBalance", "amount")] @@ -486,67 +536,215 @@ def estimated_balance( } -def scan(api_base_url: str, auth_path: Path) -> dict[str, Any]: - config = read_config() - api_key, account_id = credentials(auth_path, config) - if not api_key: - return base_record(usageStatusText="Fireworks unavailable", authHelpText=AUTH_HELP) +def current_balance( + client: FireworksClient, + account_id: str, + account: dict[str, Any], + config: dict[str, Any], +) -> dict[str, Any] | None: + live = live_balance(client, account_id) + if live is None: + return estimated_balance(client, account_id, account, config) + try: + funded = Decimal(str(config.get("fundedAmount") or "0")) + if not funded.is_finite() or funded < 0: + funded = Decimal("0") + except InvalidOperation: + funded = Decimal("0") + return { + "remaining": float(live), + "funded": float(funded), + "spent": float(max(Decimal("0"), funded - live)), + "currency": "USD", + "estimated": False, + } - client = FireworksClient(api_key, api_base_url) - account: dict[str, Any] = {} - if account_id: - account_id = normalize_account_id(account_id) - else: - account_id, account = client.discover_account() - today = datetime.now().astimezone().date() - usage = client.usage(account_id, today - timedelta(days=29), today + timedelta(days=1)) - record = base_record(ready=True, hasLocalStats=True) - record.update(summarize_usage(usage, today)) +# ---------------------------------------------------------------- caches - live = live_balance(client, account_id) - if live is not None: - try: - funded = Decimal(str(config.get("fundedAmount") or "0")) - if not funded.is_finite() or funded < 0: - funded = Decimal("0") - except InvalidOperation: - funded = Decimal("0") - record["balance"] = { - "remaining": float(live), - "funded": float(funded), - "spent": float(max(Decimal("0"), funded - live)), - "currency": "USD", - "estimated": False, - } - return record +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 credential_digest(secret: str) -> str: + return hashlib.sha256(secret.encode("utf-8")).hexdigest()[:16] + + +def read_json(path: Path) -> dict[str, Any] | None: try: - balance = estimated_balance(client, account_id, account, config) - if balance: - record["balance"] = balance - except FireworksError as error: - record["usageStatusText"] = "Balance unavailable" - record["authHelpText"] = str(error) + parsed = json.loads(path.read_text(encoding="utf-8")) + except Exception: + return None + return parsed if isinstance(parsed, dict) else None + + +def write_json(path: Path, payload: dict[str, Any]) -> None: + # A temp name unique to this writer: several collector runs can overlap. + 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 epoch_ms(value: Any) -> int: + if isinstance(value, bool): + return 0 + try: + parsed = float(value or 0) + except (TypeError, ValueError): + return 0 + return int(parsed) if math.isfinite(parsed) and parsed > 0 else 0 + + +# The last good answer of one kind (the account's usage rows, or its balance) +# for this API key and account, whatever its file time: a clock that went +# back doesn't lose it. An answer for another account never stands in. +def cached_part(path: Path, account_id: str) -> dict[str, Any] | None: + cached = read_json(path) + if not cached or not epoch_ms(cached.get("fetchedAtMs")): + return None + if account_id and normalize_account_id(str(cached.get("accountId") or "")) != account_id: + return None + return dict(cached, fetchedAtMs=epoch_ms(cached.get("fetchedAtMs"))) + +def recent(part: dict[str, Any] | None, max_age: float) -> bool: + if part is None: + return False + age = time_module.time() - part["fetchedAtMs"] / 1000 + return 0 <= age < max_age + + +# Refreshes can overlap (the panel's own and a sooner retry), so an answer +# only replaces one fetched before it. +def remember(path: Path, part: dict[str, Any]) -> None: + try: + current = read_json(path) or {} + if epoch_ms(current.get("fetchedAtMs")) <= part["fetchedAtMs"]: + write_json(path, part) + except OSError as exc: + print(f"omarchy-agent-usage-fireworks: could not write {path.name} ({exc})", file=sys.stderr) + + +def now_ms() -> int: + # Floored, never rounded: a stamp half a millisecond ahead reads as a clock + # that went back and fails the 0 <= age guard on the very next run. + return int(time_module.time() * 1000) + + +# ---------------------------------------------------------------- record + + +def scan(api_base_url: str, auth_path: Path, force: bool = False, limits_only: bool = False) -> dict[str, Any]: + config = read_config() + api_key, account_id = credentials(auth_path, config) + # No key anywhere: Fireworks isn't used here, and the panel skips it. + if not api_key: + return base_record(tierLabel="") + + account_id = normalize_account_id(account_id) if account_id else "" + digest = credential_digest(api_key) + usage_cache = cache_root() / f"fireworks-usage-{digest}.json" + balance_cache = cache_root() / f"fireworks-limits-{digest}.json" + usage_part = cached_part(usage_cache, account_id) + balance_part = cached_part(balance_cache, account_id) + + # The usage rows are this record's stats, so they are reused like a local + # scan: briefly on a normal run, for 15 minutes on the panel's near-limit + # refresh. That refresh is about the balance, which comes from separate + # calls, so it is probed like a normal run. --force reuses nothing. + usage_age = 0 if force else (LIMITS_ONLY_REUSE_SECONDS if limits_only else SCAN_REUSE_SECONDS) + balance_age = 0 if force else PROBE_MIN_INTERVAL_SECONDS + usage_fresh = recent(usage_part, usage_age) + balance_fresh = recent(balance_part, balance_age) + + errors: list[FireworksError] = [] + failed: set[str] = set() + today = datetime.now().astimezone().date() + if not (usage_fresh and balance_fresh): + client = FireworksClient(api_key, api_base_url) + account: dict[str, Any] = {} + try: + if not account_id: + account_id, account = client.discover_account() + # A key that now reaches another account starts over. + if usage_part and normalize_account_id(str(usage_part.get("accountId") or "")) != account_id: + usage_part, usage_fresh = None, False + if balance_part and normalize_account_id(str(balance_part.get("accountId") or "")) != account_id: + balance_part, balance_fresh = None, False + except FireworksError as error: + errors.append(error) + failed.update(name for name, fresh in (("usage", usage_fresh), ("balance", balance_fresh)) if not fresh) + else: + if not usage_fresh: + try: + usage = client.usage(account_id, today - timedelta(days=29), today + timedelta(days=1)) + usage_part = {"accountId": account_id, "fetchedAtMs": now_ms(), "serverlessCosts": usage.get("serverlessCosts") or []} + remember(usage_cache, usage_part) + except FireworksError as error: + errors.append(error) + failed.add("usage") + if not balance_fresh: + try: + balance_part = {"accountId": account_id, "fetchedAtMs": now_ms(), "balance": current_balance(client, account_id, account, config)} + remember(balance_cache, balance_part) + except FireworksError as error: + errors.append(error) + failed.add("balance") + + record = base_record(ready=True) + # Kept rows are summarized against today, so a day-old answer still files + # its tokens under the right days. + if usage_part is not None: + rows = usage_part.get("serverlessCosts") + record.update(summarize_usage({"serverlessCosts": rows if isinstance(rows, list) else []}, today)) + record["hasLocalStats"] = True + if balance_part is not None and isinstance(balance_part.get("balance"), dict): + record["balance"] = balance_part["balance"] + kept = [part for part in (usage_part, balance_part) if part is not None] + if kept: + record["limitsFetchedAt"] = min(part["fetchedAtMs"] for part in kept) + # Anything that didn't come back fresh marks the record stale, kept or not; + # the panel shows the kept numbers' age on hover. + record["limitsStale"] = bool(failed) + if any(isinstance(error, FireworksTransportError) for error in errors): + # Usually the seconds after login before the network is up; the panel + # tries once more soon instead of waiting out its refresh interval. + record["retryAdvised"] = True + + if any(isinstance(error, FireworksAuthError) for error in errors): + record["usageStatusText"] = "Sign-in expired" + record["authHelpText"] = "Fireworks rejected the API key. " + AUTH_HELP + elif "usage" in failed and usage_part is None: + record["usageStatusText"] = "Fireworks unavailable" + record["authHelpText"] = str(errors[0]) + elif "balance" in failed and balance_part is None: + record["usageStatusText"] = "Balance unavailable" + record["authHelpText"] = str(errors[0]) + elif errors: + record["authHelpText"] = str(errors[0]) return record def main() -> int: parser = argparse.ArgumentParser(description="Print the Fireworks usage record as JSON") - # Stats and balance come from the same few API calls, so there is no cache - # to force past and no faster limits-only path. The flags exist so every - # collector accepts the same invocation. - parser.add_argument("--force", action="store_true") - parser.add_argument("--limits-only", action="store_true") + parser.add_argument("--force", action="store_true", help="re-fetch usage and balance, ignoring caches") + parser.add_argument("--limits-only", action="store_true", help="reuse any recent usage answer; only the balance must be fresh") parser.add_argument("--auth-path", default=os.environ.get("FIREWORKS_AUTH_PATH", "~/.fireworks/auth.ini")) parser.add_argument("--api-base-url", default=os.environ.get("FIREWORKS_API_BASE_URL", API_BASE_URL)) args = parser.parse_args() try: - record = scan(args.api_base_url, Path(args.auth_path).expanduser()) - except FireworksError as error: - record = base_record(usageStatusText="Fireworks unavailable", authHelpText=str(error)) + record = scan(args.api_base_url, Path(args.auth_path).expanduser(), args.force, args.limits_only) except Exception as error: record = base_record(usageStatusText="Fireworks unavailable", authHelpText="Fireworks usage scan failed") print(f"omarchy-agent-usage-fireworks: {type(error).__name__}", file=sys.stderr) diff --git a/test/shell.d/agent-usage-fireworks-scanner-test.sh b/test/shell.d/agent-usage-fireworks-scanner-test.sh index 6440ce1164f..8f9bfda0c74 100755 --- a/test/shell.d/agent-usage-fireworks-scanner-test.sh +++ b/test/shell.d/agent-usage-fireworks-scanner-test.sh @@ -7,6 +7,8 @@ require_command python3 TEST_HOME=$(mktemp -d) trap 'rm -rf "$TEST_HOME"' EXIT +# The collector caches its last good answers; keep them out of the real cache. +export XDG_CACHE_HOME="$TEST_HOME/.cache" auth_file="$TEST_HOME/auth.ini" cat >"$auth_file" <<'EOF' @@ -33,6 +35,14 @@ no_key=$(HOME="$TEST_HOME" XDG_CONFIG_HOME="$TEST_HOME/.config" XDG_DATA_HOME="$ fail "Fireworks collector prints a valid record without credentials" "$no_key" pass "Fireworks collector prints a valid record without credentials" +# An unused machine is not trouble: no status, no help, no balance. +[[ $(jq -c '{usageStatusText, authHelpText, tierLabel, limits, balance, limitsStale, limitsFetchedAt}' <<<"$no_key") == '{"usageStatusText":"","authHelpText":"","tierLabel":"","limits":[],"balance":null,"limitsStale":false,"limitsFetchedAt":0}' ]] || + fail "Fireworks collector reports no trouble on an unused machine" "$no_key" +pass "Fireworks collector reports no trouble on an unused machine" +[[ $(jq -c '[.todayPrompts, .todaySessions, .todayTotalTokens, .todayTokensByModel, .recentDays, .totalPrompts, .totalSessions, .activeDays, .activeDates, .modelUsage]' <<<"$no_key") == '[0,0,0,{},[],0,0,0,[],{}]' ]] || + fail "Fireworks collector prints the full empty stats block on an unused machine" "$no_key" +pass "Fireworks collector prints the full empty stats block on an unused machine" + result=$(python3 - "$ROOT/bin/omarchy-agent-usage-fireworks" "$auth_file" "$TEST_HOME/.config" "$TEST_HOME/.local/share" "$TEST_HOME/.pi/agent" <<'PY' import importlib.machinery import importlib.util @@ -153,7 +163,7 @@ class LiveBalanceClient(WorkingClient): os.environ["FIREWORKS_API_KEY"] = "fw_test" scanner.FireworksClient = WorkingClient -record = scanner.scan("https://example.invalid", auth_path) +record = scanner.scan("https://example.invalid", auth_path, force=True) summary["record"] = { "schemaVersion": record["schemaVersion"], "id": record["id"], @@ -184,16 +194,29 @@ summary["paginationMerges"] = ( ) scanner.FireworksClient = LiveBalanceClient -live = scanner.scan("https://example.invalid", auth_path) +live = scanner.scan("https://example.invalid", auth_path, force=True) summary["liveBalance"] = live["balance"] +# A failed balance check keeps the last good balance, stale, with the status +# left to limitsStale; with nothing kept it says the balance is unavailable. scanner.FireworksClient = BalanceFailureClient -scanned = scanner.scan("https://example.invalid", auth_path) +kept = scanner.scan("https://example.invalid", auth_path, force=True) +summary["balanceFailureKeepsBalance"] = ( + kept["balance"] == live["balance"] + and kept["limitsStale"] is True + and kept["usageStatusText"] == "" + and kept["authHelpText"] == "Billing scope denied" + and kept.get("retryAdvised") is None +) +for cached in scanner.cache_root().glob("fireworks-limits-*.json"): + cached.unlink() +scanned = scanner.scan("https://example.invalid", auth_path, force=True) summary["balanceFailurePreservesTokens"] = ( scanned["ready"] is True and "balance" not in scanned and scanned["modelUsage"]["kimi-k2.5"]["outputTokens"] == 120 and scanned["usageStatusText"] == "Balance unavailable" + and scanned["limitsStale"] is True ) # East of Greenwich, a local-midnight bucket starts on the previous UTC date; @@ -251,6 +274,10 @@ pass "Fireworks collector estimates the balance from configured funding" fail "Fireworks collector prefers the live getBalance ledger when the key can read it" "$result" pass "Fireworks collector prefers the live getBalance ledger when the key can read it" +[[ $(jq -r '.balanceFailureKeepsBalance' <<<"$result") == "true" ]] || + fail "Fireworks collector keeps the last balance, stale, when the check fails" "$result" +pass "Fireworks collector keeps the last balance, stale, when the check fails" + [[ $(jq -r '.balanceFailurePreservesTokens' <<<"$result") == "true" ]] || fail "Fireworks collector preserves tokens when balance lookup fails" "$result" pass "Fireworks collector preserves tokens when balance lookup fails" @@ -278,3 +305,164 @@ pass "Fireworks collector treats missing env references as unresolved" [[ $(jq -r '.piCommandKeySkipped' <<<"$result") == "true" ]] || fail "Fireworks collector skips command-lookup pi keys and falls through" "$result" pass "Fireworks collector skips command-lookup pi keys and falls through" + +# The cache contract: last good answers per key, reused briefly, kept stale +# past a failed check. Each check starts from an empty cache. +contract=$(python3 - "$ROOT/bin/omarchy-agent-usage-fireworks" "$auth_file" "$TEST_HOME" <<'PY' +import importlib.machinery, importlib.util, json, os, shutil, sys, time +from decimal import Decimal +from pathlib import Path + +home = Path(sys.argv[3]) +os.environ["XDG_CONFIG_HOME"] = str(home / ".config") +os.environ["XDG_DATA_HOME"] = str(home / ".local/share") +os.environ["PI_CODING_AGENT_DIR"] = str(home / ".pi/agent") +os.environ["TZ"] = "UTC" +time.tzset() +os.environ.pop("FIREWORKS_ACCOUNT_ID", None) +os.environ["FIREWORKS_API_KEY"] = "fw_test" +auth_path = Path(sys.argv[2]) + +loader = importlib.machinery.SourceFileLoader("fireworks_contract", sys.argv[1]) +spec = importlib.util.spec_from_loader(loader.name, loader) +scanner = importlib.util.module_from_spec(spec) +loader.exec_module(scanner) + +today = time.strftime("%Y-%m-%d") +ROWS = {"serverlessCosts": [{"startTime": today + "T00:00:00Z", "promptTokens": "100", "completionTokens": "20", "group": {"model_name": "accounts/fireworks/models/kimi-k2p5"}}]} +LIVE = {"balance": {"units": "12", "nanos": 500000000}} +state = {"usage": ROWS, "balance": LIVE, "calls": []} + +class StubClient: + def __init__(self, api_key, base_url): + self.api_key = api_key + def answer(self, kind): + state["calls"].append(kind) + value = state[kind] + if isinstance(value, Exception): + raise value + return value + def request(self, path, query=None, body=None): + return self.answer("balance") + def usage(self, account_id, start_day, end_day): + return self.answer("usage") + def account(self, account_id): + return {} + def spent(self, account_id, start_at, end_at): + return Decimal("0") +scanner.FireworksClient = StubClient + +def fresh_cache(): + shutil.rmtree(scanner.cache_root(), ignore_errors=True) + state.update(usage=ROWS, balance=LIVE, calls=[]) + +def run(**flags): + return scanner.scan("https://example.invalid", auth_path, **flags) + +def pick(r): + picked = {k: r.get(k) for k in ("ready", "usageStatusText", "authHelpText", "limitsStale", "retryAdvised", "todayTotalTokens", "hasLocalStats")} + picked["balance"] = (r.get("balance") or {}).get("remaining") + picked["fetched"] = r.get("limitsFetchedAt", 0) > 0 + return picked + +out = {} +UNREACHABLE = scanner.FireworksTransportError("Could not reach the Fireworks API") +SERVER = scanner.FireworksError("Fireworks API returned HTTP 500") +REFUSED = scanner.FireworksAuthError("Fireworks rejected the API key") + +fresh_cache() +out["fresh"] = pick(run(force=True)) + +# Usage and balance both fail: the last good stats and balance stand, stale. +state.update(usage=UNREACHABLE, balance=UNREACHABLE) +out["stale"] = pick(run(force=True)) + +# A server error is no reason for a sooner retry. +state.update(usage=SERVER, balance=SERVER) +out["serverError"] = pick(run(force=True)) + +# Nothing kept: the stats source is unavailable. +fresh_cache() +state.update(usage=UNREACHABLE, balance=UNREACHABLE) +out["nothingKept"] = pick(run(force=True)) + +# A refused key asks for a new one. +fresh_cache() +run(force=True) +state.update(usage=REFUSED, balance=REFUSED) +out["refused"] = pick(run(force=True)) + +# Answers are reused briefly and --force reuses nothing; --limits-only keeps +# the usage rows for 15 minutes but re-checks a balance older than the probe +# interval. +fresh_cache() +run() +first = list(state["calls"]) +state["calls"].clear() +run() +reused = list(state["calls"]) +run(force=True) +forced = list(state["calls"]) +state["calls"].clear() +for path in scanner.cache_root().glob("fireworks-*.json"): + cached = json.loads(path.read_text()) + cached["fetchedAtMs"] -= 60_000 + path.write_text(json.dumps(cached)) +run(limits_only=True) +limits_only = list(state["calls"]) +state["calls"].clear() +run() +normal = list(state["calls"]) +out["reuse"] = {"first": first, "reused": reused, "forced": forced, "limitsOnly": limits_only, "normal": normal} + +# A cache dated in the future (the clock went back) is kept but never counts +# as recent. +fresh_cache() +run(force=True) +for path in scanner.cache_root().glob("fireworks-*.json"): + cached = json.loads(path.read_text()) + cached["fetchedAtMs"] += 3_600_000 + path.write_text(json.dumps(cached)) +state["calls"].clear() +run() +out["futureCache"] = sorted(set(state["calls"])) + +# Another key, or the same key now pointed at another account, never shows +# these numbers. +fresh_cache() +run(force=True) +state.update(usage=UNREACHABLE, balance=UNREACHABLE) +os.environ["FIREWORKS_API_KEY"] = "fw_other" +out["otherKey"] = pick(run(force=True)) +os.environ["FIREWORKS_API_KEY"] = "fw_test" +os.environ["FIREWORKS_ACCOUNT_ID"] = "elsewhere" +out["otherAccount"] = pick(run(force=True)) +os.environ.pop("FIREWORKS_ACCOUNT_ID") + +# Untrusted token fields skip, never crash or inflate. +out["robust"] = scanner.summarize_usage({"serverlessCosts": [ + {"startTime": today + "T00:00:00Z", "promptTokens": True, "completionTokens": "NaN", "group": {"model_name": "a"}}, + {"startTime": today + "T00:00:00Z", "promptTokens": 1e300, "completionTokens": "inf", "group": {"model_name": "b"}}, + {"startTime": today + "T00:00:00Z", "promptTokens": -5, "completionTokens": 7, "group": {"model_name": "c"}}, +]})["modelUsage"] +print(json.dumps(out, separators=(",", ":"))) +PY +) + +check_contract() { + local filter="$1" expected="$2" label="$3" actual + actual=$(jq -c "$filter" <<<"$contract") + [[ $actual == "$expected" ]] || fail "$label (got: $actual)" "$contract" + pass "$label" +} + +check_contract '.fresh' '{"ready":true,"usageStatusText":"","authHelpText":"","limitsStale":false,"retryAdvised":null,"todayTotalTokens":120,"hasLocalStats":true,"balance":12.5,"fetched":true}' "Fireworks collector reports a fresh check as healthy" +check_contract '.stale' '{"ready":true,"usageStatusText":"","authHelpText":"Could not reach the Fireworks API","limitsStale":true,"retryAdvised":true,"todayTotalTokens":120,"hasLocalStats":true,"balance":12.5,"fetched":true}' "Fireworks collector keeps the last stats and balance, stale, past an unreachable API" +check_contract '.serverError | {limitsStale, retryAdvised, usageStatusText}' '{"limitsStale":true,"retryAdvised":null,"usageStatusText":""}' "Fireworks collector asks for a sooner retry only when the API is unreachable" +check_contract '.nothingKept | {ready, usageStatusText, limitsStale, retryAdvised, balance}' '{"ready":true,"usageStatusText":"Fireworks unavailable","limitsStale":true,"retryAdvised":true,"balance":null}' "Fireworks collector says it is unavailable when nothing was kept" +check_contract '.refused | {usageStatusText, limitsStale, balance}' '{"usageStatusText":"Sign-in expired","limitsStale":true,"balance":12.5}' "Fireworks collector reads a rejected key as an expired sign-in" +check_contract '.reuse' '{"first":["usage","balance"],"reused":[],"forced":["usage","balance"],"limitsOnly":["balance"],"normal":["usage"]}' "Fireworks collector reuses recent answers and --force re-fetches" +check_contract '.futureCache' '["balance","usage"]' "Fireworks collector never treats a future-dated cache as recent" +check_contract '.otherKey | {balance, todayTotalTokens, usageStatusText}' '{"balance":null,"todayTotalTokens":0,"usageStatusText":"Fireworks unavailable"}' "Fireworks collector keeps each key's answers apart" +check_contract '.otherAccount | {balance, todayTotalTokens}' '{"balance":null,"todayTotalTokens":0}' "Fireworks collector keeps each account's answers apart" +check_contract '.robust' '{"c":{"inputTokens":0,"outputTokens":7,"cacheReadInputTokens":0,"cacheCreationInputTokens":0}}' "Fireworks collector skips untrusted token fields" From 9f0f20bfc10081494f81ec9e356bcfeb2a52934a Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:13:23 -0400 Subject: [PATCH 061/546] Bring the Hermes collector in line with the collector contract A machine without Hermes state prints an empty status with ready false, and an unreadable state.db reads as Hermes unavailable with the detail in the help text. The scan is cached per day under a lock, honouring --force and --limits-only, every active date is emitted, and broken token counts skip their row. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-hermes | 215 +++++++++++++++++++----- test/shell.d/agent-usage-hermes-test.sh | 120 +++++++++++++ 2 files changed, 296 insertions(+), 39 deletions(-) diff --git a/bin/omarchy-agent-usage-hermes b/bin/omarchy-agent-usage-hermes index 33051b27e92..30c26315dbb 100755 --- a/bin/omarchy-agent-usage-hermes +++ b/bin/omarchy-agent-usage-hermes @@ -14,15 +14,35 @@ them as first-class sessions and so does this record. """ import argparse +import fcntl +import hashlib import json +import math import os -import shutil import sqlite3 +import sys +import tempfile +import time from datetime import datetime, timedelta, timezone from pathlib import Path RECENT_DAYS = 7 +# A scan this recent is only reused to dedup concurrent collector runs; +# --limits-only promises only fresh limits, and Hermes has none, so it may +# reuse the scan for up to 15 minutes. --force reuses nothing. +SCAN_REUSE_SECONDS = 20 +LIMITS_ONLY_REUSE_SECONDS = 900 + +# Token counts beyond this are a corrupt row, not usage. +MAX_TOKENS = 10**12 + + +def cache_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 + # Hermes anchors named profiles under ~/.hermes whatever HERMES_HOME says, so # they are looked for there. A profile linked to the home counts once. @@ -44,22 +64,40 @@ def state_paths(): return paths +# Counts from disk are untrusted: a negative, non-finite, absurdly large or +# non-numeric value counts as nothing. +def tokens(value): + if isinstance(value, bool): + return 0 + try: + parsed = float(value or 0) + except (TypeError, ValueError, OverflowError): + return 0 + if not math.isfinite(parsed) or parsed <= 0 or parsed > MAX_TOKENS: + return 0 + return int(parsed) + + def day_start(day): return datetime(day.year, day.month, day.day).astimezone().timestamp() +def empty_bucket(): + return {"inputTokens": 0, "outputTokens": 0, "cacheReadInputTokens": 0, "cacheCreationInputTokens": 0} + + # 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"): + # A start time SQLite can't read as an epoch dates to NULL and names no day. + for date in conn.execute("SELECT DISTINCT date(started_at, 'unixepoch', 'localtime') FROM sessions WHERE started_at > 0"): if date[0]: totals["dates"].add(date[0]) @@ -73,72 +111,171 @@ def scan(path, totals, today, week): 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) + counts = [tokens(uncached), tokens(output), tokens(cache_read), tokens(cache_write)] + total = sum(counts) + if not total: + continue + model = str(model or "hermes") + bucket = totals["models"].setdefault(model, empty_bucket()) + for field, count in zip(("inputTokens", "outputTokens", "cacheReadInputTokens", "cacheCreationInputTokens"), counts): + bucket[field] += count if day in totals["byDay"]: - totals["byDay"][day] += tokens - if day == today.isoformat() and tokens: - totals["todayByModel"][model] = totals["todayByModel"].get(model, 0) + tokens + totals["byDay"][day] += total + if day == today.isoformat(): + totals["todayByModel"][model] = totals["todayByModel"].get(model, 0) + total finally: conn.close() -def collect(): - now = datetime.now().astimezone() - today = now.date() +# The stats from every state.db, with the paths that couldn't be read. +def collect_stats(paths): + today = datetime.now().astimezone().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 + unreadable = [] for path in paths: try: scan(path, totals, today, week) - read += 1 - except sqlite3.Error: - status = "Hermes state unreadable" + except sqlite3.Error as error: + unreadable.append(f"{path}: {error}") 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, + stats = { "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()], + "recentDays": [{"date": day, "messageCount": count} for day, count in totals["byDay"].items()], "totalPrompts": totals["prompts"], "totalSessions": totals["sessions"], "activeDays": len(dates), - "activeDates": dates[-400:], + # Every date travels: merging several machines' snapshots needs their + # union, which a count alone cannot give. + "activeDates": dates, "modelUsage": totals["models"], - "limits": [], + } + return stats, unreadable + + +def read_fresh_json(path, max_age_seconds): + if max_age_seconds <= 0: + 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: + 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, payload): + # A temp name unique to this writer: several collector runs can overlap. + 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 + + +# The scan is cached per day and per set of databases, rebuilt under a lock +# so concurrent runs wait for one scan. A scan that couldn't read every +# database isn't cached. +def cached_stats(paths, max_age): + scan_date = datetime.now().astimezone().date().isoformat() + + def usable(cached): + if isinstance(cached, dict) and cached.get("schemaVersion") == 1 and cached.get("scanDate") == scan_date and isinstance(cached.get("stats"), dict): + return cached["stats"] + return None + + try: + digest = hashlib.sha1("\n".join(str(path.resolve()) for path in paths).encode("utf-8")).hexdigest()[:16] + root = cache_root() + cache_file, lock_file = root / f"hermes-scan-{digest}.json", root / f"hermes-scan-{digest}.lock" + stats = usable(read_fresh_json(cache_file, max_age)) + if stats is not None: + return stats, [] + with lock_file.open("w") as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + stats = usable(read_fresh_json(cache_file, max_age)) + if stats is not None: + return stats, [] + stats, unreadable = collect_stats(paths) + if not unreadable: + try: + write_json(cache_file, {"schemaVersion": 1, "scanDate": scan_date, "stats": stats}) + except OSError as exc: + print(f"omarchy-agent-usage-hermes: could not write usage cache ({exc})", file=sys.stderr) + return stats, unreadable + except OSError as exc: + # The cache is an optimization; the record is the contract. + print(f"omarchy-agent-usage-hermes: cache unavailable ({exc}); scanning directly", file=sys.stderr) + return collect_stats(paths) + + +def collect(force=False, limits_only=False): + record = { + "schemaVersion": 1, + "id": "hermes", + "name": "Hermes", + "updatedAt": datetime.now(timezone.utc).isoformat(), + "ready": False, + "hasLocalStats": False, + "hasPromptStats": True, "tierLabel": "", - "usageStatusText": status, + "usageStatusText": "", "authHelpText": "", + "limits": [], + "limitsStale": False, + "limitsFetchedAt": 0, + "todayPrompts": 0, + "todaySessions": 0, + "todayTotalTokens": 0, + "todayTokensByModel": {}, + "recentDays": [], + "totalPrompts": 0, + "totalSessions": 0, + "activeDays": 0, + "activeDates": [], + "modelUsage": {}, } + # No state.db: Hermes isn't used on this machine, which is not trouble -- + # with sync on, another machine's Hermes still shows through its stats. + paths = state_paths() + if not paths: + return record + + max_age = 0 if force else (LIMITS_ONLY_REUSE_SECONDS if limits_only else SCAN_REUSE_SECONDS) + stats, unreadable = cached_stats(paths, max_age) + record.update(stats) + record["ready"] = True + record["hasLocalStats"] = len(unreadable) < len(paths) + if unreadable: + record["usageStatusText"] = "Hermes unavailable" + record["authHelpText"] = "Couldn't read Hermes' state.db (" + "; ".join(unreadable) + ")." + return record + 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=(",", ":"))) + parser = argparse.ArgumentParser(description="Print the Hermes Agent usage record as JSON") + parser.add_argument("--force", action="store_true", help="rescan every state.db, ignoring the cache") + parser.add_argument("--limits-only", action="store_true", help="reuse any recent scan; Hermes has no limits to probe") + args = parser.parse_args() + print(json.dumps(collect(args.force, args.limits_only), separators=(",", ":"))) if __name__ == "__main__": diff --git a/test/shell.d/agent-usage-hermes-test.sh b/test/shell.d/agent-usage-hermes-test.sh index 9fcd0fcfc47..a5edf26eac4 100755 --- a/test/shell.d/agent-usage-hermes-test.sh +++ b/test/shell.d/agent-usage-hermes-test.sh @@ -7,6 +7,8 @@ require_command python3 TEST_HOME=$(mktemp -d) trap 'rm -rf "$TEST_HOME"' EXIT +# The collector caches its scan; keep it out of the real cache. +export XDG_CACHE_HOME="$TEST_HOME/.cache" # 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 @@ -95,3 +97,121 @@ 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" + +# No state.db is an unused machine, not trouble: with sync on, another +# machine's Hermes must not show red here. +[[ $(jq -c '{ready, usageStatusText, authHelpText, limitsStale, limitsFetchedAt, hasLocalStats}' <<<"$record") == '{"ready":false,"usageStatusText":"","authHelpText":"","limitsStale":false,"limitsFetchedAt":0,"hasLocalStats":false}' ]] || + fail "Hermes collector reports no trouble without Hermes state" "$record" +pass "Hermes collector reports no trouble without Hermes state" +[[ $(jq -c '[.todayPrompts, .todaySessions, .todayTotalTokens, .todayTokensByModel, .recentDays, .totalPrompts, .totalSessions, .activeDays, .activeDates, .modelUsage]' <<<"$record") == '[0,0,0,{},[],0,0,0,[],{}]' ]] || + fail "Hermes collector prints the full empty stats block without Hermes state" "$record" +pass "Hermes collector prints the full empty stats block without Hermes state" + +# Hermes on the machine is ready whether or not the hermes command is on PATH. +record=$(HOME="$TEST_HOME" HERMES_HOME="" PATH=/usr/bin:/bin "$ROOT/bin/omarchy-agent-usage-hermes" --force) +[[ $(jq -c '{ready, usageStatusText, authHelpText}' <<<"$record") == '{"ready":true,"usageStatusText":"","authHelpText":""}' ]] || + fail "Hermes collector is ready once Hermes state exists" "$record" +pass "Hermes collector is ready once Hermes state exists" + +# An unreadable state.db says Hermes is unavailable, with the detail in the +# help text, and the readable profiles still count. +bad_home=$(mktemp -d) +make_state "$bad_home/.hermes/profiles/work/state.db" gpt-6-sol:1000 +mkdir -p "$bad_home/.hermes" +printf 'not a database at all, just text that is long enough to read' >"$bad_home/.hermes/state.db" +record=$(HOME="$bad_home" HERMES_HOME="" XDG_CACHE_HOME="$bad_home/.cache" "$ROOT/bin/omarchy-agent-usage-hermes" --force) +[[ $(jq -c '{ready, usageStatusText, hasLocalStats, totalSessions}' <<<"$record") == '{"ready":true,"usageStatusText":"Hermes unavailable","hasLocalStats":true,"totalSessions":1}' ]] || + fail "Hermes collector reports an unreadable state.db" "$record" +[[ $(jq -r .authHelpText <<<"$record") == "Couldn't read Hermes' state.db ($bad_home/.hermes/state.db: "* ]] || + fail "Hermes collector names the unreadable state.db" "$record" +pass "Hermes collector reports an unreadable state.db" +[[ -z $(find "$bad_home/.cache" -name 'hermes-scan-*.json' 2>/dev/null) ]] || + fail "Hermes collector doesn't cache a scan that missed a database" +pass "Hermes collector doesn't cache a scan that missed a database" +rm -rf "$bad_home" + +# The scan is reused for SCAN_REUSE_SECONDS, for LIMITS_ONLY_REUSE_SECONDS on +# --limits-only, and never on --force. +cache_home=$(mktemp -d) +make_state "$cache_home/.hermes/state.db" gpt-6-astra:10 +run_cached() { + HOME="$cache_home" HERMES_HOME="" XDG_CACHE_HOME="$cache_home/.cache" "$ROOT/bin/omarchy-agent-usage-hermes" "$@" | jq -r .totalSessions +} +add_session() { + python3 - "$cache_home/.hermes/state.db" "$1" <<'PY' +import sqlite3, sys, time +conn = sqlite3.connect(sys.argv[1]) +conn.execute("INSERT INTO sessions VALUES (?, ?)", (sys.argv[2], time.time())) +conn.commit() +PY +} +age_cache() { + local file + for file in "$cache_home"/.cache/omarchy/agent-usage/hermes-scan-*.json; do + touch -d "-$1 seconds" "$file" + done +} +[[ $(run_cached) == "1" ]] || fail "Hermes collector scans on first run" +add_session session-2 +[[ $(run_cached) == "1" ]] || fail "Hermes collector reuses a scan from seconds ago" +pass "Hermes collector reuses a scan from seconds ago" +[[ $(run_cached --force) == "2" ]] || fail "Hermes collector rescans on --force" +pass "Hermes collector rescans on --force" +add_session session-3 +age_cache 120 +[[ $(run_cached --limits-only) == "2" ]] || fail "Hermes collector reuses a recent scan on --limits-only" +pass "Hermes collector reuses a recent scan on --limits-only" +[[ $(run_cached) == "3" ]] || fail "Hermes collector rescans once the reuse window passes" +pass "Hermes collector rescans once the reuse window passes" +add_session session-4 +age_cache -3600 +[[ $(run_cached --limits-only) == "4" ]] || fail "Hermes collector never reuses a future-dated scan" +pass "Hermes collector never reuses a future-dated scan" + +# The scan is rebuilt under its lock: a concurrent run waits for it. +lock_file=$(find "$cache_home/.cache/omarchy/agent-usage" -name 'hermes-scan-*.lock' | head -n 1) +[[ -n $lock_file ]] || fail "Hermes collector keeps a scan lock" +waited=$(python3 - "$lock_file" "$ROOT/bin/omarchy-agent-usage-hermes" "$cache_home" <<'PY' +import fcntl, os, subprocess, sys, time +lock = open(sys.argv[1], "w") +fcntl.flock(lock, fcntl.LOCK_EX) +env = dict(os.environ, HOME=sys.argv[3], HERMES_HOME="", XDG_CACHE_HOME=sys.argv[3] + "/.cache") +started = time.time() +proc = subprocess.Popen([sys.argv[2], "--force"], env=env, stdout=subprocess.DEVNULL) +time.sleep(1.2) +still_running = proc.poll() is None +fcntl.flock(lock, fcntl.LOCK_UN) +proc.wait() +print("true" if still_running else "false") +PY +) +[[ $waited == "true" ]] || fail "Hermes collector rebuilds the scan under its lock" +pass "Hermes collector rebuilds the scan under its lock" +rm -rf "$cache_home" + +# Every active date travels, not only the latest 400, and broken token counts +# or start times skip their row instead of inflating the record. +dates_home=$(mktemp -d) +make_state "$dates_home/.hermes/state.db" +python3 - "$dates_home/.hermes/state.db" <<'PY' +import sqlite3, sys, time +conn = sqlite3.connect(sys.argv[1]) +day = 86400 +now = time.time() +conn.executemany("INSERT INTO sessions VALUES (?, ?)", [(f"old-{n}", now - (n + 1) * day) for n in range(450)]) +conn.execute("INSERT INTO sessions VALUES ('broken', 'not a time')") +conn.executemany("INSERT INTO session_model_usage VALUES (?, ?, ?, ?, ?, ?)", [ + ("old-1", "negative", -500, 0, 0, 0), + ("old-2", "huge", 10**15, 0, 0, 0), + ("broken", "untimed", 5, 0, 0, 0), +]) +conn.commit() +PY +record=$(HOME="$dates_home" HERMES_HOME="" XDG_CACHE_HOME="$dates_home/.cache" "$ROOT/bin/omarchy-agent-usage-hermes" --force) +rm -rf "$dates_home" +[[ $(jq -c '[.activeDays, (.activeDates | length)]' <<<"$record") == '[451,451]' ]] || + fail "Hermes collector emits every active date" "$record" +pass "Hermes collector emits every active date" +[[ $(jq -c '.modelUsage | keys' <<<"$record") == '["untimed"]' ]] || + fail "Hermes collector skips broken token counts" "$record" +pass "Hermes collector skips broken token counts" From 14a6a576fcc3fe786b2fc6b3f96f246cbc6f7215 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:13:23 -0400 Subject: [PATCH 062/546] Add Hermes Agent's mark to the agents panel The glyph is lifted from the app icons in NousResearch/hermes-agent (assets/icon-master-dark.svg and icon-master.svg, MIT), without their rounded tile: a white mark plus a dark twin for light surfaces. Co-Authored-By: Claude Opus 5.5 --- shell/plugins/agents/assets/hermes-light.svg | 4 ++++ shell/plugins/agents/assets/hermes.svg | 4 ++++ 2 files changed, 8 insertions(+) create mode 100644 shell/plugins/agents/assets/hermes-light.svg create mode 100644 shell/plugins/agents/assets/hermes.svg diff --git a/shell/plugins/agents/assets/hermes-light.svg b/shell/plugins/agents/assets/hermes-light.svg new file mode 100644 index 00000000000..de9159de9bf --- /dev/null +++ b/shell/plugins/agents/assets/hermes-light.svg @@ -0,0 +1,4 @@ + + Hermes Agent + + diff --git a/shell/plugins/agents/assets/hermes.svg b/shell/plugins/agents/assets/hermes.svg new file mode 100644 index 00000000000..0383a7624d8 --- /dev/null +++ b/shell/plugins/agents/assets/hermes.svg @@ -0,0 +1,4 @@ + + Hermes Agent + + From de5958b7b34132f661b05e3e1aaf794c4af3f7b2 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 22:56:51 -0400 Subject: [PATCH 063/546] Name Claude's scan reuse windows after the collector contract Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-claude | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/bin/omarchy-agent-usage-claude b/bin/omarchy-agent-usage-claude index e615c5836eb..0baf805db13 100755 --- a/bin/omarchy-agent-usage-claude +++ b/bin/omarchy-agent-usage-claude @@ -35,6 +35,11 @@ AGENT_NAME = "Claude Code" AUTH_HELP = "Run `claude auth login` to restore authoritative usage." USAGE_ENDPOINT = "https://api.anthropic.com/api/oauth/usage" PROBE_MIN_INTERVAL_SECONDS = 15 +# A local scan this recent is reused rather than redone: a normal run only +# dedups overlapping collector runs, while --limits-only promises fresh limits +# alone and may lean on a scan for much longer. +SCAN_REUSE_SECONDS = 20 +LIMITS_ONLY_REUSE_SECONDS = 900 # Anthropic rate-limits its usage endpoint readily. A refused re-check of # numbers this recent still describes the account; only older ones are # reported as stale. @@ -1216,11 +1221,11 @@ def main() -> int: parser = argparse.ArgumentParser() parser.add_argument("--force", action="store_true", help="rescan transcripts from scratch and re-probe limits, ignoring caches") parser.add_argument("--limits-only", action="store_true", help="reuse any recent transcript scan; only the limits probe must be fresh") - parser.add_argument("--cache-seconds", type=float, default=20) + parser.add_argument("--cache-seconds", type=float, default=SCAN_REUSE_SECONDS) args = parser.parse_args() claude_dir = config_dir() - scan_age = 0 if args.force else (900 if args.limits_only else args.cache_seconds) + scan_age = 0 if args.force else (LIMITS_ONLY_REUSE_SECONDS if args.limits_only else args.cache_seconds) stats = cached_scan(claude_dir / "projects", scan_age, use_index=not args.force) if number(stats.get("totalPrompts")) <= 0: From c7b4dbf83425fa4a40d23663a2b402ca324077f3 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 22:57:48 -0400 Subject: [PATCH 064/546] Count pi usage from PI_CODING_AGENT_DIR and PI_CODING_AGENT_SESSION_DIR The Claude and Codex collectors scanned only ~/.pi and ~/.omp, so a pi whose agent tree or sessions were moved by its documented overrides had its usage dropped from the panel. Roots are canonicalised, with Codex session and parentSession paths through realpath, so an override aliasing a default root through a symlink is counted once and forks still find their parent. A root nested inside another is dropped, since the outer one already reads it, and the Codex scan cache is keyed by the overrides too. Builds on #14550. Co-Authored-By: Claude Opus 5.5 Co-authored-by: Omarchybot <317366263+omarchybot@users.noreply.github.com> --- bin/omarchy-agent-usage-claude | 12 +++++- bin/omarchy-agent-usage-codex | 22 +++++++--- .../agent-usage-claude-scanner-test.sh | 34 +++++++++++++++ .../shell.d/agent-usage-codex-scanner-test.sh | 41 +++++++++++++++++++ 4 files changed, 103 insertions(+), 6 deletions(-) diff --git a/bin/omarchy-agent-usage-claude b/bin/omarchy-agent-usage-claude index 0baf805db13..de7ccd1dd9a 100755 --- a/bin/omarchy-agent-usage-claude +++ b/bin/omarchy-agent-usage-claude @@ -541,7 +541,17 @@ def pi_session_roots() -> list[Path]: except OSError: # No profiles directory, or it is unreadable: the default root stands. pass - return roots + # pi moves its whole agent tree with PI_CODING_AGENT_DIR and its sessions + # alone with PI_CODING_AGENT_SESSION_DIR, both expanding a leading ~. + if os.environ.get("PI_CODING_AGENT_DIR"): + roots.append(Path(os.environ["PI_CODING_AGENT_DIR"]).expanduser() / "sessions") + if os.environ.get("PI_CODING_AGENT_SESSION_DIR"): + roots.append(Path(os.environ["PI_CODING_AGENT_SESSION_DIR"]).expanduser()) + # Canonical roots, so an override that aliases a default root through a + # symlink is scanned under one spelling. A root inside another is already + # read with it, so it is dropped rather than counted twice. + canonical = list(dict.fromkeys(root.resolve() for root in roots)) + return [root for root in canonical if not any(other != root and root.is_relative_to(other) for other in canonical)] def scan_pi_usage(max_age_seconds: float) -> dict[str, Any] | None: diff --git a/bin/omarchy-agent-usage-codex b/bin/omarchy-agent-usage-codex index c76359943c0..1374c1a1490 100755 --- a/bin/omarchy-agent-usage-codex +++ b/bin/omarchy-agent-usage-codex @@ -230,7 +230,7 @@ def pi_parent_session(path): parent = header.get("parentSession") if not isinstance(parent, str) or not parent: return None - return os.path.abspath(os.path.expanduser(parent)) + return os.path.realpath(os.path.expanduser(parent)) # `omp --profile=` (and pi's equivalent) relocates the whole agent tree @@ -249,7 +249,17 @@ def pi_session_roots(): except OSError: # No profiles directory, or it is unreadable: the default root stands. pass - return roots + # pi moves its whole agent tree with PI_CODING_AGENT_DIR and its sessions + # alone with PI_CODING_AGENT_SESSION_DIR, both expanding a leading ~. + if os.environ.get("PI_CODING_AGENT_DIR"): + roots.append(Path(os.environ["PI_CODING_AGENT_DIR"]).expanduser() / "sessions") + if os.environ.get("PI_CODING_AGENT_SESSION_DIR"): + roots.append(Path(os.environ["PI_CODING_AGENT_SESSION_DIR"]).expanduser()) + # Canonical roots, so an override that aliases a default root through a + # symlink is scanned under one spelling. A root inside another is already + # read with it, so it is dropped rather than counted twice. + canonical = list(dict.fromkeys(root.resolve() for root in roots)) + return [root for root in canonical if not any(other != root and root.is_relative_to(other) for other in canonical)] def scan_pi_sessions(): @@ -284,7 +294,7 @@ def scan_pi_sessions(): continue line = event.get("data", {}).get("lines", {}).get("text", "") raw_path = event.get("data", {}).get("path", {}).get("text", "pi-session") - path = os.path.abspath(os.path.expanduser(raw_path)) + path = os.path.realpath(os.path.expanduser(raw_path)) entry = json.loads(line) except Exception: continue @@ -546,8 +556,10 @@ def scan_cache_digest(): codex_home = Path(os.environ.get("CODEX_HOME") or (Path.home() / ".codex")) db = Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) / "opencode" / "opencode.db" # The digest covers every data path the scan reads: the codex session - # roots, the opencode DB, and (via Path.home()) the pi/omp session roots. - return hashlib.sha1((str(Path.home()) + "\n" + str(codex_home) + "\n" + str(db)).encode("utf-8")).hexdigest()[:16] + # roots, the opencode DB, and the pi/omp session roots (under Path.home(), + # or wherever pi's own overrides moved them). + pi_roots = "\n".join(os.environ.get(name, "") for name in ("PI_CODING_AGENT_DIR", "PI_CODING_AGENT_SESSION_DIR")) + return hashlib.sha1((str(Path.home()) + "\n" + str(codex_home) + "\n" + str(db) + "\n" + pi_roots).encode("utf-8", "surrogatepass")).hexdigest()[:16] def scan_cache_paths(): diff --git a/test/shell.d/agent-usage-claude-scanner-test.sh b/test/shell.d/agent-usage-claude-scanner-test.sh index 533c2ad6ab1..c927842607d 100644 --- a/test/shell.d/agent-usage-claude-scanner-test.sh +++ b/test/shell.d/agent-usage-claude-scanner-test.sh @@ -9,6 +9,8 @@ require_command python3 SCRATCH=$(mktemp -d) trap 'rm -rf "$SCRATCH"' EXIT TEST_HOME=$(mktemp -d "$SCRATCH/home.XXXXXX") +# pi's directory overrides would point every scan at the developer's own sessions. +unset PI_CODING_AGENT_DIR PI_CODING_AGENT_SESSION_DIR projects="$TEST_HOME/.claude/projects/example" mkdir -p "$projects" @@ -351,6 +353,38 @@ result=$(HOME="$PI_HOME" XDG_CACHE_HOME="$PI_HOME/.cache" XDG_DATA_HOME="$PI_HOM fail "Claude collector filters pi and omp sessions to Anthropic providers" "$result" pass "Claude collector counts pi and omp subscription usage, profiles included" +# pi moves its agent tree with PI_CODING_AGENT_DIR and its sessions alone with +# PI_CODING_AGENT_SESSION_DIR, so usage under either must be counted, and an +# override reaching the default root through a symlink must count it once. +CUSTOM_PI_HOME=$(mktemp -d "$SCRATCH/home.XXXXXX") +mkdir -p "$CUSTOM_PI_HOME/.config/pi/sessions/project" "$CUSTOM_PI_HOME/pi-sessions" "$CUSTOM_PI_HOME/.pi/agent/sessions/project" +cat >"$CUSTOM_PI_HOME/.config/pi/sessions/project/agent-dir.jsonl" <"$CUSTOM_PI_HOME/pi-sessions/session-dir.jsonl" <"$CUSTOM_PI_HOME/.pi/agent/sessions/project/default.jsonl" <"$collector" chmod +x "$collector" TEST_HOME=$(mktemp -d "$SCRATCH/home.XXXXXX") +# pi's directory overrides would point every scan at the developer's own sessions. +unset PI_CODING_AGENT_DIR PI_CODING_AGENT_SESSION_DIR # A fixture home signed in to Codex, with an empty bin/ for its CLI. signed_in_home() { @@ -172,6 +174,45 @@ result=$(HOME="$GIT_HOME" CODEX_HOME="$GIT_HOME/.codex" XDG_DATA_HOME="$GIT_HOME fail "Codex collector counts pi sessions when HOME is a git checkout" "$result" pass "Codex collector counts pi sessions when HOME is a git checkout" +# pi moves its agent tree with PI_CODING_AGENT_DIR and its sessions alone with +# PI_CODING_AGENT_SESSION_DIR, so usage under either must be counted. An +# override reaching the default root through a symlink must count it once and +# still match a fork whose parent is spelled through that symlink. +CUSTOM_PI_HOME=$(signed_in_home) +cp "$TEST_HOME/bin/codex" "$CUSTOM_PI_HOME/bin/codex" +mkdir -p "$CUSTOM_PI_HOME/.config/pi/sessions/project" "$CUSTOM_PI_HOME/pi-sessions" "$CUSTOM_PI_HOME/.pi/agent/sessions/project" +cat >"$CUSTOM_PI_HOME/.config/pi/sessions/project/agent-dir.jsonl" <"$CUSTOM_PI_HOME/pi-sessions/session-dir.jsonl" <"$CUSTOM_PI_HOME/.pi/agent/sessions/project/default.jsonl" <"$CUSTOM_PI_HOME/.pi/agent/sessions/project/default-fork.jsonl" < Date: Wed, 14 Oct 2026 22:59:02 -0400 Subject: [PATCH 065/546] Count every Codex pi line that has no id pi lines without an id were all keyed by their file alone, so every one after the first in a session was dropped as a duplicate. They are keyed by their line now, so each counts once. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-codex | 11 ++++++++--- .../shell.d/agent-usage-codex-scanner-test.sh | 19 +++++++++++++++++++ 2 files changed, 27 insertions(+), 3 deletions(-) diff --git a/bin/omarchy-agent-usage-codex b/bin/omarchy-agent-usage-codex index 1374c1a1490..860c51a4449 100755 --- a/bin/omarchy-agent-usage-codex +++ b/bin/omarchy-agent-usage-codex @@ -295,19 +295,24 @@ def scan_pi_sessions(): line = event.get("data", {}).get("lines", {}).get("text", "") raw_path = event.get("data", {}).get("path", {}).get("text", "pi-session") path = os.path.realpath(os.path.expanduser(raw_path)) + line_number = event.get("data", {}).get("line_number") entry = json.loads(line) except Exception: continue - if entry.get("type") != "message": + if not isinstance(entry, dict) or entry.get("type") != "message": continue message_id = str(entry.get("id") or "") message_timestamp = str(entry.get("timestamp") or "") - # Forks retain both fields; IDs alone can collide across sessions. + # Forks retain both fields; IDs alone can collide across sessions. A + # line without an id has nothing to tie it to a copy elsewhere, so it + # always counts, keyed by where it sits. if message_id and message_timestamp: message_key = ("entry", message_id, message_timestamp) - else: + elif message_id: message_key = ("path", path, message_id) + else: + message_key = ("line", path, line_number) duplicate = message_key in seen_pi_messages seen_pi_messages.add(message_key) message = entry.get("message") or {} diff --git a/test/shell.d/agent-usage-codex-scanner-test.sh b/test/shell.d/agent-usage-codex-scanner-test.sh index 5c9ea43b8a5..2dba8ce93f0 100644 --- a/test/shell.d/agent-usage-codex-scanner-test.sh +++ b/test/shell.d/agent-usage-codex-scanner-test.sh @@ -213,6 +213,25 @@ result=$(HOME="$CUSTOM_PI_HOME" CODEX_HOME="$CUSTOM_PI_HOME/.codex" XDG_CACHE_HO fail "Codex collector counts a default root reached through a symlinked override once" "$result" pass "Codex collector follows pi's agent and session directory overrides" +# A line without an id always counts, even beside an identical one, and a +# session override nested inside a default root reads it only once. +IDLESS_PI_HOME=$(signed_in_home) +cp "$TEST_HOME/bin/codex" "$IDLESS_PI_HOME/bin/codex" +mkdir -p "$IDLESS_PI_HOME/.pi/agent/sessions/project" +cat >"$IDLESS_PI_HOME/.pi/agent/sessions/project/idless.jsonl" < Date: Wed, 14 Oct 2026 22:59:02 -0400 Subject: [PATCH 066/546] Count a forked pi message once in Claude's usage A pi fork repeats its parent's messages, and Claude keyed them by file, so every copied message counted again under the fork. Messages are keyed by id and timestamp as the Codex collector does, with an id-less line still counting by its place in the file. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-claude | 15 +++++++++++-- .../agent-usage-claude-scanner-test.sh | 22 +++++++++++++++++++ 2 files changed, 35 insertions(+), 2 deletions(-) diff --git a/bin/omarchy-agent-usage-claude b/bin/omarchy-agent-usage-claude index de7ccd1dd9a..c5d4362c5e1 100755 --- a/bin/omarchy-agent-usage-claude +++ b/bin/omarchy-agent-usage-claude @@ -569,7 +569,7 @@ def scan_pi_usage(max_age_seconds: float) -> dict[str, Any] | None: today_sessions: set[str] = set() today_tokens: dict[str, int] = {} usage_by_model: dict[str, dict[str, int]] = {} - seen: set[str] = set() + seen: set[tuple[str, str, Any]] = set() prompts = 0 today_prompt_count = 0 today_token_total = 0 @@ -590,7 +590,18 @@ def scan_pi_usage(max_age_seconds: float) -> dict[str, Any] | None: provider = str(message.get("provider") or "") if provider != "anthropic": continue - unique_key = f"{path}:{entry.get('id') or line_number}" + # A fork repeats its parent's messages with their ids and + # timestamps, so the pair counts a message once across the + # tree; an id alone can collide between sessions. A line + # without an id always counts. + message_id = str(entry.get("id") or "") + stamp = str(entry.get("timestamp") or "") + if message_id and stamp: + unique_key = ("entry", message_id, stamp) + elif message_id: + unique_key = ("path", str(path), message_id) + else: + unique_key = ("line", str(path), line_number) if unique_key in seen: continue seen.add(unique_key) diff --git a/test/shell.d/agent-usage-claude-scanner-test.sh b/test/shell.d/agent-usage-claude-scanner-test.sh index c927842607d..c5e3a1f7536 100644 --- a/test/shell.d/agent-usage-claude-scanner-test.sh +++ b/test/shell.d/agent-usage-claude-scanner-test.sh @@ -385,6 +385,28 @@ result=$(HOME="$CUSTOM_PI_HOME" CLAUDE_CONFIG_DIR="$CUSTOM_PI_HOME/.claude" XDG_ fail "Claude collector counts a default root reached through a symlinked override once" "$result" pass "Claude collector follows pi's agent and session directory overrides" +# A fork repeats its parent's messages with their ids and timestamps, so they +# count once; a line without an id always counts, even beside an identical one. +FORK_PI_HOME=$(mktemp -d "$SCRATCH/home.XXXXXX") +mkdir -p "$FORK_PI_HOME/.pi/agent/sessions/project" +cat >"$FORK_PI_HOME/.pi/agent/sessions/project/parent.jsonl" <"$FORK_PI_HOME/.pi/agent/sessions/project/fork.jsonl" < Date: Wed, 14 Oct 2026 23:01:37 -0400 Subject: [PATCH 067/546] Read Codex usage from opencode's V2 store and OPENCODE_DB The Codex collector read only opencode's original message table at the default path, so usage written by a V2 build, or to a database OPENCODE_DB names, never reached it. It now reads both stores as the opencode collector does: V2 first, a migrated message counted once by its id, a fork's copies of its source's messages left out, and compactions counted like replies. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-codex | 185 ++++++++++++------ .../shell.d/agent-usage-codex-scanner-test.sh | 71 +++++-- 2 files changed, 188 insertions(+), 68 deletions(-) diff --git a/bin/omarchy-agent-usage-codex b/bin/omarchy-agent-usage-codex index 860c51a4449..d4017ea0997 100755 --- a/bin/omarchy-agent-usage-codex +++ b/bin/omarchy-agent-usage-codex @@ -369,6 +369,126 @@ def scan_pi_sessions(): today_sessions.add(path) +def opencode_db(): + # OPENCODE_DB points opencode at another database, as the CLI itself reads + # it; a relative name lives beside the default one. Only a build whose + # channel was literally "v2" wrote opencode-v2.db, so it is a fallback for + # those installs alone. + root = Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) / "opencode" + configured = os.environ.get("OPENCODE_DB", "").strip() + if configured: + path = Path(configured).expanduser() + return path if path.is_absolute() else root / path + normal = root / "opencode.db" + legacy = root / "opencode-v2.db" + return legacy if not normal.is_file() and legacy.is_file() else normal + + +def table_columns(conn, table): + try: + return {row[1] for row in conn.execute(f"PRAGMA table_info({table})")} + except sqlite3.Error: + return set() + + +def opencode_row(message_id, session_id, entry, row_type, created, seen): + # Shared by both message stores. The original `message` table keeps the + # role, provider and model at the top level; the V2 `session_message` table + # nests the provider and model under `model` and keeps the role in its + # `type` column, with a compaction's tokens spent like any reply's. + if not isinstance(entry, dict): + return + if (entry.get("role") or row_type) not in ("assistant", "compaction"): + return + model = entry.get("model") if isinstance(entry.get("model"), dict) else {} + provider = entry.get("providerID") or model.get("providerID") + # Exact match: opencode provider ids are free-form, and a custom + # "openai-local" gateway is not this subscription. + if str(provider or "") != "openai": + return + tokens = entry.get("tokens") if isinstance(entry.get("tokens"), dict) else {} + cache = tokens.get("cache") if isinstance(tokens.get("cache"), dict) else {} + 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): + return + # A migrated message can sit in both stores under the same id. + if message_id in seen: + return + seen.add(message_id) + times = entry.get("time") if isinstance(entry.get("time"), dict) else {} + day = local_day(times.get("created") or created) + model_id = entry.get("modelID") or model.get("id") + add_usage(day, "opencode:" + str(session_id), model_name(str(model_id or "").rstrip("/").split("/")[-1]), input_tokens, output_tokens, cache_read, cache_write) + + +def scan_opencode_session_messages(conn, tables, seen): + # The V2 store. The LIKE gate skips rows whose JSON cannot name the + # provider before any parse. + columns = table_columns(conn, "session_message") + if not {"id", "session_id", "type", "data"} <= columns: + return + created = "m.time_created" if "time_created" in columns else "NULL" + source = "session_message m" + fork_filter = "" + if "time_created" in columns and "session_v2" in tables and {"id", "fork_session_id", "time_created"} <= table_columns(conn, "session_v2"): + # A fork copies its source session's messages under new ids; only what + # came after the fork was spent in it. + source += " LEFT JOIN session_v2 s ON s.id = m.session_id" + fork_filter = " AND (s.fork_session_id IS NULL OR m.time_created >= s.time_created)" + for message_id, session_id, row_type, time_created, raw in conn.execute( + f"SELECT m.id, m.session_id, m.type, {created}, m.data FROM {source}" + " WHERE m.type IN ('assistant', 'compaction')" + " AND m.data LIKE '%\"providerID\"%:%\"openai\"%'" + fork_filter + ): + # One malformed row must not abort the scan. + try: + opencode_row(str(message_id), session_id, json.loads(raw), str(row_type or ""), time_created, seen) + except Exception: + continue + + +def scan_opencode_messages(conn, seen): + # The original store. OpenCode DBs grow huge (every historical message, + # JSON included), and Python-side json.loads of every row wasted ~600 MB of + # RSS on machines whose subscription never ran on OpenAI. The json_extract + # conditions are the authority for the rows that reach them: role == + # "assistant" and providerID == "openai", the same exact-match values the + # Python filter checks. The guards in front are pure acceleration, not a + # perfect proxy for the Python filter: + # - The LIKE gates skip rows whose JSON cannot contain the two key/value + # pairs, avoiding the JSON parse of tens-of-MB blobs. They can differ + # from json.loads on duplicate keys (Python keeps the last, SQLite + # json_extract keeps the first) and ASCII-escaped values + # (\"ass\\u0069stant\" decodes for Python but not for LIKE), so a row + # the authority would accept can be gated out. Both cases are + # vanishingly rare in real opencode data. + # - json_valid guards the parse itself: json_extract() RAISES on + # malformed JSON instead of returning NULL, and one such row would + # otherwise abort the whole scan. SQLite does not promise that AND + # terms evaluate left to right, so the guard is a CASE around each + # json_extract rather than a separate AND term. + columns = table_columns(conn, "message") + if not {"session_id", "data"} <= columns: + return + message_id = "id" if "id" in columns else "rowid" + for row_id, session_id, raw in conn.execute( + f"SELECT {message_id}, session_id, data FROM message" + " WHERE data LIKE '%\"role\"%:%\"assistant\"%'" + " AND data LIKE '%\"providerID\"%:%\"openai\"%'" + " 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 = 'openai'" + ): + # One row that still fails json.loads must not abort the scan. + try: + opencode_row(str(row_id), session_id, json.loads(raw), "", None, seen) + except Exception: + continue + + def scan_opencode_sessions(): # A subscription burned entirely through opencode leaves no native session # files, but opencode records per-message provider, model, and token usage @@ -376,7 +496,7 @@ def scan_opencode_sessions(): # 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 = Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) / "opencode" / "opencode.db" + db = opencode_db() if not db.is_file(): return True try: @@ -385,59 +505,14 @@ def scan_opencode_sessions(): return False try: conn.execute("PRAGMA query_only = ON") - # OpenCode DBs grow huge (every historical message, JSON included), and - # Python-side json.loads of every row wasted ~600 MB of RSS on machines - # whose subscription never ran on OpenAI. The json_extract conditions are - # the authority for the rows that reach them: role == "assistant" and - # providerID == "openai", the same exact-match values the Python filter - # below checks. The guards in front are pure acceleration, not a perfect - # proxy for the Python filter: - # - The LIKE gates skip rows whose JSON cannot contain the two - # key/value pairs, avoiding the JSON parse of tens-of-MB blobs. They - # can differ from json.loads on duplicate keys (Python keeps the - # last, SQLite json_extract keeps the first) and ASCII-escaped - # values (\"ass\\u0069stant\" decodes for Python but not for LIKE), - # so a row the authority would accept can be gated out. Both cases - # are vanishingly rare in real opencode data. - # - json_valid guards the parse itself: json_extract() RAISES on - # malformed JSON instead of returning NULL, and one such row would - # otherwise abort the whole scan. SQLite does not promise that AND - # terms evaluate left to right, so the guard is a CASE around each - # json_extract rather than a separate AND term. Rows that are not - # well-formed JSON are skipped here; the per-row try/except below - # stays as the final safety net for rows that pass the SQL filter - # but fail json.loads. - for session_id, raw in conn.execute( - "SELECT session_id, data FROM message" - " WHERE data LIKE '%\"role\"%:%\"assistant\"%'" - " AND data LIKE '%\"providerID\"%:%\"openai\"%'" - " 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 = 'openai'" - ): - # One malformed row must not abort the scan, so every shape assumption - # lives inside the try. - try: - entry = json.loads(raw) - # Exact match: opencode provider ids are free-form, and a custom - # "openai-local" gateway is not this subscription. - if not isinstance(entry, dict) or entry.get("role") != "assistant": - continue - if str(entry.get("providerID") or "") != "openai": - continue - tokens = entry.get("tokens") or {} - cache = tokens.get("cache") or {} - input_tokens = number(tokens.get("input")) - # opencode keeps thinking tokens out of output; both are generated. - output_tokens = number(tokens.get("output")) + number(tokens.get("reasoning")) - cache_read = number(cache.get("read")) - cache_write = number(cache.get("write")) - if not (input_tokens or output_tokens or cache_read or cache_write): - continue - day = local_day((entry.get("time") or {}).get("created")) - model = model_name(str(entry.get("modelID") or "").rstrip("/").split("/")[-1]) - except Exception: - continue - add_usage(day, "opencode:" + str(session_id), model, input_tokens, output_tokens, cache_read, cache_write) + tables = {name for (name,) in conn.execute("SELECT name FROM sqlite_master WHERE type = 'table'")} + # Both stores, V2 first, deduplicated by message id: a build that + # migrated keeps counting, and a message in both counts once. + seen = set() + if "session_message" in tables: + scan_opencode_session_messages(conn, tables, seen) + if "message" in tables: + scan_opencode_messages(conn, seen) except sqlite3.Error: # Transient lock, schema migration, corruption: the numbers stop here, # incomplete. @@ -559,7 +634,7 @@ def cache_root(): def scan_cache_digest(): codex_home = Path(os.environ.get("CODEX_HOME") or (Path.home() / ".codex")) - db = Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) / "opencode" / "opencode.db" + db = opencode_db() # The digest covers every data path the scan reads: the codex session # roots, the opencode DB, and the pi/omp session roots (under Path.home(), # or wherever pi's own overrides moved them). diff --git a/test/shell.d/agent-usage-codex-scanner-test.sh b/test/shell.d/agent-usage-codex-scanner-test.sh index 2dba8ce93f0..81a36373509 100644 --- a/test/shell.d/agent-usage-codex-scanner-test.sh +++ b/test/shell.d/agent-usage-codex-scanner-test.sh @@ -281,6 +281,60 @@ pass "Codex collector counts OpenAI usage, reasoning included, from opencode ses fail "Codex collector ignores prefix-colliding providers, user messages, and malformed rows" "$result" pass "Codex collector ignores prefix-colliding providers, user messages, and malformed rows" +# OpenCode V2 nests the provider under `model` and spends tokens on +# compactions too. A message migrated into both stores counts once, and a +# fork's copies of its source's messages don't count again. OPENCODE_DB +# points at the database the way opencode itself reads it. +V2_HOME=$(signed_in_home) +cp "$TEST_HOME/bin/codex" "$V2_HOME/bin/codex" + +python3 - "$V2_HOME/elsewhere/opencode-dev.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)") +conn.execute("CREATE TABLE session_message (id text PRIMARY KEY, session_id text NOT NULL, type text NOT NULL, time_created integer NOT NULL, data text NOT NULL)") +conn.execute("CREATE TABLE session_v2 (id text PRIMARY KEY, fork_session_id text, time_created integer NOT NULL)") +now_ms = int(time.time() * 1000) +forked_ms = now_ms - 1000 + +def v2(id, session, kind, provider, model, input, created=now_ms): + return (id, session, kind, created, json.dumps({ + "model": {"providerID": provider, "id": model}, + "tokens": {"input": input, "output": 0, "reasoning": 0, "cache": {"read": 0, "write": 0}}, + "time": {"created": created}, + })) + +conn.executemany("INSERT INTO session_v2 VALUES (?, ?, ?)", [("ses_a", None, now_ms - 5000), ("ses_b", "ses_a", forked_ms)]) +conn.executemany("INSERT INTO session_message VALUES (?, ?, ?, ?, ?)", [ + v2("msg_1", "ses_a", "assistant", "openai", "gpt-5.2-codex", 100, now_ms - 4000), + v2("msg_2", "ses_a", "compaction", "openai", "gpt-5.2-codex", 10, now_ms - 3000), + v2("msg_3", "ses_a", "assistant", "openai-local", "gpt-5.2-codex", 999), + v2("msg_4", "ses_b", "assistant", "openai", "gpt-5.2-codex", 100, forked_ms - 4000), + v2("msg_5", "ses_b", "assistant", "openai", "gpt-5.2-codex", 1), +]) +conn.execute("INSERT INTO message VALUES ('msg_1', 'ses_a', ?, ?, ?)", (now_ms, now_ms, json.dumps({ + "role": "assistant", "providerID": "openai", "modelID": "gpt-5.2-codex", + "tokens": {"input": 100, "output": 0, "reasoning": 0, "cache": {"read": 0, "write": 0}}, + "time": {"created": now_ms - 4000}, +}))) +conn.commit() +conn.close() +PY + +result=$(HOME="$V2_HOME" CODEX_HOME="$V2_HOME/.codex" XDG_CACHE_HOME="$V2_HOME/.cache" XDG_DATA_HOME="$V2_HOME/.local/share" \ + OPENCODE_DB="$V2_HOME/elsewhere/opencode-dev.db" PATH="$V2_HOME/bin:$PATH" "$collector") + +[[ $(jq -c '{todayTotalTokens, totalPrompts, totalSessions, models: (.modelUsage | keys)}' <<<"$result") == '{"todayTotalTokens":111,"totalPrompts":3,"totalSessions":2,"models":["gpt-5.2-codex"]}' ]] || + fail "Codex collector reads opencode's V2 records once each, without fork copies" "$result" +pass "Codex collector reads opencode's V2 records once each, without fork copies" + # A warm cache makes --limits-only cheap: local stats come from the last scan # instead of another walk over the opencode database, and --force bypasses it. CACHE_HOME=$(signed_in_home) @@ -688,19 +742,9 @@ pass "Codex collector still prints a complete record when the cache is unwritabl INTERRUPTED_HOME=$(signed_in_home) cp "$TEST_HOME/bin/codex" "$INTERRUPTED_HOME/bin/codex" -# A database without the message table makes the scan fail mid-flight. -python3 - "$INTERRUPTED_HOME/.local/share/opencode/opencode.db" <<'PY' -import sqlite3 -import sys -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 unrelated (id text PRIMARY KEY)") -conn.commit() -conn.close() -PY +# A file that isn't a database makes the scan fail mid-flight. +mkdir -p "$INTERRUPTED_HOME/.local/share/opencode" +printf 'not a sqlite database, just bytes in the way\n' >"$INTERRUPTED_HOME/.local/share/opencode/opencode.db" result=$(HOME="$INTERRUPTED_HOME" CODEX_HOME="$INTERRUPTED_HOME/.codex" XDG_CACHE_HOME="$INTERRUPTED_HOME/.cache" XDG_DATA_HOME="$INTERRUPTED_HOME/.local/share" \ PATH="$INTERRUPTED_HOME/bin:$PATH" "$collector") @@ -720,6 +764,7 @@ import time from pathlib import Path db = Path(sys.argv[1]) +db.unlink() 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) From 212cc87ab25425b6d5808b807bb0e00e8bcdb4c2 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:02:50 -0400 Subject: [PATCH 068/546] Read every Codex rate-limit bucket The collector read only rateLimits, the app-server's backward-compatible single-bucket view, so a model's own quota in rateLimitsByLimitId never showed and a map-only answer showed no limits at all. Every bucket is read now, the general one first and untitled, and each other titled " Session" or " Weekly" so the panel shows it beside the general window of that span. A percentage or reset time that doesn't parse costs that window or field rather than the whole probe, and percentages are kept within 0..1. Builds on #14630. Co-Authored-By: Claude Opus 5.5 Co-authored-by: Harrison Kairalla --- bin/omarchy-agent-usage-codex | 79 +++++++++-- test/shell.d/agent-usage-codex-limits-test.sh | 129 ++++++++++++++++++ 2 files changed, 196 insertions(+), 12 deletions(-) create mode 100644 test/shell.d/agent-usage-codex-limits-test.sh diff --git a/bin/omarchy-agent-usage-codex b/bin/omarchy-agent-usage-codex index d4017ea0997..df2518dcf16 100755 --- a/bin/omarchy-agent-usage-codex +++ b/bin/omarchy-agent-usage-codex @@ -14,6 +14,7 @@ import argparse import fcntl import hashlib import json +import math import os import select import shutil @@ -894,8 +895,12 @@ def rpc_request(proc, request_id, method, params=None, timeout=8): def limit_window(window): if not isinstance(window, dict): return None - used = window.get("usedPercent") - if used is None: + try: + used = float(window.get("usedPercent")) + except (TypeError, ValueError): + return None + # A percentage that isn't one (NaN, infinite, a boolean) says nothing. + if isinstance(window.get("usedPercent"), bool) or not math.isfinite(used): return None mins = number(window.get("windowDurationMins")) if mins == 10080: @@ -906,14 +911,64 @@ def limit_window(window): label = f"{mins}m window" else: label = "Limit" - reset = window.get("resetsAt") + reset = number(window.get("resetsAt")) + try: + resets_at = datetime.fromtimestamp(reset, timezone.utc).isoformat() if reset > 0 else "" + except (OverflowError, OSError, ValueError): + resets_at = "" return { "label": label, - "percent": float(used) / 100.0, - "resetsAt": datetime.fromtimestamp(number(reset), timezone.utc).isoformat() if reset else "", + "percent": max(0.0, min(1.0, used / 100.0)), + "resetsAt": resets_at, + "minutes": mins, } +# The panel reads a window from its title, and a model's own allowance titled +# " Session" or " Weekly" sits beside the general window of the +# same span rather than as a row of its own. +def window_title(minutes): + if minutes >= 28 * 24 * 60: + return "Monthly" + if minutes >= 24 * 60: + return "Weekly" + if minutes > 0: + return "Session" + return "" + + +def rate_limit_buckets(result): + # New app-servers return every metered bucket; rateLimits is only the + # backward-compatible single-bucket view. Never count both views. + buckets = result.get("rateLimitsByLimitId") + if isinstance(buckets, dict): + entries = [(str(key), value) for key, value in buckets.items() if isinstance(value, dict)] + if entries: + # Keep the subscription's general limits and plan ahead of model limits. + return sorted(entries, key=lambda entry: entry[0] != "codex") + legacy = result.get("rateLimits") + return [("codex", legacy)] if isinstance(legacy, dict) else [] + + +def bucket_limits(buckets): + limits = [] + for limit_id, bucket in buckets: + # The general bucket's windows are the plain ones. Any other bucket is a + # model's own allowance, named the way the app-server names it. + general = limit_id == "codex" or (len(buckets) == 1 and not bucket.get("limitName")) + name = str(bucket.get("limitName") or limit_id).strip() + for window in (bucket.get("primary"), bucket.get("secondary")): + entry = limit_window(window) + if not entry: + continue + minutes = entry.pop("minutes") + if not general: + span = window_title(minutes) + entry["title"] = f"{name} {span}" if span else name + limits.append(entry) + return limits + + def codex_rpc_help(proc, stderr_file, exc): """Explain an app-server RPC failure: a live server's stall, a dead one's own stderr, or the login hint when it died saying nothing.""" @@ -1032,9 +1087,12 @@ def fetch_codex_rpc(home=None): result["usageStatusText"] = "Codex limits unavailable" result["authHelpText"] = message return result - limits = (limits_msg.get("result") or {}).get("rateLimits") or {} + rpc_result = limits_msg.get("result") if isinstance(limits_msg.get("result"), dict) else {} + buckets = rate_limit_buckets(rpc_result) result["authHelpText"] = "" - plan = limits.get("planType") or "" + plan = next((bucket.get("planType") for _, bucket in buckets if bucket.get("planType")), "") + if not plan and isinstance(rpc_result.get("rateLimits"), dict): + plan = rpc_result["rateLimits"].get("planType") or "" # The limits name the plan themselves. account/read is only a fallback # for when they don't, and never a reason to lose them: Codex 0.158's @@ -1048,14 +1106,11 @@ def fetch_codex_rpc(home=None): pass result["tierLabel"] = str(plan) if plan else "" - for window in (limits.get("primary"), limits.get("secondary")): - entry = limit_window(window) - if entry: - result["limits"].append(entry) + result["limits"] = bucket_limits(buckets) # OpenAI hands out free "full reset" credits that wipe the rate limits on # demand and lapse after a month. Worth knowing about before a limit bites. - credits = reset_credits(limits_msg.get("result") or {}) + credits = reset_credits(rpc_result) if credits: result["resetCredits"] = credits except Exception as exc: diff --git a/test/shell.d/agent-usage-codex-limits-test.sh b/test/shell.d/agent-usage-codex-limits-test.sh new file mode 100644 index 00000000000..1bd18f442ae --- /dev/null +++ b/test/shell.d/agent-usage-codex-limits-test.sh @@ -0,0 +1,129 @@ +#!/bin/bash + +set -euo pipefail + +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" + +require_command python3 + +SCRATCH=$(mktemp -d) +trap 'rm -rf "$SCRATCH"' EXIT + +HOME="$SCRATCH" python3 - "${1:-$ROOT/bin/omarchy-agent-usage-codex}" <<'PY' +import json +import os +import runpy +import sys +from pathlib import Path + +collector = sys.argv[1] +home = Path.home() +codex_home = home / ".codex" +codex_home.mkdir() +(codex_home / "auth.json").write_text("{}") +bin_dir = home / "bin" +bin_dir.mkdir() +stub = bin_dir / "codex" +stub.write_text("#!" + sys.executable + "\n" + ''' +import json +import os +import sys +from pathlib import Path + +for line in sys.stdin: + request = json.loads(line) + method = request["method"] + if "id" not in request: + continue + if method == "account/read": + with Path(os.environ["RPC_CALLS"]).open("a") as log: + log.write(method + "\\n") + reply = {"account": {"planType": "plus"}} + elif method == "account/rateLimits/read": + reply = json.loads(os.environ["RPC_LIMITS"]) + else: + reply = {} + print(json.dumps({"id": request["id"], "result": reply}), flush=True) +''') +stub.chmod(0o755) +calls = home / "calls" +os.environ.update(PATH=str(bin_dir) + os.pathsep + os.environ["PATH"], + CODEX_HOME=str(codex_home), RPC_CALLS=str(calls)) + + +def window(percent, minutes=300): + return {"usedPercent": percent, "windowDurationMins": minutes} + + +def check(name, payload, expected, plan="pro", account_read=False, reset_credits=None): + calls.write_text("") + os.environ["RPC_LIMITS"] = json.dumps(payload) + loaded = runpy.run_path(collector, run_name="limits_test") + record = loaded["fetch_codex_rpc"](codex_home) + assert record["usageStatusText"] == "", (name, record) + assert record["authHelpText"] == "", (name, record) + assert record["tierLabel"] == plan, (name, record) + assert record["limits"] == expected, (name, record) + assert bool(calls.read_text()) == account_read, (name, calls.read_text()) + assert record.get("resetCredits") == reset_credits, (name, record) + print("ok - " + name) + + +def entry(percent, label="5h window", title=None): + result = {"label": label, "percent": percent / 100, "resetsAt": ""} + if title: + result["title"] = title + return result + + +legacy = {"planType": "pro", "primary": window(12)} +check("legacy single-bucket limits keep their display shape", {"rateLimits": legacy}, [entry(12)]) +check("map-only limits supply both windows and the plan", { + "rateLimitsByLimitId": {"codex": {"planType": "pro", "primary": window(25), + "secondary": window(42, 10080)}} +}, [entry(25), entry(42, "Weekly (7-day)")]) +check("multi-bucket limits supersede the legacy view without duplicates", { + "rateLimits": legacy, + "rateLimitsByLimitId": {"codex": {"planType": "pro", "primary": window(36)}} +}, [entry(36)]) +check("named buckets are titled for their model and window, general Codex limits first", { + "rateLimits": {"planType": "plus", "primary": window(99)}, + "rateLimitsByLimitId": { + "codex_other": {"planType": "plus", "limitName": "Other model", "primary": window(80, 60)}, + "codex": {"planType": "pro", "primary": window(25)}} +}, [entry(25), entry(80, "1h window", "Other model Session")]) +check("a single named bucket keeps its model label", { + "rateLimitsByLimitId": {"codex_other": {"planType": "pro", "limitName": "Other model", + "primary": window(80)}} +}, [entry(80, title="Other model Session")]) +check("unnamed model buckets are titled with their identifiers", { + "rateLimitsByLimitId": {"codex": legacy, "codex_other": {"primary": window(80)}} +}, [entry(12), entry(80, title="codex_other Session")]) +for unavailable in (None, {}, {"codex": None}, []): + check("unavailable multi-bucket view falls back to legacy: " + repr(unavailable), { + "rateLimits": legacy, "rateLimitsByLimitId": unavailable + }, [entry(12)]) +check("invalid map entries do not hide valid buckets", { + "rateLimitsByLimitId": {"invalid": None, "codex": legacy} +}, [entry(12)]) +check("legacy plan metadata is retained when map buckets omit it", { + "rateLimits": legacy, "rateLimitsByLimitId": {"codex": {"primary": window(33)}} +}, [entry(33)]) +check("account metadata remains an optional plan fallback", { + "rateLimitsByLimitId": {"codex": {"primary": window(33)}} +}, [entry(33)], plan="plus", account_read=True) +check("malformed legacy metadata does not erase valid map limits", { + "rateLimits": [], "rateLimitsByLimitId": {"codex": {"primary": window(33)}} +}, [entry(33)], plan="plus", account_read=True) +check("a weekly model bucket attaches to the weekly window", { + "rateLimitsByLimitId": {"codex": legacy, "codex_spark": {"limitName": "Spark", "secondary": window(30, 10080)}} +}, [entry(12), entry(30, "Weekly (7-day)", "Spark Weekly")]) +check("an untrusted percentage or reset time costs only that window or field", { + "rateLimits": {"planType": "pro", "primary": {"usedPercent": "lots", "windowDurationMins": 300}, + "secondary": {"usedPercent": 250, "windowDurationMins": 10080, "resetsAt": -5}} +}, [entry(100, "Weekly (7-day)")]) +check("multi-bucket limits retain earned reset credits", { + "rateLimitsByLimitId": {"codex": legacy}, + "rateLimitResetCredits": {"credits": [{"status": "available", "expiresAt": 1900000000}]} +}, [entry(12)], reset_credits={"available": 1, "nextExpiresAt": "2030-03-17T17:46:40+00:00"}) +PY From 9435cfd25de824c23d06b355a7542c3638741269 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:07:55 -0400 Subject: [PATCH 069/546] Keep the last Codex limits, marked stale, when a probe fails Any failed app-server probe replaced the record's limits with an empty list, so the panel lost limits that were right a refresh ago, and every panel open started another app-server. Each signed-in login now keeps its last good answer in its own codex-limits-.json, keyed by home and by the workspace and user in auth.json, or by CODEX_ACCESS_TOKEN; a keyring login, which can't be told apart, keeps nothing, and neither does a probe a login change ran across. A failed probe serves the kept limits as stale, with elapsed windows at zero, and records and accounts carry limitsStale/limitsFetchedAt and stale/fetchedAt. A probe within 15 seconds is reused unless --force asks again. A stalled app-server, or one that died unable to reach OpenAI, sets retryAdvised. A missing CLI reads as "Codex limits unavailable", the contract's word for it. Builds on #15249 and #14298. Co-Authored-By: Claude Opus 5.5 Co-authored-by: Omarchybot <317366263+omarchybot@users.noreply.github.com> Co-authored-by: AnPod --- bin/omarchy-agent-usage-codex | 188 ++++++++++++++++- test/shell.d/agent-usage-codex-limits-test.sh | 190 +++++++++++++++++- .../shell.d/agent-usage-codex-scanner-test.sh | 8 +- 3 files changed, 373 insertions(+), 13 deletions(-) diff --git a/bin/omarchy-agent-usage-codex b/bin/omarchy-agent-usage-codex index df2518dcf16..7cecd34982d 100755 --- a/bin/omarchy-agent-usage-codex +++ b/bin/omarchy-agent-usage-codex @@ -11,6 +11,7 @@ panel only ever reads the JSON this prints. """ import argparse +import base64 import fcntl import hashlib import json @@ -44,6 +45,22 @@ NATIVE_PROVIDER = "openai" SCAN_REUSE_SECONDS = 20 LIMITS_ONLY_REUSE_SECONDS = 900 +# Each probe starts an app-server, which syncs the plugin list on the way up. +# A panel opened and shut repeatedly reuses an answer this recent instead. +PROBE_MIN_INTERVAL_SECONDS = 15 + +# What a dead app-server's stderr says when it never reached OpenAI at all, +# rather than being turned away by it. +TRANSPORT_ERRORS = ( + "error sending request", + "dns error", + "failed to lookup address", + "connection refused", + "connection reset", + "network is unreachable", + "timed out", +) + # Per-file totals of native Codex sessions survive between runs, so a rescan # only reads the session files that changed. Bumping the schema retires every # record written by an older layout. @@ -1031,7 +1048,7 @@ def fetch_codex_rpc(home=None): env = dict(ENV, CODEX_HOME=str(home)) if home else ENV codex = find_codex_binary() if not codex: - result["usageStatusText"] = "Codex unavailable" + result["usageStatusText"] = "Codex limits unavailable" result["authHelpText"] = "codex not found in PATH" return result @@ -1065,8 +1082,8 @@ def fetch_codex_rpc(home=None): except Exception as exc: if stderr_file: stderr_file.close() - result["usageStatusText"] = "Codex unavailable" - result["authHelpText"] = str(exc) + result["usageStatusText"] = "Codex limits unavailable" + result["authHelpText"] = str(exc)[:300] return result try: @@ -1113,9 +1130,19 @@ def fetch_codex_rpc(home=None): credits = reset_credits(rpc_result) if credits: result["resetCredits"] = credits + result["ok"] = True except Exception as exc: result["usageStatusText"] = "Codex limits unavailable" result["authHelpText"] = codex_rpc_help(proc, stderr_file, exc) + # A live app-server that never answered, or one that died saying it + # couldn't reach OpenAI, is a network that isn't up yet (the first probe + # after login often beats DHCP), so the panel asks again soon. One that + # was turned away, or died for its own reasons, would only fail again. + help_text = result["authHelpText"].lower() + if isinstance(exc, TimeoutError) or ( + help_text.startswith("codex app-server exited:") and any(error in help_text for error in TRANSPORT_ERRORS) + ): + result["transport"] = True finally: try: proc.terminate() @@ -1130,6 +1157,141 @@ def fetch_codex_rpc(home=None): return result +# Who a home's probe answers for, or "" when only Codex could say: keyring and +# auto stores keep the login out of auth.json. +def signed_in_identity(codex_home): + token = ENV.get("CODEX_ACCESS_TOKEN") + if token: + return "token:" + hashlib.sha1(token.encode("utf-8", "surrogatepass")).hexdigest() + try: + with (codex_home / "config.toml").open("rb") as handle: + config = tomllib.load(handle) + except FileNotFoundError: + config = {} + except Exception: + return "" + if str(config.get("cli_auth_credentials_store") or "file") != "file": + return "" + # account_id names the workspace, whose members each have their own limits, + # so the user comes from the ID token's claims, read as Codex reads them. + try: + tokens = json.loads((codex_home / "auth.json").read_text(encoding="utf-8"))["tokens"] + payload = tokens["id_token"].split(".")[1] + claims = json.loads(base64.urlsafe_b64decode(payload + "=" * (-len(payload) % 4)))["https://api.openai.com/auth"] + workspace = str(tokens["account_id"] or "") + user = str(claims.get("chatgpt_user_id") or claims.get("user_id") or "") + except Exception: + return "" + return f"{workspace}\n{user}" if workspace and user else "" + + +# Keyed by the home and the account it is signed in to, so a home signed in to +# someone else never falls back on the last one's numbers. None when the +# account can't be told, since then no numbers can be vouched for. +def limits_cache_path(home): + codex_home = Path(home or ENV.get("CODEX_HOME") or (Path.home() / ".codex")) + identity = signed_in_identity(codex_home) + if not identity: + return None + # surrogatepass: a lone surrogate from JSON or the environment must not stop the record. + digest = hashlib.sha1(f"{codex_home}\n{identity}".encode("utf-8", "surrogatepass")).hexdigest()[:16] + try: + return cache_root() / f"codex-limits-{digest}.json" + except Exception as exc: + print(f"omarchy-agent-usage-codex: limits cache unavailable ({exc})", file=sys.stderr) + return None + + +def read_json(path): + try: + value = json.loads(path.read_text(encoding="utf-8")) + except Exception: + return {} + return value if isinstance(value, dict) else {} + + +# Cached windows reset to zero when their measured period ends. Without a +# fresh probe, the next window has no known reset time yet. +def usable_cached_limits(cached): + entries = cached.get("limits") + if not isinstance(entries, list): + return [] + now = datetime.now(timezone.utc) + usable = [] + for entry in entries: + if not isinstance(entry, dict): + continue + try: + resets = datetime.fromisoformat(str(entry.get("resetsAt") or "").replace("Z", "+00:00")) + except ValueError: + resets = None + if resets and resets.tzinfo is None: + resets = resets.replace(tzinfo=timezone.utc) + usable.append(dict(entry, percent=0.0, resetsAt="") if resets and resets <= now else entry) + return usable + + +# The app-server's answer, with the last limits it gave kept as the answer of +# record when a later probe fails: one slow or dropped call is no reason to +# blank limits that were right a refresh ago. `live` says the limits came from +# this probe, or one recent enough to stand in for it. +def collect_limits(home=None, force=False): + cache = limits_cache_path(home) + cached = read_json(cache) if cache else {} + fallback = usable_cached_limits(cached) + fetched_at_ms = number(cached.get("fetchedAtMs")) + + # --force is a person asking for fresh numbers, so it skips the reuse + # window; the interval is there to absorb repeated panel opens. + fetched_age = time.time() - fetched_at_ms / 1000 + if fallback and not force and 0 <= fetched_age < PROBE_MIN_INTERVAL_SECONDS: + result = { + "limits": fallback, + "tierLabel": str(cached.get("tierLabel") or ""), + "usageStatusText": "", + "authHelpText": "", + "live": True, + "fetchedAtMs": fetched_at_ms, + } + if isinstance(cached.get("resetCredits"), dict): + result["resetCredits"] = cached["resetCredits"] + return result + + rpc = fetch_codex_rpc(home) + ok = rpc.pop("ok", False) + transport = rpc.pop("transport", False) + # A login that changed while the probe ran leaves no telling whose these are. + same_login = cache is not None and cache == limits_cache_path(home) + + if ok: + rpc["live"] = True + rpc["fetchedAtMs"] = round(time.time() * 1000) + if same_login and rpc["limits"]: + try: + write_json(cache, { + "fetchedAtMs": rpc["fetchedAtMs"], + "limits": rpc["limits"], + "tierLabel": rpc["tierLabel"], + "resetCredits": rpc.get("resetCredits"), + }) + except Exception as exc: + print(f"omarchy-agent-usage-codex: could not write limits cache ({exc})", file=sys.stderr) + return rpc + + rpc["live"] = False + rpc["fetchedAtMs"] = 0 + if transport: + rpc["retryAdvised"] = True + if fallback and same_login: + # A sign-in to restore still says so; any other failure leaves the kept + # numbers to speak for themselves, with their age, and the detail of what + # went wrong in the help text. + rpc.update(limits=fallback, fetchedAtMs=fetched_at_ms, tierLabel=rpc["tierLabel"] or str(cached.get("tierLabel") or "")) + if rpc["usageStatusText"] == "Codex limits unavailable": + rpc["usageStatusText"] = "" + return rpc + + # The accounts `omarchy agent account` registered, read straight from its # registry. Only a registry holding a second account matters: with one, the # record is exactly what it always was. @@ -1156,10 +1318,10 @@ def registered_accounts(): # Each account's app-server runs in that account's home, so it answers for that # login and refreshes that login's tokens there. -def account_limits(account): +def account_limits(account, force): # The primary is ~/.codex by definition, not whatever CODEX_HOME this run # happened to inherit — that would report another account's limits as Main's. - rpc = fetch_codex_rpc(account.get("home") or (Path.home() / ".codex")) + rpc = collect_limits(account.get("home") or (Path.home() / ".codex"), force) return { "id": account["id"], "label": str(account.get("label") or account["id"]), @@ -1168,9 +1330,11 @@ def account_limits(account): "active": account["active"], "primary": bool(account.get("primary")), "limits": rpc["limits"], - "stale": not rpc["limits"], + "stale": not rpc["live"], + "fetchedAt": rpc["fetchedAtMs"], "usageStatusText": rpc["usageStatusText"], "authHelpText": rpc["authHelpText"], + "retryAdvised": bool(rpc.get("retryAdvised")), "resetCredits": rpc.get("resetCredits"), } @@ -1189,7 +1353,7 @@ def main(): max_age = 0 if args.force else (LIMITS_ONLY_REUSE_SECONDS if args.limits_only else SCAN_REUSE_SECONDS) stats = cached_local_stats(max_age) registered = registered_accounts() - accounts = [account_limits(account) for account in registered] + accounts = [account_limits(account, args.force) for account in registered] current = next((a for a in accounts if a["active"]), None) if current: # The top-level fields keep describing the account new sessions use. @@ -1201,8 +1365,10 @@ def main(): } if current.get("resetCredits"): rpc["resetCredits"] = current["resetCredits"] + stale, fetched_at, retry = current["stale"], current["fetchedAt"], current["retryAdvised"] else: - rpc = fetch_codex_rpc() + rpc = collect_limits(force=args.force) + stale, fetched_at, retry = not rpc.pop("live"), rpc.pop("fetchedAtMs"), rpc.pop("retryAdvised", False) record = { "schemaVersion": 1, @@ -1214,6 +1380,12 @@ def main(): } record.update(stats) record.update(rpc) + # Numbers kept past a failed probe, and when they were measured, so the + # panel can tell them apart from ones checked just now. + record["limitsStale"] = stale + record["limitsFetchedAt"] = fetched_at + if retry: + record["retryAdvised"] = True if accounts: record["accountSwitch"] = registered[0]["switch"] record["accounts"] = accounts diff --git a/test/shell.d/agent-usage-codex-limits-test.sh b/test/shell.d/agent-usage-codex-limits-test.sh index 1bd18f442ae..a68ea099696 100644 --- a/test/shell.d/agent-usage-codex-limits-test.sh +++ b/test/shell.d/agent-usage-codex-limits-test.sh @@ -9,7 +9,7 @@ require_command python3 SCRATCH=$(mktemp -d) trap 'rm -rf "$SCRATCH"' EXIT -HOME="$SCRATCH" python3 - "${1:-$ROOT/bin/omarchy-agent-usage-codex}" <<'PY' +HOME="$SCRATCH" XDG_CACHE_HOME="$SCRATCH/.cache" python3 - "${1:-$ROOT/bin/omarchy-agent-usage-codex}" <<'PY' import json import os import runpy @@ -127,3 +127,191 @@ check("multi-bucket limits retain earned reset credits", { "rateLimitResetCredits": {"credits": [{"status": "available", "expiresAt": 1900000000}]} }, [entry(12)], reset_credits={"available": 1, "nextExpiresAt": "2030-03-17T17:46:40+00:00"}) PY + +require_command jq + +test_tmp="$SCRATCH/probe" +mkdir -p "$test_tmp" +export HOME="$test_tmp/home" +export XDG_STATE_HOME="$test_tmp/state" +export XDG_CACHE_HOME="$test_tmp/cache" +export XDG_DATA_HOME="$test_tmp/data" +unset CODEX_HOME CODEX_ACCESS_TOKEN OPENCODE_DB PI_CODING_AGENT_DIR PI_CODING_AGENT_SESSION_DIR + +# A scratch copy of the collector with its 8s RPC timeout shortened, so a +# stalled app-server gives up quickly. +collector="$test_tmp/omarchy-agent-usage-codex" +sed 's/timeout=8/timeout=0.5/g' "$ROOT/bin/omarchy-agent-usage-codex" >"$collector" +chmod +x "$collector" + +mkdir -p "$test_tmp/bin" "$test_tmp/empty-bin" + +mkdir -p "$HOME/.codex/sessions" + +# Signs ~/.codex in as a user of a workspace, the way Codex writes it: the +# workspace as account_id, the user in the ID token's claims. +cat >"$test_tmp/login" <<'PY' +#!/usr/bin/python3 +import base64, json, os, sys +claims = {"https://api.openai.com/auth": {"chatgpt_account_id": sys.argv[1], "chatgpt_user_id": sys.argv[2]}} +payload = base64.urlsafe_b64encode(json.dumps(claims).encode("utf-8", "surrogatepass")).decode().rstrip("=") +home = os.environ.get("CODEX_HOME") or os.path.expanduser("~/.codex") +with open(os.path.join(home, "auth.json"), "w") as handle: + json.dump({"tokens": {"id_token": f"e30.{payload}.sig", "account_id": sys.argv[1]}}, handle) +PY +chmod +x "$test_tmp/login" +login() { + "$test_tmp/login" "$@" +} +login ws-one user-a + +# A stand-in app-server whose rate-limit answer CODEX_FAKE picks: limits, a +# backend error, a sign-out, dying with or without saying why, or never +# answering. Every start is logged, so a reused probe can be told apart. +cat >"$test_tmp/bin/codex" <<'PY' +#!/usr/bin/python3 +import json, os, sys, time +mode = os.environ.get("CODEX_FAKE", "ok") +with open(os.environ["CODEX_STARTS"], "a") as log: + log.write(mode + "\n") +now = int(time.time()) +for line in sys.stdin: + message = json.loads(line) + if "id" not in message: + continue + if message["method"] == "account/rateLimits/read": + if mode == "died": + print("error sending request for url (https://chatgpt.com/backend-api/wham/usage)", file=sys.stderr) + sys.exit(1) + if mode == "crashed": + print("thread 'main' panicked at config.rs: unexpected argument", file=sys.stderr) + sys.exit(1) + if mode == "silent": + sys.exit(0) + if mode == "stall": + time.sleep(5) + continue + if mode == "backend": + print(json.dumps({"id": message["id"], "error": {"code": -32603, "message": "failed to fetch codex rate limits"}}), flush=True) + continue + if mode == "switch": + os.system(os.environ["CODEX_LOGIN"] + " ws-three user-c") + if mode == "auth": + print(json.dumps({"id": message["id"], "error": {"code": -32603, "message": "not authenticated"}}), flush=True) + continue + print(json.dumps({"id": message["id"], "result": {"rateLimits": {"planType": "pro", + "primary": {"usedPercent": 55 if mode == "switch" else 40, "windowDurationMins": 300, "resetsAt": now + 3600}, + "secondary": {"usedPercent": 70, "windowDurationMins": 10080, "resetsAt": now + 2}}}}), flush=True) + continue + print(json.dumps({"id": message["id"], "result": {}}), flush=True) +PY +chmod +x "$test_tmp/bin/codex" +export CODEX_STARTS="$test_tmp/starts" + +collect() { + local mode=$1 + shift + CODEX_FAKE="$mode" CODEX_LOGIN="$test_tmp/login" PATH="$test_tmp/bin:$PATH" "$collector" "${@---force}" +} + +starts() { + wc -l <"$CODEX_STARTS" 2>/dev/null || echo 0 +} + +fresh=$(collect ok) +[[ $(jq -c '{limitsStale, fresh: (.limitsFetchedAt > 0), used: [.limits[].percent], tierLabel, usageStatusText, authHelpText}' <<<"$fresh") == '{"limitsStale":false,"fresh":true,"used":[0.4,0.7],"tierLabel":"pro","usageStatusText":"","authHelpText":""}' ]] || + fail "fresh Codex limits are live and say when they were measured" "$fresh" +pass "fresh Codex limits are live and say when they were measured" +fetched_at=$(jq '.limitsFetchedAt' <<<"$fresh") + +# A panel opened again within the probe interval reuses the answer rather than +# starting another app-server; a person pressing refresh does not. +before=$(starts) +reused=$(collect backend --limits-only) +[[ $(starts) == "$before" ]] || fail "a Codex probe this recent is reused" "$(cat "$CODEX_STARTS")" +[[ $(jq -c --argjson at "$fetched_at" '{limitsStale, at: (.limitsFetchedAt == $at), used: [.limits[].percent], tierLabel, authHelpText}' <<<"$reused") == '{"limitsStale":false,"at":true,"used":[0.4,0.7],"tierLabel":"pro","authHelpText":""}' ]] || + fail "a reused Codex probe reads as current" "$reused" +collect ok >/dev/null +[[ $(starts) == $((before + 1)) ]] || fail "--force probes Codex past the reuse window" "$(cat "$CODEX_STARTS")" +pass "a recent Codex probe is reused, and --force probes again" +fetched_at=$(jq '.fetchedAtMs' <"$(ls "$XDG_CACHE_HOME"/omarchy/agent-usage/codex-limits-*.json)") + +# Let the weekly window reset before the next probe fails. +sleep 3 + +for mode in backend died crashed; do + kept=$(collect "$mode") + [[ $(jq -c --argjson at "$fetched_at" '{limitsStale, at: (.limitsFetchedAt == $at), used: [.limits[].percent * 100 | round], resets: [.limits[].resetsAt == ""], status: .usageStatusText, tierLabel}' <<<"$kept") == '{"limitsStale":true,"at":true,"used":[40,0],"resets":[false,true],"status":"","tierLabel":"pro"}' ]] || + fail "a failed Codex probe ($mode) keeps the last limits, stale, with an elapsed window at zero" "$kept" + [[ $(jq -r '.authHelpText' <<<"$kept") != "" ]] || fail "a failed Codex probe ($mode) says what went wrong" "$kept" +done +pass "a failed Codex probe keeps the last limits, stale, with an elapsed window at zero" + +# Only a network that isn't there yet asks the panel to retry soon: a stalled +# app-server, or one that died unable to reach OpenAI. A refusal, or a crash +# of its own, would only fail again. +for mode in stall died; do + retry=$(collect "$mode") + [[ $(jq -c '{retryAdvised, limitsStale, kept: (.limits | length)}' <<<"$retry") == '{"retryAdvised":true,"limitsStale":true,"kept":2}' ]] || + fail "an unreachable Codex backend ($mode) advises a retry" "$retry" +done +for mode in backend crashed auth; do + refused=$(collect "$mode") + [[ $(jq -r '.retryAdvised // false' <<<"$refused") == false ]] || + fail "a Codex backend that answered ($mode) advises no retry" "$refused" +done +pass "only an unreachable Codex backend advises a retry" + +auth=$(collect auth) +[[ $(jq -c '{status: .usageStatusText, limitsStale, kept: (.limits | length)}' <<<"$auth") == '{"status":"Waiting for auth","limitsStale":true,"kept":2}' ]] || + fail "a Codex sign-out is reported beside the last limits" "$auth" +pass "a Codex sign-out is reported beside the last limits" + +for who in "ws-two user-a" "ws-one user-b"; do + login $who + other=$(collect backend) + [[ $(jq -c '{limits, status: .usageStatusText}' <<<"$other") == '{"limits":[],"status":"Codex limits unavailable"}' ]] || + fail "a home signed in to someone else never shows the last one's limits ($who)" "$other" +done +pass "a home signed in to someone else, in another workspace or the same one, never shows the last one's limits" +login ws-one user-a + +# The login auth.json names is only who answers when nothing overrides it. +token=$(CODEX_ACCESS_TOKEN=another collect backend) +[[ $(jq -c '.limits' <<<"$token") == '[]' ]] || + fail "an access token in the environment never shows auth.json's account's limits" "$token" +pass "an access token in the environment never shows auth.json's account's limits" + +printf 'cli_auth_credentials_store = "keyring"\n' >"$HOME/.codex/config.toml" +keyring=$(collect backend) +[[ $(jq -c '.limits' <<<"$keyring") == '[]' ]] || + fail "a keyring login, which can't be told apart, keeps no limits" "$keyring" +rm "$HOME/.codex/config.toml" +pass "a keyring login, which can't be told apart, keeps no limits" + +file=$(collect backend) +[[ $(jq -c '{limitsStale, used: (.limits[0].percent * 100 | round)}' <<<"$file") == '{"limitsStale":true,"used":40}' ]] || + fail "the file store's own account keeps its limits" "$file" +pass "the file store's own account keeps its limits" + +# Signed in to someone else while the probe ran: the answer is shown, but kept +# for neither login. +switched=$(collect switch) +[[ $(jq '.limitsStale' <<<"$switched") == false ]] || fail "a probe across a login change still reports what it got" "$switched" +after=$(collect backend) +[[ $(jq -c '.limits' <<<"$after") == '[]' ]] || + fail "limits from a probe across a login change are kept for neither login" "$after" +login ws-one user-a +before=$(collect backend) +[[ $(jq '.limits[0].percent * 100 | round' <<<"$before") == 40 ]] || + fail "limits from a probe across a login change are kept for neither login" "$before" +pass "limits from a probe across a login change are kept for neither login" + +login ws-one $'\xed\xa0\x80' +for auth in '{"tokens":["unexpected"]}' '{"tokens":{"account_id":"ws-one"}}' "$(cat "$HOME/.codex/auth.json")"; do + printf '%s\n' "$auth" >"$HOME/.codex/auth.json" + odd=$(collect backend) || fail "an auth.json of another shape still gives a record ($auth)" + [[ $(jq -r '.usageStatusText' <<<"$odd") == "Codex limits unavailable" ]] || + fail "an auth.json of another shape still gives a record ($auth)" "$odd" +done +pass "an auth.json of another shape still gives a record" diff --git a/test/shell.d/agent-usage-codex-scanner-test.sh b/test/shell.d/agent-usage-codex-scanner-test.sh index 81a36373509..3200105ce36 100644 --- a/test/shell.d/agent-usage-codex-scanner-test.sh +++ b/test/shell.d/agent-usage-codex-scanner-test.sh @@ -905,8 +905,8 @@ chmod +x "$LAUNCH_HOME/bin/mise" result=$(HOME="$LAUNCH_HOME" CODEX_HOME="$LAUNCH_HOME/.codex" MISE_CALLS_FILE="$LAUNCH_HOME/mise-calls" XDG_DATA_HOME="$LAUNCH_HOME/.local/share" \ PATH="$LAUNCH_HOME/bin:$SAFE_PATH" "$collector" --limits-only) -[[ $(jq -r '.usageStatusText' <<<"$result") == "Codex unavailable" ]] || - fail "Codex collector reports Codex unavailable when only the launcher exists" "$result" +[[ $(jq -r '.usageStatusText' <<<"$result") == "Codex limits unavailable" ]] || + fail "Codex collector reports its limits unavailable when only the launcher exists" "$result" # The launcher would log `use -g ...` through the mise stub if it ever ran; # `which codex` is the only permitted call. [[ ! -s $LAUNCH_HOME/mise-calls || $(cat "$LAUNCH_HOME/mise-calls") == "which codex" ]] || @@ -966,8 +966,8 @@ ln -s "$SHIM_HOME/bin/mise" "$SHIM_HOME/.local/share/mise/shims/codex" result=$(HOME="$SHIM_HOME" CODEX_HOME="$SHIM_HOME/.codex" MISE_CALLS_FILE="$SHIM_HOME/mise-calls" XDG_DATA_HOME="$SHIM_HOME/.local/share" \ PATH="$SHIM_HOME/bin:$SAFE_PATH" "$collector" --limits-only) -[[ $(jq -r '.usageStatusText' <<<"$result") == "Codex unavailable" ]] || - fail "Codex collector reports Codex unavailable when only a mise shim exists" "$result" +[[ $(jq -r '.usageStatusText' <<<"$result") == "Codex limits unavailable" ]] || + fail "Codex collector reports its limits unavailable when only a mise shim exists" "$result" [[ ! -s $SHIM_HOME/mise-calls || $(cat "$SHIM_HOME/mise-calls") == "which codex" ]] || fail "Codex collector must not execute a mise shim" "$(cat "$SHIM_HOME/mise-calls" 2>/dev/null)" pass "Codex collector treats a shim symlink to mise as lazy" From d16b3fa92ad76d8dd1b998276268e38ea3957036 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:08:01 -0400 Subject: [PATCH 070/546] Leave Codex out of the panel on a machine that never ran it The Codex record was always ready, so a machine with no Codex sign-in and no local usage still showed a Codex tab asking for auth. It is ready once there is local usage, limits, or a registered account, as the other collectors are. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-codex | 4 +++- test/shell.d/agent-usage-codex-limits-test.sh | 8 +++++++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/bin/omarchy-agent-usage-codex b/bin/omarchy-agent-usage-codex index 7cecd34982d..1cccb52461c 100755 --- a/bin/omarchy-agent-usage-codex +++ b/bin/omarchy-agent-usage-codex @@ -1375,7 +1375,9 @@ def main(): "id": AGENT_ID, "name": AGENT_NAME, "updatedAt": datetime.now(timezone.utc).isoformat(), - "ready": True, + # In use here once it has local usage or limits to show; a machine that + # never ran Codex prints an empty record the panel skips. + "ready": number(stats.get("totalPrompts")) > 0 or bool(rpc["limits"]) or bool(accounts), "hasLocalStats": True, } record.update(stats) diff --git a/test/shell.d/agent-usage-codex-limits-test.sh b/test/shell.d/agent-usage-codex-limits-test.sh index a68ea099696..9957a29c1ee 100644 --- a/test/shell.d/agent-usage-codex-limits-test.sh +++ b/test/shell.d/agent-usage-codex-limits-test.sh @@ -146,6 +146,12 @@ chmod +x "$collector" mkdir -p "$test_tmp/bin" "$test_tmp/empty-bin" +# A machine that never ran Codex prints an empty record the panel skips. +unused=$(PATH="$test_tmp/empty-bin:/usr/bin:/bin" "$collector" --force) +[[ $(jq -c '{ready, limits, limitsStale, limitsFetchedAt, totalPrompts, todayTotalTokens, recentDays: (.recentDays | length), modelUsage, retry: (.retryAdvised // false)}' <<<"$unused") == '{"ready":false,"limits":[],"limitsStale":true,"limitsFetchedAt":0,"totalPrompts":0,"todayTotalTokens":0,"recentDays":7,"modelUsage":{},"retry":false}' ]] || + fail "a machine that never ran Codex prints an empty record" "$unused" +pass "a machine that never ran Codex prints an empty record" + mkdir -p "$HOME/.codex/sessions" # Signs ~/.codex in as a user of a workspace, the way Codex writes it: the @@ -219,7 +225,7 @@ starts() { } fresh=$(collect ok) -[[ $(jq -c '{limitsStale, fresh: (.limitsFetchedAt > 0), used: [.limits[].percent], tierLabel, usageStatusText, authHelpText}' <<<"$fresh") == '{"limitsStale":false,"fresh":true,"used":[0.4,0.7],"tierLabel":"pro","usageStatusText":"","authHelpText":""}' ]] || +[[ $(jq -c '{ready, limitsStale, fresh: (.limitsFetchedAt > 0), used: [.limits[].percent], tierLabel, usageStatusText, authHelpText}' <<<"$fresh") == '{"ready":true,"limitsStale":false,"fresh":true,"used":[0.4,0.7],"tierLabel":"pro","usageStatusText":"","authHelpText":""}' ]] || fail "fresh Codex limits are live and say when they were measured" "$fresh" pass "fresh Codex limits are live and say when they were measured" fetched_at=$(jq '.limitsFetchedAt' <<<"$fresh") From c2086405d435700272d1caa8e96ba153e765a515 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:09:03 -0400 Subject: [PATCH 071/546] Skip untrusted Codex timestamps and token counts, not the record A pi timestamp out of range raised past the scan and lost the whole Codex record, and a negative or boolean token count was summed as if it were one. A time that isn't one now dates its row today, and a count that isn't a positive finite number counts as nothing. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-codex | 18 ++++++++++++++---- .../shell.d/agent-usage-codex-scanner-test.sh | 19 +++++++++++++++++++ 2 files changed, 33 insertions(+), 4 deletions(-) diff --git a/bin/omarchy-agent-usage-codex b/bin/omarchy-agent-usage-codex index 1cccb52461c..993a559b13f 100755 --- a/bin/omarchy-agent-usage-codex +++ b/bin/omarchy-agent-usage-codex @@ -67,14 +67,19 @@ TRANSPORT_ERRORS = ( FILE_CACHE_SCHEMA = 1 +# Timestamps on disk are untrusted: one that is no time at all (a boolean, +# NaN, or out of range) dates its row today rather than stopping the scan. def local_day(value): - if value is None: + if value is None or isinstance(value, bool): return datetime.now().strftime("%Y-%m-%d") if isinstance(value, (int, float)): # pi message timestamps are milliseconds; Codex timestamps are usually seconds. if value > 10_000_000_000: value = value / 1000 - return datetime.fromtimestamp(value).strftime("%Y-%m-%d") + try: + return datetime.fromtimestamp(value).strftime("%Y-%m-%d") + except (OverflowError, OSError, ValueError): + return datetime.now().strftime("%Y-%m-%d") text = str(value) try: if text.endswith("Z"): @@ -88,11 +93,16 @@ def local_day(value): return datetime.now().strftime("%Y-%m-%d") +# Counts from disk and the app-server are untrusted: a boolean, a negative, +# or a value that isn't a finite number counts as nothing. def number(value): + if isinstance(value, bool): + return 0 try: - return int(value or 0) - except Exception: + n = float(value or 0) + except (TypeError, ValueError, OverflowError): return 0 + return int(n) if math.isfinite(n) and n > 0 else 0 def model_name(raw): diff --git a/test/shell.d/agent-usage-codex-scanner-test.sh b/test/shell.d/agent-usage-codex-scanner-test.sh index 3200105ce36..8e8edd01dae 100644 --- a/test/shell.d/agent-usage-codex-scanner-test.sh +++ b/test/shell.d/agent-usage-codex-scanner-test.sh @@ -232,6 +232,25 @@ result=$(HOME="$IDLESS_PI_HOME" CODEX_HOME="$IDLESS_PI_HOME/.codex" XDG_CACHE_HO fail "Codex collector counts every pi line without an id, once per line" "$result" pass "Codex collector counts every pi line without an id, once per line" +# Timestamps and token counts on disk are untrusted: a time out of range +# dates its row today, and a negative, boolean or non-numeric count is no +# count, without losing the rest of the record. +ODD_PI_HOME=$(signed_in_home) +cp "$TEST_HOME/bin/codex" "$ODD_PI_HOME/bin/codex" +mkdir -p "$ODD_PI_HOME/.pi/agent/sessions/project" +cat >"$ODD_PI_HOME/.pi/agent/sessions/project/odd.jsonl" < Date: Wed, 14 Oct 2026 23:10:04 -0400 Subject: [PATCH 072/546] Zero Grok's elapsed windows without inventing their next reset A failed Grok check served its cached credits as they were, so a week that had reset since still read as spent. Every path that serves kept numbers, the failed check included, now brings an elapsed window back at zero. Its next reset stays unknown until xAI answers, rather than being computed a whole number of weeks on, since reset times are only ever the service's own. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-grok | 38 ++++++++++++++------------- test/shell.d/agent-usage-grok-test.sh | 34 +++++++++++++++++++----- 2 files changed, 48 insertions(+), 24 deletions(-) diff --git a/bin/omarchy-agent-usage-grok b/bin/omarchy-agent-usage-grok index f2d4d7fe340..fa1997af42b 100755 --- a/bin/omarchy-agent-usage-grok +++ b/bin/omarchy-agent-usage-grok @@ -193,20 +193,20 @@ def credits_limit(data): } -# A window whose reset has passed starts over empty. A weekly one's next -# reset is a whole number of weeks later; a monthly one's isn't known. -def rolled_forward(limits): +# A kept window whose reset has passed starts over empty. Its next reset is +# xAI's to say, so it stays unknown until a probe answers. +def usable_cached_limits(cached): + limits = cached.get("limits") + if not isinstance(limits, list): + return [] now = datetime.now(timezone.utc) - rolled = [] + usable = [] for limit in limits: + if not isinstance(limit, dict): + continue resets = parse_time(limit.get("resetsAt")) - if resets and resets <= now: - limit = dict(limit, percent=0.0, resetsAt="") - if limit.get("label") == "Weekly": - weeks = (now - resets) // timedelta(weeks=1) + 1 - limit["resetsAt"] = (resets + weeks * timedelta(weeks=1)).isoformat() - rolled.append(limit) - return rolled + usable.append(dict(limit, percent=0.0, resetsAt="") if resets and resets <= now else limit) + return usable def collect_limits(home, key, force): @@ -218,19 +218,21 @@ def collect_limits(home, key, force): return result cache = cache_root() / f"grok-limits-{key}.json" - cached = read_json(cache) or {} + cached = read_json(cache) + cached = cached if isinstance(cached, dict) else {} + fallback = usable_cached_limits(cached) fetched_at = float(cached.get("fetchedAtMs") or 0) result["fetchedAtMs"] = fetched_at expires = parse_time(entry.get("expires_at")) if expires and expires.timestamp() <= time.time(): - result["limits"] = cached.get("limits") or [] + result["limits"] = fallback # The access token lives six hours and Grok mints a new one from its # refresh token whenever it starts, so a lapsed one is routine, not a # sign-out. And while Grok hasn't run, nothing here has spent its - # allowance, so the last numbers still stand, rolled past any reset. + # allowance, so the last numbers still stand, an elapsed window at zero. if entry.get("refresh_token") and time.time() - expires.timestamp() < REFRESH_LIFETIME_SECONDS: - result.update(limits=rolled_forward(result["limits"]), live=True) + result["live"] = True if not result["limits"]: result["usageStatusText"] = "Limits paused" result["authHelpText"] = "Start Grok to see its limits; signing in again is not needed." @@ -239,8 +241,8 @@ def collect_limits(home, key, force): result["authHelpText"] = "Grok's saved sign-in expired. Start Grok to refresh it." return result - if cached.get("limits") and not force and time.time() - fetched_at / 1000 < PROBE_REUSE_SECONDS: - result.update(limits=cached["limits"], live=True, authHelpText="") + if fallback and not force and time.time() - fetched_at / 1000 < PROBE_REUSE_SECONDS: + result.update(limits=fallback, live=True, authHelpText="") return result request = urllib.request.Request(CREDITS_URL, headers={"Authorization": f"Bearer {token}", "Accept": "application/json"}) @@ -258,7 +260,7 @@ def collect_limits(home, key, force): result.update(limits=[limit], live=True, fetchedAtMs=round(time.time() * 1000), authHelpText="") write_json(cache, {"fetchedAtMs": result["fetchedAtMs"], "limits": result["limits"]}) else: - result["limits"] = cached.get("limits") or [] + result["limits"] = fallback if not result["usageStatusText"] and not result["limits"]: result["usageStatusText"] = "Grok limits unavailable" return result diff --git a/test/shell.d/agent-usage-grok-test.sh b/test/shell.d/agent-usage-grok-test.sh index eeaf985eaf3..dbd6c3d47c6 100755 --- a/test/shell.d/agent-usage-grok-test.sh +++ b/test/shell.d/agent-usage-grok-test.sh @@ -43,6 +43,11 @@ end = os.environ["PERIOD_END"] # percentage of zero. def urlopen(request, timeout=None): assert request.full_url == collector.CREDITS_URL + failure = os.environ.get("GROK_FAIL") + if failure == "unreachable": + raise collector.urllib.error.URLError(OSError("Name or service not known")) + if failure: + raise collector.urllib.error.HTTPError(request.full_url, int(failure), "refused", {}, None) token = request.get_header("Authorization").split(" ", 1)[1] config = {"currentPeriod": {"type": "USAGE_PERIOD_TYPE_WEEKLY", "end": end}, "billingPeriodEnd": end} if used.get(token): @@ -96,17 +101,16 @@ record=$(collect) fail "a lapsed access token with a refresh token keeps showing the last numbers" "$record" pass "a lapsed access token with a refresh token keeps showing the last numbers" -# A week that reset while Grok sat idle starts over at 0%, a week later. +# A week that reset while Grok sat idle starts over at 0%, its next reset +# unknown until xAI says. cache=$(ls "$XDG_CACHE_HOME"/omarchy/agent-usage/grok-limits-*.json | head -1) reset_past=$(python3 -c 'import datetime as dt; print((dt.datetime.now(dt.timezone.utc) - dt.timedelta(hours=1)).isoformat())') jq --arg at "$reset_past" '.limits[0].resetsAt = $at' "$cache" >"$test_tmp/cache.json" mv "$test_tmp/cache.json" "$cache" record=$(collect) -next_reset=$(jq -r '.limits[0].resetsAt' <<<"$record") -[[ $(jq -r '.limits[0].percent' <<<"$record") == 0.0 ]] && - python3 -c 'import datetime as dt, sys; n = dt.datetime.fromisoformat(sys.argv[1]); d = n - dt.datetime.now(dt.timezone.utc); sys.exit(0 if dt.timedelta(days=6) < d < dt.timedelta(days=7) else 1)' "$next_reset" || - fail "a week that reset while Grok was idle starts over a week later" "$record" -pass "a week that reset while Grok was idle starts over a week later" +[[ $(jq -c '.limits[0] | {percent, resetsAt}' <<<"$record") == '{"percent":0.0,"resetsAt":""}' ]] || + fail "a week that reset while Grok was idle starts over at zero" "$record" +pass "a week that reset while Grok was idle starts over at zero" # Without cached numbers there's nothing to show, so it says how to get them. rm -f "$XDG_CACHE_HOME"/omarchy/agent-usage/grok-limits-*.json @@ -124,6 +128,24 @@ record=$(collect) fail "a sign-in lapsed past Grok's 30 days asks for a sign-in" "$record" pass "a sign-in lapsed past Grok's 30 days asks for a sign-in" +# A failed check keeps the last numbers, stale, and a window that has reset +# since comes back at zero rather than at what it read then. +signed_in "$HOME/.grok" token-main u-main "$future" "X Premium+" +rm -f "$XDG_CACHE_HOME"/omarchy/agent-usage/grok-limits-*.json +fresh=$(collect) +fetched_at=$(jq '.limitsFetchedAt' <<<"$fresh") +record=$(GROK_FAIL=500 collect) +[[ $(jq -c --argjson at "$fetched_at" '{stale: .limitsStale, at: (.limitsFetchedAt == $at), percent: .limits[0].percent, status: .usageStatusText}' <<<"$record") == '{"stale":true,"at":true,"percent":0.42,"status":""}' ]] || + fail "a failed Grok check keeps the last credits, stale" "$record" +pass "a failed Grok check keeps the last credits, stale" +cache=$(ls "$XDG_CACHE_HOME"/omarchy/agent-usage/grok-limits-*.json | head -1) +jq --arg at "$reset_past" '.limits[0].resetsAt = $at' "$cache" >"$test_tmp/cache.json" +mv "$test_tmp/cache.json" "$cache" +record=$(GROK_FAIL=500 collect) +[[ $(jq -c '{stale: .limitsStale, window: (.limits[0] | {percent, resetsAt})}' <<<"$record") == '{"stale":true,"window":{"percent":0.0,"resetsAt":""}}' ]] || + fail "a failed Grok check zeroes a kept window that has reset since" "$record" +pass "a failed Grok check zeroes a kept window that has reset since" + # A period with nothing used yet comes without a percentage. signed_in "$HOME/.grok" token-fresh u-main "$future" "X Premium+" rm -f "$XDG_CACHE_HOME"/omarchy/agent-usage/grok-limits-*.json From 5b49b74962c8a9f884b8dbf0de54bf7d66b826fe Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:10:54 -0400 Subject: [PATCH 073/546] Tell Grok's unreachable, refused and garbled checks apart Every failed Grok check looked the same: no retry was advised when the endpoint couldn't be reached, a refused token read as "Waiting for auth", and a NaN percentage passed the clamp as a full allowance. An endpoint no server answered for now sets retryAdvised, a 401 or 403 reads as "Sign-in expired", a percentage that isn't a number is no answer, and the help text says which of these happened. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-grok | 38 ++++++++++++++++++++++----- test/shell.d/agent-usage-grok-test.sh | 25 ++++++++++++++++++ 2 files changed, 57 insertions(+), 6 deletions(-) diff --git a/bin/omarchy-agent-usage-grok b/bin/omarchy-agent-usage-grok index fa1997af42b..bff74906763 100755 --- a/bin/omarchy-agent-usage-grok +++ b/bin/omarchy-agent-usage-grok @@ -14,6 +14,7 @@ by model. The agents panel only ever reads the JSON this prints. import argparse import json +import math import os import re import tempfile @@ -186,9 +187,16 @@ def plan(home): def credits_limit(data): config = data["config"] period = config["currentPeriod"] + used = config.get("creditUsagePercent") or 0 + if isinstance(used, bool): + raise ValueError("creditUsagePercent is not a number") + used = float(used) + # NaN would pass any clamp as 100%; a percentage that isn't one is no answer. + if not math.isfinite(used): + raise ValueError("creditUsagePercent is not a number") return { "label": "Weekly" if "WEEK" in str(period.get("type")) else "Monthly", - "percent": max(0.0, min(1.0, float(config.get("creditUsagePercent") or 0) / 100)), + "percent": max(0.0, min(1.0, used / 100)), "resetsAt": parse_time(period["end"]).isoformat(), } @@ -246,15 +254,30 @@ def collect_limits(home, key, force): return result request = urllib.request.Request(CREDITS_URL, headers={"Authorization": f"Bearer {token}", "Accept": "application/json"}) + limit = None try: with urllib.request.urlopen(request, timeout=8) as response: - limit = credits_limit(json.loads(response.read())) + body = response.read() except urllib.error.HTTPError as error: - limit = None + body = None if error.code in (401, 403): - result["usageStatusText"] = "Waiting for auth" - except Exception: - limit = None + # The token was turned away before it lapsed: a sign-in to redo. + result["usageStatusText"] = "Sign-in expired" + result["authHelpText"] = "xAI refused Grok's saved sign-in. Start Grok, or run `grok login`, to sign in again." + else: + result["authHelpText"] = f"Grok's usage endpoint returned status {error.code}. Local Grok stats are still shown." + except (urllib.error.URLError, OSError): + # No server answered at all (no route, no DNS, a timeout), often a + # network not up yet after login, so the panel asks again soon. Any real + # answer, an error status included, is a server to stop pestering. + body = None + result["retryAdvised"] = True + result["authHelpText"] = "Couldn't reach Grok's usage endpoint. Retrying shortly. Local Grok stats are still shown." + if body is not None: + try: + limit = credits_limit(json.loads(body)) + except Exception: + result["authHelpText"] = "Grok's usage endpoint answered without a usable limit. Local Grok stats are still shown." if limit: result.update(limits=[limit], live=True, fetchedAtMs=round(time.time() * 1000), authHelpText="") @@ -313,6 +336,7 @@ def account_limits(account, force): "fetchedAt": limits["fetchedAtMs"], "usageStatusText": limits["usageStatusText"], "authHelpText": limits["authHelpText"], + "retryAdvised": bool(limits.get("retryAdvised")), } @@ -352,6 +376,8 @@ def main(): "limitsStale": stale, "limitsFetchedAt": fetched_at, } + if limits.get("retryAdvised"): + record["retryAdvised"] = True record.update(cached_local_stats(stats_home, LIMITS_ONLY_SCAN_REUSE_SECONDS if args.limits_only and not args.force else 0)) if accounts: record["accountSwitch"] = registered[0]["switch"] diff --git a/test/shell.d/agent-usage-grok-test.sh b/test/shell.d/agent-usage-grok-test.sh index dbd6c3d47c6..04b96b70117 100755 --- a/test/shell.d/agent-usage-grok-test.sh +++ b/test/shell.d/agent-usage-grok-test.sh @@ -52,6 +52,8 @@ def urlopen(request, timeout=None): config = {"currentPeriod": {"type": "USAGE_PERIOD_TYPE_WEEKLY", "end": end}, "billingPeriodEnd": end} if used.get(token): config["creditUsagePercent"] = used[token] + if os.environ.get("GROK_PERCENT"): + config["creditUsagePercent"] = json.loads(os.environ["GROK_PERCENT"]) return io.BytesIO(json.dumps({"config": config}).encode()) collector.urllib.request.urlopen = urlopen @@ -146,6 +148,29 @@ record=$(GROK_FAIL=500 collect) fail "a failed Grok check zeroes a kept window that has reset since" "$record" pass "a failed Grok check zeroes a kept window that has reset since" +# Only an endpoint that couldn't be reached at all asks the panel to retry +# soon; one that answered with an error would only answer the same again. +record=$(GROK_FAIL=unreachable collect) +[[ $(jq -c '{retryAdvised, stale: .limitsStale, kept: (.limits | length), status: .usageStatusText}' <<<"$record") == '{"retryAdvised":true,"stale":true,"kept":1,"status":""}' ]] || + fail "an unreachable Grok endpoint advises a retry and keeps the last credits" "$record" +record=$(GROK_FAIL=500 collect) +[[ $(jq -r '.retryAdvised // false' <<<"$record") == false ]] || + fail "a Grok endpoint that answered with an error advises no retry" "$record" +pass "only an unreachable Grok endpoint advises a retry" + +# A token turned away before it lapsed is a sign-in to redo. +record=$(GROK_FAIL=401 collect) +[[ $(jq -c '{status: .usageStatusText, stale: .limitsStale, retry: (.retryAdvised // false)}' <<<"$record") == '{"status":"Sign-in expired","stale":true,"retry":false}' ]] || + fail "a refused Grok sign-in says it expired" "$record" +pass "a refused Grok sign-in says it expired" + +# A percentage that isn't a number is no answer, not a full allowance. +rm -f "$XDG_CACHE_HOME"/omarchy/agent-usage/grok-limits-*.json +record=$(GROK_PERCENT=NaN collect) +[[ $(jq -c '{status: .usageStatusText, limits}' <<<"$record") == '{"status":"Grok limits unavailable","limits":[]}' ]] || + fail "a Grok percentage that isn't a number reads as unavailable" "$record" +pass "a Grok percentage that isn't a number reads as unavailable" + # A period with nothing used yet comes without a percentage. signed_in "$HOME/.grok" token-fresh u-main "$future" "X Premium+" rm -f "$XDG_CACHE_HOME"/omarchy/agent-usage/grok-limits-*.json From 8db5467b3fcbc7f0bdd336989f0839d471fcc990 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:11:58 -0400 Subject: [PATCH 074/546] Reuse Grok's scan and probe for the collector contract's windows Grok reused a probe for 60 seconds, after its CLI's once-a-minute watch rather than anything xAI's endpoint asks for, rescanned on every normal run, and trusted a cache whose time was in the future. A probe is now reused for PROBE_MIN_INTERVAL_SECONDS (15) and a scan for SCAN_REUSE_SECONDS (20), or LIMITS_ONLY_REUSE_SECONDS (900) on a limits-only refresh, each only while 0 <= age < N, with --force reusing neither. The scan is rebuilt under a lock and keyed by its scanDate. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-grok | 66 ++++++++++++++++++++------- test/shell.d/agent-usage-grok-test.sh | 37 ++++++++++++++- 2 files changed, 86 insertions(+), 17 deletions(-) diff --git a/bin/omarchy-agent-usage-grok b/bin/omarchy-agent-usage-grok index bff74906763..b0f42d5040f 100755 --- a/bin/omarchy-agent-usage-grok +++ b/bin/omarchy-agent-usage-grok @@ -13,6 +13,7 @@ by model. The agents panel only ever reads the JSON this prints. """ import argparse +import fcntl import json import math import os @@ -28,14 +29,17 @@ AGENT_ID = "grok" AGENT_NAME = "Grok" AUTH_HELP = "Start Grok, or run `grok login`, to sign in." CREDITS_URL = "https://cli-chat-proxy.grok.com/v1/billing?format=credits" -# Grok's own CLI watches its subscription once a minute; a panel opened and -# shut repeatedly reuses an answer for about that long. -PROBE_REUSE_SECONDS = 60 +# A panel opened and shut repeatedly reuses an answer this recent rather than +# asking xAI again. +PROBE_MIN_INTERVAL_SECONDS = 15 # Grok's docs give a saved sign-in 30 days; a refresh token older than that # can't renew the access token, so it no longer counts as signed in. REFRESH_LIFETIME_SECONDS = 30 * 24 * 3600 -# A limits-only refresh needs fresh limits, not a fresh history scan. -LIMITS_ONLY_SCAN_REUSE_SECONDS = 900 +# A local scan this recent is reused: a normal run only dedups overlapping +# collector runs, while a limits-only refresh needs fresh limits, not a fresh +# history scan. +SCAN_REUSE_SECONDS = 20 +LIMITS_ONLY_REUSE_SECONDS = 900 def read_json(path): @@ -151,16 +155,40 @@ def local_stats(home): } -# A scan from before midnight would count yesterday's sessions as today's. -def cached_local_stats(home, max_age): - cache = cache_root() / "grok-stats.json" - cached = read_json(cache) or {} +# A scan from before midnight would count yesterday's sessions as today's, +# and one from a clock that has since gone back can't say how old it is. +def fresh_stats(cache, home, max_age): + cached = read_json(cache) + if not isinstance(cached, dict) or not isinstance(cached.get("stats"), dict): + return None today = datetime.now().astimezone().date().isoformat() - if cached.get("home") == str(home) and cached.get("day") == today and time.time() - float(cached.get("at") or 0) < max_age: + try: + age = time.time() - float(cached.get("at") or 0) + except (TypeError, ValueError): + return None + if cached.get("home") == str(home) and cached.get("scanDate") == today and 0 <= age < max_age: return cached["stats"] - stats = local_stats(home) - write_json(cache, {"home": str(home), "day": today, "at": time.time(), "stats": stats}) - return stats + return None + + +# Rebuilt under a lock, so overlapping runs (the update command backgrounds +# one per agent while the panel refreshes on its own) scan once between them. +def cached_local_stats(home, max_age): + root = cache_root() + root.mkdir(parents=True, exist_ok=True) + cache = root / "grok-stats.json" + stats = fresh_stats(cache, home, max_age) + if stats is not None: + return stats + with (root / "grok-stats.lock").open("w") as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + stats = fresh_stats(cache, home, max_age) + if stats is not None: + return stats + stats = local_stats(home) + today = datetime.now().astimezone().date().isoformat() + write_json(cache, {"home": str(home), "scanDate": today, "at": time.time(), "stats": stats}) + return stats # One login per issuer in auth.json; the first is the CLI's own. @@ -229,7 +257,11 @@ def collect_limits(home, key, force): cached = read_json(cache) cached = cached if isinstance(cached, dict) else {} fallback = usable_cached_limits(cached) - fetched_at = float(cached.get("fetchedAtMs") or 0) + try: + fetched_at = float(cached.get("fetchedAtMs") or 0) + except (TypeError, ValueError): + fetched_at = 0 + fetched_at = round(fetched_at) if math.isfinite(fetched_at) else 0 result["fetchedAtMs"] = fetched_at expires = parse_time(entry.get("expires_at")) @@ -249,7 +281,8 @@ def collect_limits(home, key, force): result["authHelpText"] = "Grok's saved sign-in expired. Start Grok to refresh it." return result - if fallback and not force and time.time() - fetched_at / 1000 < PROBE_REUSE_SECONDS: + fetched_age = time.time() - fetched_at / 1000 + if fallback and not force and 0 <= fetched_age < PROBE_MIN_INTERVAL_SECONDS: result.update(limits=fallback, live=True, authHelpText="") return result @@ -378,7 +411,8 @@ def main(): } if limits.get("retryAdvised"): record["retryAdvised"] = True - record.update(cached_local_stats(stats_home, LIMITS_ONLY_SCAN_REUSE_SECONDS if args.limits_only and not args.force else 0)) + scan_age = 0 if args.force else (LIMITS_ONLY_REUSE_SECONDS if args.limits_only else SCAN_REUSE_SECONDS) + record.update(cached_local_stats(stats_home, scan_age)) if accounts: record["accountSwitch"] = registered[0]["switch"] record["accounts"] = accounts diff --git a/test/shell.d/agent-usage-grok-test.sh b/test/shell.d/agent-usage-grok-test.sh index 04b96b70117..2bad21cb87f 100755 --- a/test/shell.d/agent-usage-grok-test.sh +++ b/test/shell.d/agent-usage-grok-test.sh @@ -43,6 +43,9 @@ end = os.environ["PERIOD_END"] # percentage of zero. def urlopen(request, timeout=None): assert request.full_url == collector.CREDITS_URL + if os.environ.get("GROK_PROBES"): + with open(os.environ["GROK_PROBES"], "a") as log: + log.write("probe\n") failure = os.environ.get("GROK_FAIL") if failure == "unreachable": raise collector.urllib.error.URLError(OSError("Name or service not known")) @@ -220,8 +223,40 @@ pass "a limits-only refresh reuses the session scan" # A scan from another day is never reused, so yesterday's sessions don't # count as today's after midnight. -jq '.day = "2000-01-01"' "$XDG_CACHE_HOME/omarchy/agent-usage/grok-stats.json" >"$test_tmp/stats.json" +jq '.scanDate = "2000-01-01"' "$XDG_CACHE_HOME/omarchy/agent-usage/grok-stats.json" >"$test_tmp/stats.json" mv "$test_tmp/stats.json" "$XDG_CACHE_HOME/omarchy/agent-usage/grok-stats.json" record=$(COLLECT_ARGS="--limits-only" collect) [[ $(jq -r '.totalSessions' <<<"$record") == 4 ]] || fail "a scan from another day is made again" "$record" pass "a scan from another day is made again" + +# A scan whose time is in the future came from a clock that has since gone +# back, so it can't say how old it is and is made again. +mkdir -p "$sessions/future" +jq -n --arg at "$now" '{last_active_at: $at, num_messages: 1}' >"$sessions/future/summary.json" +jq '.at = 99999999999' "$XDG_CACHE_HOME/omarchy/agent-usage/grok-stats.json" >"$test_tmp/stats.json" +mv "$test_tmp/stats.json" "$XDG_CACHE_HOME/omarchy/agent-usage/grok-stats.json" +record=$(COLLECT_ARGS="--limits-only" collect) +[[ $(jq -r '.totalSessions' <<<"$record") == 5 ]] || fail "a future-dated Grok scan is made again" "$record" +pass "a future-dated Grok scan is made again" + +# A normal run reuses a scan for a few seconds, so overlapping runs read once; +# --force reads again. +mkdir -p "$sessions/soon" +jq -n --arg at "$now" '{last_active_at: $at, num_messages: 1}' >"$sessions/soon/summary.json" +record=$(COLLECT_ARGS=" " collect) +[[ $(jq -r '.totalSessions' <<<"$record") == 5 ]] || fail "a normal Grok run reuses a scan this recent" "$record" +record=$(collect) +[[ $(jq -r '.totalSessions' <<<"$record") == 6 ]] || fail "--force scans Grok's sessions again" "$record" +pass "a normal Grok run reuses a recent scan, and --force scans again" + +# A probe this recent is reused rather than asking xAI again, until a person +# presses refresh. +rm -f "$XDG_STATE_HOME/omarchy/agents/accounts/grok.json" "$test_tmp/probes" +collect >/dev/null +record=$(GROK_PROBES="$test_tmp/probes" COLLECT_ARGS="--limits-only" collect) +[[ ! -e $test_tmp/probes && $(jq -c '{stale: .limitsStale, percent: .limits[0].percent}' <<<"$record") == '{"stale":false,"percent":0.42}' ]] || + fail "a Grok probe this recent is reused" "$record" +GROK_PROBES="$test_tmp/probes" collect >/dev/null +[[ $(wc -l <"$test_tmp/probes") == 1 ]] || fail "--force probes Grok past the reuse window" "$(cat "$test_tmp/probes")" +pass "a recent Grok probe is reused, and --force probes again" + From 77eb4c7685a92621104acfc39c6993148e9c682a Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:12:37 -0400 Subject: [PATCH 075/546] Count Grok as in use when it has sessions here Grok was ready only with a sign-in, so a machine whose sessions outlived its login dropped the tab and its history. Local sessions now make it ready too, reading "Waiting for auth" for the sign-in it lacks. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-grok | 12 +++++++----- test/shell.d/agent-usage-grok-test.sh | 11 +++++++++++ 2 files changed, 18 insertions(+), 5 deletions(-) diff --git a/bin/omarchy-agent-usage-grok b/bin/omarchy-agent-usage-grok index b0f42d5040f..a064784bdfa 100755 --- a/bin/omarchy-agent-usage-grok +++ b/bin/omarchy-agent-usage-grok @@ -392,16 +392,19 @@ def main(): limits = collect_limits(home, cache_key(home, "main"), args.force) stale, fetched_at = not limits["live"], limits["fetchedAtMs"] - # Nobody signed in to Grok here: an empty record, which the panel skips. # With accounts registered, the primary is ~/.grok by definition. stats_home = Path.home() / ".grok" if accounts else default_home() - signed_in = bool(accounts) or bool(login(stats_home)) + scan_age = 0 if args.force else (LIMITS_ONLY_REUSE_SECONDS if args.limits_only else SCAN_REUSE_SECONDS) + stats = cached_local_stats(stats_home, scan_age) + # In use here once someone signed in or ran a session; otherwise an empty + # record, which the panel skips. + in_use = bool(accounts) or bool(login(stats_home)) or stats["totalSessions"] > 0 record = { "schemaVersion": 1, "id": AGENT_ID, "name": AGENT_NAME, "updatedAt": datetime.now(timezone.utc).isoformat(), - "ready": signed_in, + "ready": in_use, "tierLabel": tier, "usageStatusText": limits["usageStatusText"], "authHelpText": limits["authHelpText"], @@ -411,8 +414,7 @@ def main(): } if limits.get("retryAdvised"): record["retryAdvised"] = True - scan_age = 0 if args.force else (LIMITS_ONLY_REUSE_SECONDS if args.limits_only else SCAN_REUSE_SECONDS) - record.update(cached_local_stats(stats_home, scan_age)) + record.update(stats) if accounts: record["accountSwitch"] = registered[0]["switch"] record["accounts"] = accounts diff --git a/test/shell.d/agent-usage-grok-test.sh b/test/shell.d/agent-usage-grok-test.sh index 2bad21cb87f..66631428cc8 100755 --- a/test/shell.d/agent-usage-grok-test.sh +++ b/test/shell.d/agent-usage-grok-test.sh @@ -69,8 +69,19 @@ export PERIOD_END="$period_end" record=$(collect) [[ $(jq -c '{ready, tierLabel, limits}' <<<"$record") == '{"ready":false,"tierLabel":"","limits":[]}' ]] || fail "a machine nobody signed in to Grok on gives an empty record" "$record" +[[ $(jq -c '{hasLocalStats, totalPrompts, totalSessions, todayTotalTokens, days: (.recentDays | length), modelUsage, status: .usageStatusText}' <<<"$record") == '{"hasLocalStats":false,"totalPrompts":0,"totalSessions":0,"todayTotalTokens":0,"days":7,"modelUsage":{},"status":"Waiting for auth"}' ]] || + fail "an empty Grok record still carries the full stats block" "$record" pass "a machine nobody signed in to Grok on gives an empty record" +# Sessions on disk put Grok in use here even with no sign-in, which is then +# what's missing. +mkdir -p "$test_tmp/orphan/sessions/%2Fhome%2Fme/one" +jq -n '{last_active_at: "2026-01-02T10:00:00Z", num_messages: 2}' >"$test_tmp/orphan/sessions/%2Fhome%2Fme/one/summary.json" +record=$(GROK_HOME="$test_tmp/orphan" collect) +[[ $(jq -c '{ready, totalSessions, status: .usageStatusText}' <<<"$record") == '{"ready":true,"totalSessions":1,"status":"Waiting for auth"}' ]] || + fail "Grok sessions with no sign-in are in use, waiting for auth" "$record" +pass "Grok sessions with no sign-in are in use, waiting for auth" + # A home chosen with GROK_HOME is the one read, as the CLI itself does. signed_in "$test_tmp/custom-grok" token-main u-main "$future" "SuperGrok" record=$(GROK_HOME="$test_tmp/custom-grok" collect) From c5bf3a3b33a33decfa2191842e8cc100e14c0120 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:13:14 -0400 Subject: [PATCH 076/546] Drop Grok's login hint while its last numbers stand With a lapsed access token Grok's CLI renews on its next start, the last numbers stand as current, but the record still carried the sign-in hint. It is empty now, as for any healthy record. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-grok | 4 +++- test/shell.d/agent-usage-grok-test.sh | 2 +- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/bin/omarchy-agent-usage-grok b/bin/omarchy-agent-usage-grok index a064784bdfa..9eac6e1414f 100755 --- a/bin/omarchy-agent-usage-grok +++ b/bin/omarchy-agent-usage-grok @@ -273,7 +273,9 @@ def collect_limits(home, key, force): # allowance, so the last numbers still stand, an elapsed window at zero. if entry.get("refresh_token") and time.time() - expires.timestamp() < REFRESH_LIFETIME_SECONDS: result["live"] = True - if not result["limits"]: + if result["limits"]: + result["authHelpText"] = "" + else: result["usageStatusText"] = "Limits paused" result["authHelpText"] = "Start Grok to see its limits; signing in again is not needed." else: diff --git a/test/shell.d/agent-usage-grok-test.sh b/test/shell.d/agent-usage-grok-test.sh index 66631428cc8..8066ac6670d 100755 --- a/test/shell.d/agent-usage-grok-test.sh +++ b/test/shell.d/agent-usage-grok-test.sh @@ -113,7 +113,7 @@ pass "a lapsed sign-in keeps the last credits and says so" jq '.[].refresh_token = "r"' "$HOME/.grok/auth.json" >"$test_tmp/auth.json" mv "$test_tmp/auth.json" "$HOME/.grok/auth.json" record=$(collect) -[[ $(jq -c '{usageStatusText, first: .limits[0].percent, stale: .limitsStale}' <<<"$record") == '{"usageStatusText":"","first":0.42,"stale":false}' ]] || +[[ $(jq -c '{usageStatusText, authHelpText, first: .limits[0].percent, stale: .limitsStale}' <<<"$record") == '{"usageStatusText":"","authHelpText":"","first":0.42,"stale":false}' ]] || fail "a lapsed access token with a refresh token keeps showing the last numbers" "$record" pass "a lapsed access token with a refresh token keeps showing the last numbers" From 4b05a5a5678b06cf196b03c8f11c2f6d8f21939b Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:02:20 -0400 Subject: [PATCH 077/546] Bring the Z.ai collector in line with the usage collector contract Zero elapsed kept windows instead of dropping them, key the limits cache by platform and key, read it whatever its time with 0 <= age guards, keep the status empty when a failed check kept limits, report a rejected key as Sign-in expired, and treat a key without a Coding Plan as healthy. Scan pi and omp profiles and PI_CODING_AGENT_DIR with rg --no-ignore, dedup forks by (id, timestamp), skip garbage timestamps and token counts, and read the key pi saves in auth.json. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-zai | 598 +++++++++++-------- test/shell.d/agent-usage-zai-limits-test.sh | 146 +++-- test/shell.d/agent-usage-zai-scanner-test.sh | 42 +- 3 files changed, 512 insertions(+), 274 deletions(-) diff --git a/bin/omarchy-agent-usage-zai b/bin/omarchy-agent-usage-zai index 989c1da33c1..bd8744d1187 100755 --- a/bin/omarchy-agent-usage-zai +++ b/bin/omarchy-agent-usage-zai @@ -7,8 +7,8 @@ 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. +generic pay-as-you-go key gets "no coding plan" and the record shows local +stats alone, which is no trouble. 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 @@ -20,12 +20,15 @@ agents panel only ever reads the JSON this prints. from __future__ import annotations import argparse +import base64 import contextlib import fcntl import hashlib import json +import math import os import stat +import subprocess import sys import tempfile import time @@ -40,7 +43,9 @@ 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." +AUTH_HELP_BAD_KEY = ( + "Z.ai rejected the saved API key. Update it in ~/.config/omarchy/agents/zai.json or ZAI_API_KEY." +) # Global (Z.ai) vs China (Zhipu / BigModel) split the same API under two hosts. PLATFORMS = { @@ -49,6 +54,10 @@ PLATFORMS = { } QUOTA_PATH = "/api/monitor/usage/quota/limit" +# The pi/omp provider names whose turns this subscription bills. Exact names: +# a custom `zai-proxy` gateway is not this subscription. +PROVIDERS = ("zai", "zhipu") + # 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"} @@ -67,37 +76,51 @@ LIMITS_ONLY_REUSE_SECONDS = 900 # 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)] +# Bounds for numbers read from session files. A timestamp outside them, or a +# token count above the cap, is corruption rather than usage. +EARLIEST_SECONDS = 946_684_800 # 2000-01-01 +LATEST_SECONDS = 4_102_444_800 # 2100-01-01 +MAX_TOKENS = 10_000_000_000 -def number(value: Any) -> int: - try: - return max(0, int(value or 0)) - except (TypeError, ValueError): +def token_count(value: Any) -> int: + # Booleans are ints to Python, and a NaN or negative count would poison + # every sum it reaches, so anything but a plain finite count reads as none. + if isinstance(value, bool) or not isinstance(value, (int, float)): + return 0 + if not math.isfinite(value) or value < 0 or value > MAX_TOKENS: return 0 + return int(value) -def local_day(value: Any) -> str: - if value is None: - return datetime.now().strftime("%Y-%m-%d") +def local_day(value: Any) -> str | None: + """The local date of an untrusted timestamp, or None when it isn't one.""" + if isinstance(value, bool): + return None if isinstance(value, (int, float)): + if not math.isfinite(value): + return None # 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) + seconds = value / 1000 if value > 10_000_000_000 else value + if not EARLIEST_SECONDS <= seconds < LATEST_SECONDS: + return None + return datetime.fromtimestamp(seconds).strftime("%Y-%m-%d") + if not isinstance(value, str) or not value.strip(): + return None 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") + parsed = datetime.fromisoformat(value.strip().replace("Z", "+00:00")) except ValueError: - return datetime.now().strftime("%Y-%m-%d") + return None + if parsed.tzinfo is not None: + parsed = parsed.astimezone() + if not 2000 <= parsed.year < 2100: + return None + return parsed.strftime("%Y-%m-%d") + + +def recent_dates() -> list[str]: + now = datetime.now() + return [(now - timedelta(days=offset)).strftime("%Y-%m-%d") for offset in range(6, -1, -1)] def empty_bucket() -> dict[str, int]: @@ -115,7 +138,7 @@ def empty_stats() -> dict[str, Any]: "todaySessions": 0, "todayTotalTokens": 0, "todayTokensByModel": {}, - "recentDays": [], + "recentDays": [{"date": day, "messageCount": 0} for day in recent_dates()], "totalPrompts": 0, "totalSessions": 0, "activeDays": 0, @@ -137,6 +160,8 @@ def base_record(**overrides: Any) -> dict[str, Any]: "usageStatusText": "", "authHelpText": "", "limits": [], + "limitsStale": False, + "limitsFetchedAt": 0, } record.update(empty_stats()) record.update(overrides) @@ -153,18 +178,14 @@ def cache_root() -> Path: return root -def read_fresh_json(path: Path, max_age_seconds: float) -> Any: - if max_age_seconds <= 0 or not path.exists(): - return None +def read_json(path: Path) -> dict[str, Any]: + # Read whatever the file's time: a clock that went back must not lose the + # last good answer. Freshness is judged by the times stored inside. 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")) + parsed = json.loads(path.read_text(encoding="utf-8")) except Exception: - return None - return None + return {} + return parsed if isinstance(parsed, dict) else {} def write_json(path: Path, payload: Any) -> None: @@ -183,6 +204,13 @@ def write_json(path: Path, payload: Any) -> None: raise +def age_below(stamp_seconds: float, limit_seconds: float) -> bool: + # A negative age is a stamp from the future: the clock moved back since it + # was written, so it proves nothing about freshness. + age = time.time() - stamp_seconds + return 0 <= age < limit_seconds + + @contextlib.contextmanager def exclusive_lock(path: Path) -> Iterator[None]: """Hold flock on `path` for the block, or run the block unlocked.""" @@ -241,11 +269,22 @@ def read_env_file(path: Path) -> dict[str, str]: 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 {} + return read_json(config_home / "omarchy" / "agents" / "zai.json") + + +def pi_agent_dir() -> Path: + configured = os.environ.get("PI_CODING_AGENT_DIR", "").strip() + return Path(os.path.expanduser(configured)) if configured else Path.home() / ".pi" / "agent" + + +def auth_json_key(path: Path, provider: str) -> str: + # pi/omp keep a key they were given with /login in auth.json. Only a literal + # key is taken: a `$VAR` or `!command` entry is resolved by pi alone. + entry = read_json(path).get(provider) + if not isinstance(entry, dict): + return "" + key = str(entry.get("key") or "").strip() + return "" if key.startswith(("!", "$")) else key def credentials() -> tuple[str, str]: @@ -264,10 +303,7 @@ def credentials() -> tuple[str, str]: platform = platform or "zhipu" if not key: - for env_path in ( - Path.home() / ".omp" / "agent" / ".env", - Path.home() / ".pi" / "agent" / ".env", - ): + for env_path in (Path.home() / ".omp" / "agent" / ".env", pi_agent_dir() / ".env"): values = read_env_file(env_path) if values.get("ZAI_API_KEY"): key, platform = values["ZAI_API_KEY"], platform or "zai" @@ -277,6 +313,15 @@ def credentials() -> tuple[str, str]: platform = platform or "zhipu" break + if not key: + for auth_path in (pi_agent_dir() / "auth.json", Path.home() / ".omp" / "agent" / "auth.json"): + if found := auth_json_key(auth_path, "zai"): + key, platform = found, platform or "zai" + break + if found := auth_json_key(auth_path, "zhipu"): + key, platform = found, platform or "zhipu" + break + if platform not in PLATFORMS: platform = "zai" return key, platform @@ -287,11 +332,13 @@ def credentials() -> tuple[str, str]: # -------------------------------------------------------------------------- def iso_from_ms(value: Any) -> str: + if isinstance(value, bool): + return "" try: ts = float(value) except (TypeError, ValueError): return "" - if ts <= 0: + if not math.isfinite(ts) or 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. @@ -306,10 +353,11 @@ def iso_from_ms(value: Any) -> str: 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 + On success returns {"ok": True, "limits": [...], "tierLabel": str}. On a + miss returns {"ok": False, "detail": str} plus one of "transport" (reached + no server, so worth an early retry), "keyRejected" (the key itself is bad) + or "noPlan" (a working key without a Coding Plan). 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 @@ -328,20 +376,20 @@ def fetch_limits(api_key: str, platform: str) -> dict[str, Any]: 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"} + return {"ok": False, "keyRejected": True, "detail": AUTH_HELP_BAD_KEY} + return {"ok": False, "detail": f"Z.ai's monitor endpoint returned HTTP {error.code}. Local stats are still shown."} + except OSError: + # URLError, timeouts and resets alike: no answer from any server. + return {"ok": False, "transport": True, "detail": "Couldn't reach Z.ai's monitor endpoint. Retrying shortly. Local stats are still shown."} + except ValueError: + return {"ok": False, "detail": "Z.ai returned an invalid monitor response. Local stats are still shown."} 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} + if "coding plan" in msg.lower(): + return {"ok": False, "noPlan": True, "detail": ""} + return {"ok": False, "detail": (msg or "Z.ai's monitor endpoint returned no limits") + ". Local stats are still shown."} limits = [] for item in data["limits"]: @@ -356,26 +404,26 @@ def fetch_limits(api_key: str, platform: str) -> dict[str, Any]: 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. + raw = item.get("percentage") + # A missing or unparseable percentage is skipped: a false 0% meter reads + # as a fresh allowance, which is worse than showing one fewer window. + if isinstance(raw, bool) or not isinstance(raw, (int, float)) or not math.isfinite(raw): continue limits.append({ "label": label, - "percent": percent, + "percent": max(0.0, min(1.0, float(raw) / 100.0)), "resetsAt": iso_from_ms(item.get("nextResetTime")), }) + if not limits: + return {"ok": False, "detail": "Z.ai's monitor endpoint returned no limits. Local stats are still shown."} tier = str(data.get("level") or "").strip()[:64] - return {"limits": limits, "tierLabel": tier[:1].upper() + tier[1:].lower() if tier else ""} + return {"ok": True, "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. +# A kept window whose reset has passed comes back at zero: the allowance it +# measured has been renewed, so it is untouched rather than unknown. A window +# with no reset time, or one that will not parse, is kept as it is. def limit_window_open(entry: dict[str, Any], now_utc: datetime) -> bool: raw = str(entry.get("resetsAt") or "") if raw == "": @@ -394,7 +442,26 @@ def usable_cached_limits(cached: dict[str, Any]) -> list[dict[str, Any]]: 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)] + return [ + entry if limit_window_open(entry, now_utc) else dict(entry, percent=0.0, resetsAt="") + for entry in entries if isinstance(entry, dict) + ] + + +def limits_cache_paths(api_key: str, platform: str) -> tuple[Path, Path]: + # Keyed by who the numbers belong to: another key, or the same key on the + # other host, never shows this one's allowance. + digest = hashlib.sha256(f"{platform}\0{api_key}".encode("utf-8")).hexdigest()[:16] + root = cache_root() + return root / f"zai-limits-{digest}.json", root / f"zai-limits-{digest}.lock" + + +def cached_answer(cache_file: Path) -> tuple[list[dict[str, Any]], str, int]: + cached = read_json(cache_file) + fetched_ms = cached.get("fetchedAtMs") + if isinstance(fetched_ms, bool) or not isinstance(fetched_ms, (int, float)) or not math.isfinite(fetched_ms): + fetched_ms = 0 + return usable_cached_limits(cached), str(cached.get("tierLabel") or ""), max(0, int(fetched_ms)) def collect_limits(api_key: str, platform: str, force: bool) -> dict[str, Any]: @@ -402,41 +469,33 @@ 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. `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. + record when a later probe fails. `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 tell kept limits from fresh ones. """ result = {"limits": [], "tierLabel": "", "usageStatusText": "", "authHelpText": "", "retryAdvised": False, "live": False, "fetchedAtMs": 0} + cache_file, lock_file = limits_cache_paths(api_key, platform) - 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 - result.update(live=True, fetchedAtMs=round(fetched_at * 1000)) + # --force is a person asking for fresh numbers, so it skips the reuse + # window; the interval absorbs repeated panel opens, not a refresh press. + fallback, cached_tier, fetched_ms = cached_answer(cache_file) + if fallback and not force and age_below(fetched_ms / 1000, PROBE_MIN_INTERVAL_SECONDS): + result.update(limits=fallback, tierLabel=cached_tier, live=True, fetchedAtMs=fetched_ms) return result - with exclusive_lock(scan_cache_paths()[2]): + with exclusive_lock(lock_file): 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 "") - result.update(live=True, fetchedAtMs=round(fetched_at * 1000)) + # Another run may have probed while this one waited for the lock. + fallback, cached_tier, fetched_ms = cached_answer(cache_file) + if fallback and age_below(fetched_ms / 1000, PROBE_MIN_INTERVAL_SECONDS): + result.update(limits=fallback, tierLabel=cached_tier, live=True, fetchedAtMs=fetched_ms) 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)) + if probe["ok"]: + result.update(limits=probe["limits"], tierLabel=probe["tierLabel"], live=True, fetchedAtMs=round(time.time() * 1000)) try: - write_json(probe_cache, { + write_json(cache_file, { "fetchedAtMs": result["fetchedAtMs"], "limits": probe["limits"], "tierLabel": result["tierLabel"], @@ -445,19 +504,24 @@ def collect_limits(api_key: str, platform: str, force: bool) -> dict[str, Any]: 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. + # A working key without a Coding Plan is a pay-as-you-go account: Z.ai has + # no limits to show for it, which is no trouble, and limits kept from a plan + # that has since lapsed no longer describe it. + if probe.get("noPlan"): + result["live"] = True + return result + + # 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["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. + result.update(limits=fallback, tierLabel=cached_tier, fetchedAtMs=fetched_ms if fallback else 0) + result["authHelpText"] = probe["detail"] if probe.get("keyRejected"): - result["authHelpText"] = AUTH_HELP_BAD_KEY + result["usageStatusText"] = "Sign-in expired" + elif not fallback: + result["usageStatusText"] = "Z.ai limits unavailable" return result @@ -465,119 +529,186 @@ def collect_limits(api_key: str, platform: str, force: bool) -> dict[str, Any]: # 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): +class Stats: + """Token and activity totals for every counted turn.""" + + def __init__(self) -> None: + self.today = datetime.now().strftime("%Y-%m-%d") + self.recent_dates = recent_dates() + self.recent = {day: 0 for day in self.recent_dates} + self.today_prompts = 0 + self.today_tokens = 0 + self.today_sessions: set[str] = set() + self.today_by_model: dict[str, int] = {} + self.total_prompts = 0 + self.sessions: set[str] = set() + self.active_dates: set[str] = set() + self.model_usage: dict[str, dict[str, int]] = {} + + def add(self, session: str, day: str, model: str, inp: int, out: int, cread: int, cwrite: int) -> None: 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()) + self.total_prompts += 1 + self.sessions.add(session) + self.active_dates.add(day) + bucket = self.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", - ] + if day in self.recent: + self.recent[day] += total + if day == self.today: + self.today_prompts += 1 + self.today_sessions.add(session) + self.today_tokens += total + self.today_by_model[model] = self.today_by_model.get(model, 0) + total + + def record(self) -> dict[str, Any]: + return { + "hasLocalStats": bool(self.model_usage), + "todayPrompts": self.today_prompts, + "todaySessions": len(self.today_sessions), + "todayTotalTokens": self.today_tokens, + "todayTokensByModel": self.today_by_model, + "recentDays": [{"date": day, "messageCount": self.recent[day]} for day in self.recent_dates], + "totalPrompts": self.total_prompts, + "totalSessions": len(self.sessions), + "activeDays": len(self.active_dates), + "activeDates": sorted(self.active_dates), + "modelUsage": self.model_usage, + } + + +# `omp --profile=` (and pi's equivalent) relocates the whole agent tree +# under /profiles//, and PI_CODING_AGENT_DIR moves pi's, so a +# subscription driven through either would otherwise go uncounted. +def pi_session_roots() -> list[Path]: + roots = [pi_agent_dir() / "sessions", Path.home() / ".omp" / "agent" / "sessions"] + for base in (Path.home() / ".pi", Path.home() / ".omp"): + try: + # Sorted so the scan order does not depend on directory order. + roots.extend(sorted(child / "agent" / "sessions" for child in (base / "profiles").iterdir() if child.is_dir())) + except OSError: + # No profiles directory, or it is unreadable: the default roots stand. + pass + unique: list[Path] = [] for root in roots: - if not root.exists(): + if root not in unique: + unique.append(root) + return unique + + +def matching_lines(root: Path, pattern: str) -> tuple[list[tuple[str, str]], bool]: + """(path, line) for each session line matching `pattern`, and whether every file was read.""" + try: + proc = subprocess.run( + # --no-ignore: these are data files, not a source tree. When $HOME is + # itself a git checkout (a common dotfiles setup), its ignore rules + # would otherwise hide every session file and count zero usage. + ["rg", "--json", "--no-ignore", "--hidden", "--glob", "*.jsonl", "-e", pattern, str(root)], + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + text=True, + errors="replace", + check=False, + ) + except OSError: + return [], False + lines = [] + for raw in proc.stdout.splitlines(): + try: + event = json.loads(raw) + if event.get("type") != "match": + continue + data = event["data"] + text = data["lines"].get("text") + if text is None: + # rg hands back lines that are not valid UTF-8 as base64. + text = base64.b64decode(data["lines"]["bytes"]).decode("utf-8", errors="replace") + path = data["path"].get("text") or "" + lines.append((path, text)) + except Exception: + continue + # 1 is "no match"; 2 means some file could not be read. + return lines, proc.returncode in (0, 1) + + +def scan_sessions() -> tuple[dict[str, Any], bool]: + stats = Stats() + complete = True + # A forked session copies its parent's messages, ids and timestamps alike, + # so (id, timestamp) counts each once across files. A line with an id but + # no timestamp dedups within its file only; a line with no id always counts. + seen: set[tuple[str, ...]] = set() + pattern = r'"provider"\s*:\s*"(' + "|".join(PROVIDERS) + r')"' + + for root in pi_session_roots(): + if not root.is_dir(): continue - for path in root.rglob("*.jsonl"): + lines, read_all = matching_lines(root, pattern) + complete = complete and read_all + for path, line in lines: try: - handle = path.open(errors="replace") - except OSError: + entry = json.loads(line) + except ValueError: + 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 + if str(message.get("provider") or "").lower() not in PROVIDERS: + continue + usage = message.get("usage") + if not isinstance(usage, dict): + continue + inp = token_count(usage.get("input")) + # pi reports reasoning tokens apart from output; both are generated. + out = token_count(usage.get("output")) + token_count(usage.get("reasoning")) + cread = token_count(usage.get("cacheRead")) + cwrite = token_count(usage.get("cacheWrite")) + if not (inp or out or cread or cwrite): + inp = token_count(usage.get("totalTokens")) + if not (inp or out or cread or cwrite): + continue + day = local_day(entry.get("timestamp")) or local_day(message.get("timestamp")) + if day is None: 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, - } + message_id = entry.get("id") + timestamp = entry.get("timestamp") + if isinstance(message_id, str) and message_id: + key = ("entry", message_id, str(timestamp)) if timestamp else ("path", path, message_id) + if key in seen: + continue + seen.add(key) + stats.add(path, day, str(message.get("model") or "glm"), inp, out, cread, cwrite) + return stats.record(), complete -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] + +def scan_cache_paths() -> tuple[Path, Path]: + # The digest covers every data path the scan reads. + identity = "\0".join(str(root) for root in pi_session_roots()) + digest = hashlib.sha1(identity.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" + return root / f"zai-scan-{digest}.json", root / f"zai-scan-{digest}.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: + cached = read_json(cache_file) + if cached.get("schemaVersion") != 2: 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: + if cached.get("scanDate") != datetime.now().strftime("%Y-%m-%d"): + return None + scanned_at = cached.get("scannedAt") + if isinstance(scanned_at, bool) or not isinstance(scanned_at, (int, float)): + return None + if not age_below(scanned_at, max_age_seconds): return None stats = cached.get("stats") if not isinstance(stats, dict): @@ -587,13 +718,6 @@ def read_cached_stats(cache_file: Path, max_age_seconds: float) -> dict[str, Any 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. @@ -601,26 +725,23 @@ def cached_local_stats(max_age: float) -> dict[str, Any]: 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): + cache_file, lock_file = scan_cache_paths() 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 + with exclusive_lock(lock_file): + cached = read_cached_stats(cache_file, max_age) + if cached is not None: + return cached + stats, complete = scan_sessions() + # A scan that missed a file would stand in for a real one until the + # next rescan; only a complete one is kept. + if complete: + write_json(cache_file, {"schemaVersion": 2, "scanDate": datetime.now().strftime("%Y-%m-%d"), "scannedAt": time.time(), "stats": stats}) + return stats + except Exception as exc: + print(f"omarchy-agent-usage-zai: cache unavailable ({exc}); scanning directly", file=sys.stderr) + return scan_sessions()[0] # -------------------------------------------------------------------------- @@ -634,23 +755,24 @@ def scan(force: bool, limits_only: bool) -> dict[str, Any]: api_key, platform = credentials() if not api_key: - record["authHelpText"] = AUTH_HELP_NO_KEY - record["ready"] = record["hasLocalStats"] + # Z.ai in use here without a key to ask for its limits; a machine that has + # never used it is simply skipped. + if record["hasLocalStats"]: + record.update(ready=True, usageStatusText="Waiting for auth", authHelpText=AUTH_HELP_NO_KEY) return record 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"] - if limits["authHelpText"]: - record["authHelpText"] = limits["authHelpText"] + record.update( + ready=True, + limits=limits["limits"], + limitsStale=bool(limits["limits"]) and not limits["live"], + limitsFetchedAt=limits["fetchedAtMs"], + tierLabel=limits["tierLabel"], + usageStatusText=limits["usageStatusText"], + authHelpText=limits["authHelpText"], + ) if limits["retryAdvised"]: record["retryAdvised"] = True - - record["ready"] = bool(record["limits"] or record["hasLocalStats"]) return record @@ -666,7 +788,7 @@ def main() -> int: 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") + record = base_record(usageStatusText="Z.ai limits 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 diff --git a/test/shell.d/agent-usage-zai-limits-test.sh b/test/shell.d/agent-usage-zai-limits-test.sh index 36915183fac..43ce6cd6a51 100755 --- a/test/shell.d/agent-usage-zai-limits-test.sh +++ b/test/shell.d/agent-usage-zai-limits-test.sh @@ -9,7 +9,8 @@ require_command python3 # 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 +TEST_HOME=$(mktemp -d) +trap 'rm -rf "$CACHE_HOME" "$TEST_HOME"' EXIT read_limits() { COLLECTOR="$ROOT/bin/omarchy-agent-usage-zai" PAYLOAD="$1" python3 - <<'PY' @@ -25,9 +26,10 @@ print(json.dumps(collector.fetch_limits("token", "zai"))) PY } -# force, cached-json, mode (success|transport|http401), payload -> {result, probes, cached} +# force, cached-json, mode (success|transport|http401|http500), payload, [key] +# -> {result, probes, cached} drive() { - COLLECTOR="$ROOT/bin/omarchy-agent-usage-zai" FORCE="$1" CACHED="$2" MODE="$3" PAYLOAD="$4" \ + COLLECTOR="$ROOT/bin/omarchy-agent-usage-zai" FORCE="$1" CACHED="$2" MODE="$3" PAYLOAD="$4" KEY="${5:-token}" \ XDG_CACHE_HOME="$CACHE_HOME" python3 - <<'PY' import importlib.machinery, importlib.util, io, json, os, urllib.error @@ -36,7 +38,8 @@ 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" +# Every case seeds the cache of the "token" key; another key must not see it. +cache = collector.limits_cache_paths("token", "zai")[0] cached = os.environ["CACHED"] if cached: cache.write_text(cached, encoding="utf-8") @@ -52,10 +55,12 @@ def urlopen(request, timeout=None): raise urllib.error.URLError("no route to host") if mode == "http401": raise urllib.error.HTTPError("u", 401, "no", {}, None) + if mode == "http500": + raise urllib.error.HTTPError("u", 500, "no", {}, None) return io.BytesIO(os.environ["PAYLOAD"].encode()) collector.urllib.request.urlopen = urlopen -result = collector.collect_limits("token", "zai", os.environ["FORCE"] == "true") +result = collector.collect_limits(os.environ["KEY"], "zai", os.environ["FORCE"] == "true") out = {"result": result, "probes": len(probes)} if cache.exists(): out["cached"] = json.loads(cache.read_text(encoding="utf-8")) @@ -63,6 +68,8 @@ print(json.dumps(out)) PY } +now_ms=$(python3 -c 'import time; print(round(time.time() * 1000))') +future_ms=$(python3 -c 'import time; print(round((time.time() + 86400) * 1000))') 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())") @@ -70,18 +77,20 @@ open_at=$(python3 -c 'import datetime as dt; print((dt.datetime.now(dt.timezone. 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. +# tagged with an unmapped (unit, number), one whose percentage will not parse, +# and one whose percentage is a boolean. 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 } + { type: "CREDIT_LIMIT", unit: 3, number: 5, percentage: null, nextResetTime: $sm }, + { type: "CREDIT_LIMIT", unit: 3, number: 5, percentage: true, nextResetTime: $sm } ] } }') -limits=$(read_limits "$payload") +limits=$(read_limits "$payload" 2>/dev/null) [[ $(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" @@ -91,39 +100,55 @@ limits=$(read_limits "$payload") 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" '{ +kept=$(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" + +# A key with no plan is a pay-as-you-go account: no trouble, no blame on the +# key, and no limits kept from a plan that lapsed. +noplan=$(drive false "$kept" success '{"success":false,"msg":"This API key has no coding plan"}') +[[ $(jq -c '.result | {usageStatusText, authHelpText, retryAdvised, limits, live}' <<<"$noplan") == '{"usageStatusText":"","authHelpText":"","retryAdvised":false,"limits":[],"live":true}' ]] || + fail "Z.ai collector treats a key without a coding plan as healthy with no limits" "$noplan" +pass "Z.ai collector treats a no-plan key as healthy, without limits" + +# A rejected key is a sign-in to redo, and the kept limits stay, stale. +rejected=$(drive false "$kept" http401 "") +[[ $(jq -r '.result.usageStatusText' <<<"$rejected") == "Sign-in expired" ]] || + fail "Z.ai collector reports a rejected key as an expired sign-in" "$rejected" +[[ $(jq -r '.result.authHelpText' <<<"$rejected") == *"rejected the saved API key"* ]] || + fail "Z.ai collector blames the key only when the key is rejected" "$rejected" +[[ $(jq -c '{live: .result.live, n: (.result.limits | length), retry: .result.retryAdvised}' <<<"$rejected") == '{"live":false,"n":2,"retry":false}' ]] || + fail "Z.ai collector keeps the last limits, stale, past a rejected key" "$rejected" +pass "Z.ai collector reports a rejected key as Sign-in expired" + +# A transport failure keeps the last known figures with no status of its own, +# brings the window that already reset back at zero, and asks for a retry. +unreachable=$(drive false "$kept" transport "") +[[ $(jq --arg open "$open_at" '[.result.limits[] | {label, percent, resetsAt}] == [{label: "Session (5-hour)", percent: 0, resetsAt: ""}, {label: "Weekly (7-day)", percent: 0.11, resetsAt: $open}]' <<<"$unreachable") == "true" ]] || + fail "Z.ai collector zeroes a kept window that has reset rather than dropping it" "$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" +[[ $(jq -c '{live: .result.live, fetchedAtMs: .result.fetchedAtMs, status: .result.usageStatusText}' <<<"$unreachable") == '{"live":false,"fetchedAtMs":1,"status":""}' ]] || + fail "Z.ai collector marks kept limits stale with when they were fetched, without a status" "$unreachable" pass "Z.ai collector serves the last known limits and advises a retry when offline" +# With nothing kept, a failed check says the limits are unavailable; only a +# transport failure asks for an early retry. +offline=$(drive false "" transport "") +[[ $(jq -c '.result | {usageStatusText, retryAdvised, n: (.limits | length)}' <<<"$offline") == '{"usageStatusText":"Z.ai limits unavailable","retryAdvised":true,"n":0}' ]] || + fail "Z.ai collector reports unavailable limits when offline with nothing kept" "$offline" +[[ $(jq -r '.result.authHelpText' <<<"$offline") == *"reach"* ]] || + fail "Z.ai collector explains an unreachable endpoint" "$offline" +server=$(drive false "" http500 "") +[[ $(jq -c '.result | {usageStatusText, retryAdvised}' <<<"$server") == '{"usageStatusText":"Z.ai limits unavailable","retryAdvised":false}' ]] || + fail "Z.ai collector does not advise a retry for a server's error status" "$server" +pass "Z.ai collector reports unavailable limits and advises a retry only for transport failures" + # 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))')" '{ +fresh=$(jq -nc --arg open "$open_at" --argjson now "$now_ms" '{ 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: [ @@ -144,12 +169,63 @@ forced=$(drive true "$fresh" success "$success_payload") fail "Z.ai collector marks a fresh probe live" "$forced" pass "Z.ai collector re-probes on --force and caches the result" +# A cache stamped in the future (the clock went back) is no proof of a recent +# probe, but it is still the last good answer to fall back on. +ahead=$(jq -nc --arg open "$open_at" --argjson ms "$future_ms" '{ + fetchedAtMs: $ms, tierLabel: "Pro", limits: [{ label: "Weekly (7-day)", percent: 0.11, resetsAt: $open }] +}') +skewed=$(drive false "$ahead" transport "") +[[ $(jq -r '.probes' <<<"$skewed") == "1" && $(jq -c '[.result.limits[].percent]' <<<"$skewed") == "[0.11]" && $(jq -r '.result.live' <<<"$skewed") == "false" ]] || + fail "Z.ai collector re-probes past a future-stamped cache and still falls back on it" "$skewed" +pass "Z.ai collector neither trusts nor loses a future-stamped cache" + +# Another key never shows this one's numbers. +other=$(drive false "$fresh" transport "" other-key) +[[ $(jq -c '.result | {n: (.limits | length), usageStatusText}' <<<"$other") == '{"n":0,"usageStatusText":"Z.ai limits unavailable"}' ]] || + fail "Z.ai collector keys its limits cache by credential" "$other" +pass "Z.ai collector keeps each key's limits apart" + # A predictable lock path must not follow and truncate a planted symlink. +lock=$(XDG_CACHE_HOME="$CACHE_HOME" COLLECTOR="$ROOT/bin/omarchy-agent-usage-zai" python3 - <<'PY' +import importlib.machinery, importlib.util, 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) +print(collector.limits_cache_paths("token", "zai")[1]) +PY +) 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" +rm -f "$lock" +ln -s "$sentinel" "$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" + +# A machine that has never used Z.ai prints the full, zeroed record the panel +# skips, with no status to alarm anyone. +unused=$(env -u ZAI_API_KEY -u ZHIPU_API_KEY -u ZHIPUAI_API_KEY -u PI_CODING_AGENT_DIR \ + HOME="$TEST_HOME" XDG_CONFIG_HOME="$TEST_HOME/.config" XDG_CACHE_HOME="$TEST_HOME/.cache" \ + "$ROOT/bin/omarchy-agent-usage-zai") +[[ $(jq -c '{ready, usageStatusText, authHelpText, limits, limitsStale, limitsFetchedAt, totalPrompts, days: (.recentDays | length), modelUsage}' <<<"$unused") == '{"ready":false,"usageStatusText":"","authHelpText":"","limits":[],"limitsStale":false,"limitsFetchedAt":0,"totalPrompts":0,"days":7,"modelUsage":{}}' ]] || + fail "Z.ai collector prints an empty, unready record on an unused machine" "$unused" +pass "Z.ai collector prints an empty, unready record on an unused machine" + +# A key pi saved with /login counts as a credential. +mkdir -p "$TEST_HOME/.pi/agent" +printf '{"zai":{"type":"api_key","key":"pi-saved"}}' >"$TEST_HOME/.pi/agent/auth.json" +picked=$(env -u ZAI_API_KEY -u ZHIPU_API_KEY -u ZHIPUAI_API_KEY -u PI_CODING_AGENT_DIR HOME="$TEST_HOME" XDG_CONFIG_HOME="$TEST_HOME/.config" \ + COLLECTOR="$ROOT/bin/omarchy-agent-usage-zai" 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) +print(json.dumps(collector.credentials())) +PY +) +[[ $picked == '["pi-saved", "zai"]' ]] || + fail "Z.ai collector reads the key pi saved in auth.json" "$picked" +pass "Z.ai collector reads the key pi saved in auth.json" diff --git a/test/shell.d/agent-usage-zai-scanner-test.sh b/test/shell.d/agent-usage-zai-scanner-test.sh index a83cab33fab..fc2fd787a5f 100755 --- a/test/shell.d/agent-usage-zai-scanner-test.sh +++ b/test/shell.d/agent-usage-zai-scanner-test.sh @@ -44,6 +44,11 @@ result=$(run) fail "Z.ai collector identifies itself and is ready on local stats alone (no key)" "$result" pass "Z.ai collector reports itself, empty limits, ready on local stats" +# In use here, but no key to ask for its limits. +[[ $(jq -r '.usageStatusText' <<<"$result") == "Waiting for auth" && $(jq -r '.authHelpText' <<<"$result") == *"ZAI_API_KEY"* ]] || + fail "Z.ai collector waits for auth when used without a key" "$result" +pass "Z.ai collector waits for auth when used without a key" + # Reasoning tokens fold into output; cache tokens split out; the repeated id is # counted once, so the glm-5.3 bucket is exactly message a. [[ $(jq -Sc '.modelUsage["glm-5.3"]' <<<"$result") == '{"cacheCreationInputTokens":2,"cacheReadInputTokens":3,"inputTokens":10,"outputTokens":10}' ]] || @@ -65,7 +70,7 @@ pass "Z.ai collector filters providers/roles and survives malformed lines" # The scan is cached so the panel's per-open --limits-only does not re-walk # every session file; --force bypasses it. cache_file=$(ls "$TEST_HOME/.cache/omarchy/agent-usage/"zai-scan-*.json 2>/dev/null | head -n 1) -[[ -n $cache_file && $(jq -r '.schemaVersion' "$cache_file") == "1" && $(jq -r '.stats.todayTotalTokens' "$cache_file") == "50" ]] || +[[ -n $cache_file && $(jq -r '.schemaVersion' "$cache_file") == "2" && $(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" @@ -88,3 +93,38 @@ result=$(run --limits-only) [[ $(jq -r '.todayTotalTokens' <<<"$result") == "150" ]] || fail "Z.ai collector --limits-only sees the cache refreshed by --force" "$result" pass "Z.ai collector --limits-only sees a refreshed cache after --force" + +# Sessions outside the default root: an omp profile, a pi profile, and a pi +# agent dir moved by PI_CODING_AGENT_DIR, all under a $HOME that is a git +# checkout ignoring everything. A forked session repeats its parent's message +# (same id and timestamp) and counts once; lines without an id always count; +# rows whose timestamp or tokens are garbage drop without taking the rest. +ROOTS_HOME=$(mktemp -d) +trap 'rm -rf "$TEST_HOME" "$ROOTS_HOME"' EXIT +mkdir -p "$ROOTS_HOME/.git" +printf '*\n' >"$ROOTS_HOME/.gitignore" +mkdir -p "$ROOTS_HOME/.omp/profiles/work/agent/sessions/p" "$ROOTS_HOME/.pi/profiles/home/agent/sessions/p" "$ROOTS_HOME/moved/sessions/p" +line() { + printf '{"type":"message",%s"timestamp":%s,"message":{"role":"assistant","provider":"zai","model":"glm-5","usage":{"input":%s}}}\n' "$1" "$2" "$3" +} +{ + line '"id":"parent",' "\"$timestamp\"" 100 + line '' "\"$timestamp\"" 1 + line '' "\"$timestamp\"" 1 +} >"$ROOTS_HOME/.omp/profiles/work/agent/sessions/p/a.jsonl" +{ + line '"id":"parent",' "\"$timestamp\"" 100 + line '"id":"child",' "\"$timestamp\"" 10 + line '"id":"bad-ts",' 0 5000 + line '"id":"bad-ts2",' '"not a date"' 5000 + line '"id":"bad-tokens",' "\"$timestamp\"" true + line '"id":"huge-tokens",' "\"$timestamp\"" 1e300 +} >"$ROOTS_HOME/.pi/profiles/home/agent/sessions/p/fork.jsonl" +line '"id":"moved",' "\"$timestamp\"" 1000 >"$ROOTS_HOME/moved/sessions/p/m.jsonl" + +roots=$(env -u ZAI_API_KEY -u ZHIPU_API_KEY -u ZHIPUAI_API_KEY \ + HOME="$ROOTS_HOME" PI_CODING_AGENT_DIR="$ROOTS_HOME/moved" XDG_CONFIG_HOME="$ROOTS_HOME/.config" XDG_CACHE_HOME="$ROOTS_HOME/.cache" \ + "$ROOT/bin/omarchy-agent-usage-zai" --force) +[[ $(jq -r '.todayTotalTokens' <<<"$roots") == "1112" && $(jq -r '.totalPrompts' <<<"$roots") == "5" ]] || + fail "Z.ai collector scans profiles and PI_CODING_AGENT_DIR past ignore rules, dedups forks, counts id-less lines, skips garbage rows" "$roots" +pass "Z.ai collector scans every pi/omp root, dedups forks, and skips garbage rows" From 1335c002558575b567dba1b279ad15e6a9966d4c Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:05:55 -0400 Subject: [PATCH 078/546] Bring the Ollama Cloud collector in line with the usage collector contract Key the limits and plan cache by credential, read it whatever its time with 0 <= age guards, report a rejected key as Sign-in expired, keep the cached plan when only the plan lookup fails, and print the full zeroed stats block with no status on an unused machine. Scan pi and omp profiles and PI_CODING_AGENT_DIR with rg --no-ignore, dedup forks by (id, timestamp), skip garbage timestamps and token counts, cache only complete scans under flock, and name the scan reuse windows. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-ollama | 534 ++++++++++++------ .../agent-usage-ollama-scanner-test.sh | 128 ++++- 2 files changed, 463 insertions(+), 199 deletions(-) diff --git a/bin/omarchy-agent-usage-ollama b/bin/omarchy-agent-usage-ollama index f4db9d81939..d7764f009fb 100755 --- a/bin/omarchy-agent-usage-ollama +++ b/bin/omarchy-agent-usage-ollama @@ -14,10 +14,14 @@ JSON this prints; it never talks to disk formats or endpoints. from __future__ import annotations import argparse +import base64 import datetime as dt import fcntl +import hashlib import json +import math import os +import subprocess import sys import tempfile import time @@ -28,31 +32,62 @@ from typing import Any AGENT_ID = "ollama" AGENT_NAME = "Ollama Cloud" +PROVIDER = "ollama-cloud" USAGE_ENDPOINT = "https://ollama.com/api/usage" ME_ENDPOINT = "https://ollama.com/api/me" +AUTH_HELP_NO_KEY = "Set OLLAMA_API_KEY, or sign in to pi with an Ollama Cloud key, to show the account limits." +AUTH_HELP_BAD_KEY = ( + "Ollama Cloud rejected the API key. Update OLLAMA_API_KEY or ~/.config/omarchy/agents/ollama.json," + " or sign in to pi again. Local stats are still shown." +) +# 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). --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 # 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. +# Bounds for numbers read from session files. A timestamp outside them, or a +# token count above the cap, is corruption rather than usage. +EARLIEST_SECONDS = 946_684_800 # 2000-01-01 +LATEST_SECONDS = 4_102_444_800 # 2100-01-01 +MAX_TOKENS = 10_000_000_000 + def config_dir() -> Path: - return Path(os.environ.get("XDG_CONFIG_HOME", Path.home() / ".config")) / "omarchy" / "agents" + return Path(os.environ.get("XDG_CONFIG_HOME") or Path.home() / ".config") / "omarchy" / "agents" def cache_root() -> Path: - root = Path(os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache")) / "omarchy" / "agent-usage" + 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_config() -> dict[str, Any]: +def read_json(path: Path) -> dict[str, Any]: + # Read whatever the file's time: a clock that went back must not lose the + # last good answer. Freshness is judged by the times stored inside. try: - data = json.loads((config_dir() / "ollama.json").read_text(encoding="utf-8")) + data = json.loads(path.read_text(encoding="utf-8")) except Exception: return {} return data if isinstance(data, dict) else {} +def read_config() -> dict[str, Any]: + return read_json(config_dir() / "ollama.json") + + +def pi_agent_dir() -> Path: + configured = os.environ.get("PI_CODING_AGENT_DIR", "").strip() + return Path(os.path.expanduser(configured)) if configured else Path.home() / ".pi" / "agent" + + def api_key() -> str: env = os.environ.get("OLLAMA_API_KEY", "").strip() if env: @@ -61,15 +96,13 @@ def api_key() -> str: 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 + # proves the account, and the usage endpoint accepts the same key. Only a + # literal key is taken: a `$VAR` or `!command` entry is resolved by pi alone. + for auth_path in (pi_agent_dir() / "auth.json", Path.home() / ".omp" / "agent" / "auth.json"): + entry = read_json(auth_path).get(PROVIDER) + key = str(entry.get("key") or "").strip() if isinstance(entry, dict) else "" + if key and not key.startswith(("!", "$")): + return key return "" @@ -86,36 +119,53 @@ 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() - +def local_date_from_timestamp(value: Any) -> str | None: + """The local date of an untrusted timestamp, or None when it isn't one.""" + if isinstance(value, bool): + return None 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() + if not math.isfinite(value): + return None + # pi/omp message timestamps are milliseconds. + seconds = value / 1000 if value > 10_000_000_000 else value + if not EARLIEST_SECONDS <= seconds < LATEST_SECONDS: + return None + return date_string(dt.datetime.fromtimestamp(seconds).date()) + if not isinstance(value, str) or not value.strip(): + return None + try: + parsed = dt.datetime.fromisoformat(value.strip().replace("Z", "+00:00")) + except ValueError: + return None + if parsed.tzinfo is not None: + parsed = parsed.astimezone() + if not 2000 <= parsed.year < 2100: + return None + return date_string(parsed.date()) - 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 token_count(value: Any) -> int: + # Booleans are ints to Python, and a NaN or negative count would poison + # every sum it reaches, so anything but a plain finite count reads as none. + if isinstance(value, bool) or not isinstance(value, (int, float)): + return 0 + if not math.isfinite(value) or value < 0 or value > MAX_TOKENS: + return 0 + return int(value) -def number(value: Any) -> int: - try: - n = float(value or 0) - return round(n) if n == n else 0 - except Exception: +def stamp(value: Any) -> int: + # An epoch-ms stamp this collector wrote, read back without trusting it. + if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value): return 0 + return max(0, int(value)) + + +def age_below(stamp_seconds: float, limit_seconds: float) -> bool: + # A negative age is a stamp from the future: the clock moved back since it + # was written, so it proves nothing about freshness. + age = time.time() - stamp_seconds + return 0 <= age < limit_seconds def empty_bucket() -> dict[str, int]: @@ -127,15 +177,19 @@ def empty_bucket() -> dict[str, int]: } -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 empty_stats() -> dict[str, Any]: + return { + "todayPrompts": 0, + "todaySessions": 0, + "todayTotalTokens": 0, + "todayTokensByModel": {}, + "recentDays": [{"date": day, "messageCount": 0} for day in recent_date_strings()], + "modelUsage": {}, + "totalPrompts": 0, + "totalSessions": 0, + "activeDays": 0, + "activeDates": [], + } def write_json(path: Path, payload: dict[str, Any]) -> None: @@ -163,123 +217,211 @@ def write_json(path: Path, payload: dict[str, Any]) -> None: # 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" - # 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 and cached.get("scanDate") == today and cached.get("modelPrefix") == model_prefix: - return cached.get("stats") +# `omp --profile=` (and pi's equivalent) relocates the whole agent tree +# under /profiles//, and PI_CODING_AGENT_DIR moves pi's, so a +# subscription driven through either would otherwise go uncounted. +def pi_session_roots() -> list[Path]: + roots = [pi_agent_dir() / "sessions", Path.home() / ".omp" / "agent" / "sessions"] + for base in (Path.home() / ".pi", Path.home() / ".omp"): + try: + # Sorted so the scan order does not depend on directory order. + roots.extend(sorted(child / "agent" / "sessions" for child in (base / "profiles").iterdir() if child.is_dir())) + except OSError: + # No profiles directory, or it is unreadable: the default roots stand. + pass + unique: list[Path] = [] + for root in roots: + if root not in unique: + unique.append(root) + return unique + +def matching_lines(root: Path, pattern: str) -> tuple[list[tuple[str, str]], bool]: + """(path, line) for each session line matching `pattern`, and whether every file was read.""" + try: + proc = subprocess.run( + # --no-ignore: these are data files, not a source tree. When $HOME is + # itself a git checkout (a common dotfiles setup), its ignore rules + # would otherwise hide every session file and count zero usage. + ["rg", "--json", "--no-ignore", "--hidden", "--glob", "*.jsonl", "-e", pattern, str(root)], + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + text=True, + errors="replace", + check=False, + ) + except OSError: + return [], False + lines = [] + for raw in proc.stdout.splitlines(): + try: + event = json.loads(raw) + if event.get("type") != "match": + continue + data = event["data"] + text = data["lines"].get("text") + if text is None: + # rg hands back lines that are not valid UTF-8 as base64. + text = base64.b64decode(data["lines"]["bytes"]).decode("utf-8", errors="replace") + lines.append((data["path"].get("text") or "", text)) + except Exception: + continue + # 1 is "no match"; 2 means some file could not be read. + return lines, proc.returncode in (0, 1) + + +def scan_sessions(model_prefix: str) -> tuple[dict[str, Any], bool]: recent_dates = recent_date_strings() recent = {day: {"date": day, "messageCount": 0} for day in recent_dates} + today = local_date_string() 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() + # A forked session copies its parent's messages, ids and timestamps alike, + # so (id, timestamp) counts each once across files. A line with an id but + # no timestamp dedups within its file only; a line with no id always counts. + seen: set[tuple[str, ...]] = set() prompts = 0 today_prompt_count = 0 today_token_total = 0 + complete = True - for root in roots: - files = root.rglob("*.jsonl") if root.is_dir() else [] - for path in files: + for root in pi_session_roots(): + if not root.is_dir(): + continue + lines, read_all = matching_lines(root, r'"provider"\s*:\s*"ollama-cloud"') + complete = complete and read_all + for path, line in lines: 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: + entry = json.loads(line) + except ValueError: + continue + if not isinstance(entry, dict) or entry.get("type") != "message": + continue + message = entry.get("message") + if not isinstance(message, dict) or message.get("role") != "assistant": + continue + if str(message.get("provider") or "") != PROVIDER: + continue + usage = message.get("usage") + if not isinstance(usage, dict): + continue + input_tokens = token_count(usage.get("input")) + output_tokens = token_count(usage.get("output")) + cache_read = token_count(usage.get("cacheRead")) + cache_write = token_count(usage.get("cacheWrite")) + total = input_tokens + output_tokens + cache_read + cache_write + if total <= 0: + total = token_count(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.lower().startswith(model_prefix): + continue + day = local_date_from_timestamp(entry.get("timestamp")) or local_date_from_timestamp(message.get("timestamp")) + if day is None: 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, {"scanDate": today, "modelPrefix": model_prefix, "stats": stats}) - return stats + message_id = entry.get("id") + timestamp = entry.get("timestamp") + if isinstance(message_id, str) and message_id: + key = ("entry", message_id, str(timestamp)) if timestamp else ("path", path, message_id) + if key in seen: + continue + seen.add(key) + + sessions.add(path) + 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(path) + today_token_total += total + today_tokens[model] = today_tokens.get(model, 0) + total + + 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), + } + return stats, complete + + +def cached_scan(max_age_seconds: float, model_prefix: str) -> dict[str, Any]: + roots = "\0".join(str(root) for root in pi_session_roots()) + digest = hashlib.sha1(roots.encode("utf-8")).hexdigest()[:16] + cache_file = cache_root() / f"ollama-scan-{digest}.json" + lock_file = cache_root() / f"ollama-scan-{digest}.lock" + + # 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. + def reusable() -> dict[str, Any] | None: + cached = read_json(cache_file) + stats = cached.get("stats") + if ( + cached.get("scanDate") == local_date_string() + and cached.get("modelPrefix") == model_prefix + and age_below(stamp(cached.get("scannedAtMs")) / 1000, max_age_seconds) + and isinstance(stats, dict) + ): + return stats + return None + + stats = reusable() + if stats is not None: + return stats + try: + with lock_file.open("w") as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + stats = reusable() + if stats is not None: + return stats + stats, complete = scan_sessions(model_prefix) + # A scan that missed a file would stand in for a real one until the + # next rescan; only a complete one is kept. + if complete: + write_json(cache_file, { + "scanDate": local_date_string(), + "scannedAtMs": round(time.time() * 1000), + "modelPrefix": model_prefix, + "stats": stats, + }) + return stats + except OSError as exc: + # The cache is an optimization; the record is the contract. + print(f"omarchy-agent-usage-ollama: cache unavailable ({exc}); scanning directly", file=sys.stderr) + return scan_sessions(model_prefix)[0] # ------------------------------------------------------------------- limits def clamp_fraction(value: Any) -> float: - try: - n = float(value) - except Exception: - return -1.0 - if not (n >= 0): + if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value) or value < 0: return -1.0 - return min(1.0, n) + return min(1.0, float(value)) def probe_limits(access_key: str) -> dict[str, Any]: @@ -292,9 +434,9 @@ def probe_limits(access_key: str) -> dict[str, Any]: 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, "auth": True, "helpText": AUTH_HELP_BAD_KEY} return {"ok": False, "helpText": f"Ollama's usage endpoint returned status {error.code}. Local stats are still shown."} - except Exception: + except OSError: # 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. @@ -303,6 +445,8 @@ def probe_limits(access_key: str) -> dict[str, Any]: "transport": True, "helpText": "Couldn't reach Ollama's usage endpoint. Retrying shortly. Local stats are still shown.", } + except ValueError: + return {"ok": False, "helpText": "Ollama's usage endpoint returned an invalid answer. Local stats are still shown."} limits = payload.get("limits") if isinstance(payload, dict) else None if not isinstance(limits, dict): @@ -324,7 +468,8 @@ def probe_limits(access_key: str) -> dict[str, Any]: return {"ok": True, "limits": out} -def probe_plan(access_key: str) -> str: +def probe_plan(access_key: str) -> str | None: + """The plan Ollama names for the key, "" when it names none, None when it couldn't be asked.""" request = urllib.request.Request( ME_ENDPOINT, headers={"Authorization": "Bearer " + access_key, "Accept": "application/json"}, @@ -334,47 +479,53 @@ def probe_plan(access_key: str) -> str: 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 None + plan = str(payload.get("Plan") or "").strip()[:64] if isinstance(payload, dict) else "" return plan[0].upper() + plan[1:] if plan else "" +def limits_cache_path(access_key: str) -> Path: + # Keyed by the credential: another key never shows this one's allowance or + # plan. + digest = hashlib.sha256(access_key.encode("utf-8")).hexdigest()[:16] + return cache_root() / f"ollama-limits-{digest}.json" + + def collect_limits(access_key: str, force: bool) -> dict[str, Any]: # `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. + # can tell kept limits from fresh ones. 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 - # 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 "") - result["fetchedAtMs"] = number(cached.get("fetchedAtMs")) - 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." + result["authHelpText"] = AUTH_HELP_NO_KEY return result + # A panel that is opened and shut repeatedly must not turn into a request + # per flick, so recent probe results are reused for a short window — and + # kept as the answer of record when a later probe fails. The endpoint names + # no reset times, so a kept window stays as it was measured. + probe_cache = limits_cache_path(access_key) + cached = read_json(probe_cache) + fallback = [entry for entry in cached.get("limits") or [] if isinstance(entry, dict)] if isinstance(cached.get("limits"), list) else [] + cached_plan = str(cached.get("plan") or "") + fetched_ms = stamp(cached.get("fetchedAtMs")) + # --force is a person asking for fresh numbers, so it skips the reuse window # entirely; the interval is there to absorb repeated panel opens, not to # overrule someone who pressed refresh. - fetched_at = number(cached.get("fetchedAtMs")) / 1000 - if fallback and not force and time.time() - fetched_at < PROBE_MIN_INTERVAL_SECONDS: - result["limits"] = fallback - result["live"] = True + if fallback and not force and age_below(fetched_ms / 1000, PROBE_MIN_INTERVAL_SECONDS): + result.update(limits=fallback, tierLabel=cached_plan, live=True, fetchedAtMs=fetched_ms) return result probe = probe_limits(access_key) if probe["ok"]: plan = probe_plan(access_key) - result["limits"] = probe["limits"] - result["tierLabel"] = plan - result.update(live=True, fetchedAtMs=round(time.time() * 1000)) + # A plan lookup that failed keeps the plan last named, rather than + # blanking a label the account still has. + plan = cached_plan if plan is None else plan + result.update(limits=probe["limits"], tierLabel=plan, live=True, fetchedAtMs=round(time.time() * 1000)) write_json(probe_cache, {"fetchedAtMs": result["fetchedAtMs"], "limits": probe["limits"], "plan": plan}) return result @@ -382,11 +533,13 @@ def collect_limits(access_key: str, force: bool) -> dict[str, Any]: # route. Ask the shell to try again sooner than its regular interval. if probe.get("transport"): result["retryAdvised"] = True + result["authHelpText"] = probe["helpText"] if fallback: - result["limits"] = fallback - else: + result.update(limits=fallback, tierLabel=cached_plan, fetchedAtMs=fetched_ms) + if probe.get("auth"): + result["usageStatusText"] = "Sign-in expired" + elif not fallback: result["usageStatusText"] = "Ollama limits unavailable" - result["authHelpText"] = probe["helpText"] return result @@ -397,33 +550,48 @@ 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) + scan_age = 0 if args.force else (LIMITS_ONLY_REUSE_SECONDS if args.limits_only else SCAN_REUSE_SECONDS) + stats = cached_scan(scan_age, model_prefix) + has_stats = bool(stats.get("modelUsage")) + key = api_key() 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": stats is not None, - "tierLabel": limits["tierLabel"], - "usageStatusText": limits["usageStatusText"], - "authHelpText": limits["authHelpText"], - "limits": limits["limits"], - "limitsStale": not limits["live"], - "limitsFetchedAt": limits["fetchedAtMs"], + "ready": False, + "hasLocalStats": has_stats, + "hasPromptStats": True, + "tierLabel": "", + "usageStatusText": "", + "authHelpText": "", + "limits": [], + "limitsStale": False, + "limitsFetchedAt": 0, } - if limits.get("retryAdvised"): - record["retryAdvised"] = True - record.update(stats or {}) + record.update(empty_stats()) + record.update(stats) + + # A machine with neither a key nor any Ollama Cloud sessions doesn't use it, + # and the panel skips the record; that is no trouble to report. + if key or has_stats: + limits = collect_limits(key, args.force) + record.update( + ready=True, + tierLabel=limits["tierLabel"], + usageStatusText=limits["usageStatusText"], + authHelpText=limits["authHelpText"], + limits=limits["limits"], + limitsStale=bool(limits["limits"]) and not limits["live"], + limitsFetchedAt=limits["fetchedAtMs"], + ) + if limits.get("retryAdvised"): + record["retryAdvised"] = True print(json.dumps(record, separators=(",", ":"), sort_keys=True)) return 0 diff --git a/test/shell.d/agent-usage-ollama-scanner-test.sh b/test/shell.d/agent-usage-ollama-scanner-test.sh index 7437010c7bf..4343f0d158c 100755 --- a/test/shell.d/agent-usage-ollama-scanner-test.sh +++ b/test/shell.d/agent-usage-ollama-scanner-test.sh @@ -8,14 +8,17 @@ 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" \ +# A machine that has never used Ollama Cloud must still print a full, +# hidden-by-default record with no status to alarm anyone: the update runner +# writes whatever valid JSON appears on stdout. +no_key=$(env -u PI_CODING_AGENT_DIR 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 + ":" + .tierLabel + ":" + (.limits | length | tostring) + ":" + (.hasLocalStats | tostring)' <<<"$no_key") == "ollama:false:Waiting for auth::0:false" ]] || +[[ $(jq -r '.id + ":" + (.ready | tostring) + ":" + .usageStatusText + ":" + .tierLabel + ":" + (.limits | length | tostring) + ":" + (.hasLocalStats | tostring)' <<<"$no_key") == "ollama:false:::0:false" ]] || fail "Ollama collector prints a valid record without credentials" "$no_key" -pass "Ollama collector prints a valid record without credentials" +[[ $(jq -c '{authHelpText, limitsStale, limitsFetchedAt, totalPrompts, days: (.recentDays | length), modelUsage, activeDates}' <<<"$no_key") == '{"authHelpText":"","limitsStale":false,"limitsFetchedAt":0,"totalPrompts":0,"days":7,"modelUsage":{},"activeDates":[]}' ]] || + fail "Ollama collector prints the full, zeroed stats block on an unused machine" "$no_key" +pass "Ollama collector prints an empty, unready record on an unused machine" # 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. @@ -42,6 +45,11 @@ result=$(HOME="$PI_HOME" XDG_CONFIG_HOME="$PI_HOME/.config" XDG_CACHE_HOME="$PI_ fail "Ollama collector counts usage from pi and omp sessions" "$result" pass "Ollama collector counts usage from pi and omp sessions" +# In use here, but no key to ask for the account's limits. +[[ $(jq -c '{ready, usageStatusText, limitsStale}' <<<"$result") == '{"ready":true,"usageStatusText":"Waiting for auth","limitsStale":false}' ]] || + fail "Ollama collector waits for auth when used without a key" "$result" +pass "Ollama collector waits for auth when used without a key" + [[ $(jq -c '.modelUsage' <<<"$result") == '{"deepseek-v4-pro":{"cacheCreationInputTokens":2,"cacheReadInputTokens":3,"inputTokens":15,"outputTokens":5},"kimi-k3":{"cacheCreationInputTokens":1,"cacheReadInputTokens":4,"inputTokens":20,"outputTokens":5}}' ]] || fail "Ollama collector filters to the ollama-cloud provider and merges model tags" "$result" pass "Ollama collector filters to the ollama-cloud provider and merges model tags" @@ -70,7 +78,7 @@ result=$(HOME="$PI_HOME" XDG_CONFIG_HOME="$PI_HOME/.config" XDG_CACHE_HOME="$PI_ 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" +cache_file=$(ls "$PI_HOME/.cache/omarchy/agent-usage/"ollama-scan-*.json | head -n 1) 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) @@ -126,8 +134,9 @@ pass "Ollama collector reports a payload without usable limits" CACHE_HOME=$(mktemp -d) trap 'rm -rf "$TEST_HOME" "$PI_HOME" "$CACHE_HOME"' EXIT +# force, cached-json (seeded for the key "key"), mode, [key] collect_limits() { - COLLECTOR="$ROOT/bin/omarchy-agent-usage-ollama" FORCE="$1" CACHED="$2" MODE="$3" \ + COLLECTOR="$ROOT/bin/omarchy-agent-usage-ollama" FORCE="$1" CACHED="$2" MODE="$3" KEY="${4:-key}" \ XDG_CACHE_HOME="$CACHE_HOME" python3 - <<'PY' import importlib.machinery, importlib.util, io, json, os @@ -136,7 +145,7 @@ 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" +cache = collector.limits_cache_path("key") cached = os.environ["CACHED"] if cached: cache.write_text(cached, encoding="utf-8") @@ -152,12 +161,16 @@ def urlopen(request, timeout=None): raise OSError("no route to host") if mode == "rejected": raise collector.urllib.error.HTTPError(request.full_url, 401, "Unauthorized", {}, io.BytesIO(b"")) + if mode == "server": + raise collector.urllib.error.HTTPError(request.full_url, 500, "Oops", {}, io.BytesIO(b"")) if request.full_url.endswith("/api/me"): + if mode == "planless": + raise OSError("plan lookup timed out") 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") +result = collector.collect_limits(os.environ["KEY"], os.environ["FORCE"] == "true") print(json.dumps({ "result": result, "probes": probes, @@ -184,17 +197,55 @@ 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" +[[ $(jq -c '{live: .result.live, fetchedAtMs: .result.fetchedAtMs, status: .result.usageStatusText}' <<<"$unreachable") == '{"live":false,"fetchedAtMs":1,"status":""}' ]] || + fail "Ollama collector marks kept limits stale with when they were fetched, without a status" "$unreachable" +# The endpoint names no reset times, so a kept window is kept as measured: +# never zeroed on a guessed reset, never given a reset time. +[[ $(jq -c '[.result.limits[] | {percent, resetsAt}]' <<<"$unreachable") == '[{"percent":0.11,"resetsAt":""}]' ]] || + fail "Ollama collector keeps a window without a reset time as it was measured" "$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" +# With nothing kept, a failed check says the limits are unavailable; only a +# transport failure asks for an early retry. +offline=$(collect_limits false "" transport) +[[ $(jq -c '.result | {usageStatusText, retryAdvised}' <<<"$offline") == '{"usageStatusText":"Ollama limits unavailable","retryAdvised":true}' ]] || + fail "Ollama collector reports unavailable limits when offline with nothing kept" "$offline" +server=$(collect_limits false "" server) +[[ $(jq -c '.result | {usageStatusText, retryAdvised}' <<<"$server") == '{"usageStatusText":"Ollama limits unavailable","retryAdvised":null}' ]] || + fail "Ollama collector does not advise a retry for a server's error status" "$server" +pass "Ollama collector advises a retry only when the endpoint can't be reached" + +# A rejected key is a sign-in to redo, and kept limits stay, stale. +rejected=$(collect_limits false "$stale" rejected) +[[ $(jq -r '.result.usageStatusText' <<<"$rejected") == "Sign-in expired" ]] || + fail "Ollama collector reports a rejected key as an expired sign-in" "$rejected" [[ $(jq -r '.result.authHelpText' <<<"$rejected") == *"rejected"* ]] || fail "Ollama collector explains a rejected key" "$rejected" -pass "Ollama collector reports a rejected key" +[[ $(jq -c '{live: .result.live, n: (.result.limits | length)}' <<<"$rejected") == '{"live":false,"n":1}' ]] || + fail "Ollama collector keeps the last limits past a rejected key" "$rejected" +pass "Ollama collector reports a rejected key as Sign-in expired" + +# Another key never shows this one's numbers or plan. +other=$(collect_limits false "$fresh" transport other-key) +[[ $(jq -c '.result | {n: (.limits | length), tierLabel, usageStatusText}' <<<"$other") == '{"n":0,"tierLabel":"","usageStatusText":"Ollama limits unavailable"}' ]] || + fail "Ollama collector keys its limits cache by credential" "$other" +pass "Ollama collector keeps each key's limits and plan apart" + +# A cache stamped in the future (the clock went back) is no proof of a recent +# probe, but it is still the last good answer to fall back on. +ahead=$(jq -nc --argjson ms "$(python3 -c 'import time; print(round((time.time() + 86400) * 1000))')" '{ + fetchedAtMs: $ms, limits: [{ label: "Weekly (7-day)", percent: 0.11, resetsAt: "" }], plan: "Pro" +}') +skewed=$(collect_limits false "$ahead" transport) +[[ $(jq -c '{probes: (.probes | length), n: (.result.limits | length), live: .result.live}' <<<"$skewed") == '{"probes":1,"n":1,"live":false}' ]] || + fail "Ollama collector re-probes past a future-stamped cache and still falls back on it" "$skewed" +pass "Ollama collector neither trusts nor loses a future-stamped cache" + +# A plan lookup that fails keeps the plan last named. +planless=$(collect_limits true "$fresh" planless) +[[ $(jq -r '.result.tierLabel' <<<"$planless") == "Pro" && $(jq -c '[.result.limits[].percent]' <<<"$planless") == "[0.44,0.11]" ]] || + fail "Ollama collector keeps the cached plan when the plan lookup fails" "$planless" +pass "Ollama collector keeps the cached plan when the plan lookup fails" # A fresh cache answers repeated panel opens without a request apiece. reused=$(collect_limits false "$fresh" normal) @@ -216,3 +267,48 @@ pass "Ollama collector re-probes on --force despite a fresh cache" [[ $(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" + +# Sessions outside the default root: an omp profile, a pi profile, and a pi +# agent dir moved by PI_CODING_AGENT_DIR, all under a $HOME that is a git +# checkout ignoring everything. A forked session repeats its parent's message +# (same id and timestamp) and counts once; lines without an id always count; +# rows whose timestamp or tokens are garbage drop without taking the rest. +ROOTS_HOME=$(mktemp -d) +trap 'rm -rf "$TEST_HOME" "$PI_HOME" "$CACHE_HOME" "$ROOTS_HOME"' EXIT +mkdir -p "$ROOTS_HOME/.git" +printf '*\n' >"$ROOTS_HOME/.gitignore" +mkdir -p "$ROOTS_HOME/.omp/profiles/work/agent/sessions/p" "$ROOTS_HOME/.pi/profiles/home/agent/sessions/p" "$ROOTS_HOME/moved/sessions/p" +line() { + printf '{"type":"message",%s"timestamp":%s,"message":{"role":"assistant","provider":"ollama-cloud","model":"kimi-k3","usage":{"input":%s}}}\n' "$1" "$2" "$3" +} +{ + line '"id":"parent",' "\"$timestamp\"" 100 + line '' "\"$timestamp\"" 1 + line '' "\"$timestamp\"" 1 +} >"$ROOTS_HOME/.omp/profiles/work/agent/sessions/p/a.jsonl" +{ + line '"id":"parent",' "\"$timestamp\"" 100 + line '"id":"child",' "\"$timestamp\"" 10 + line '"id":"bad-ts",' -5 5000 + line '"id":"bad-ts2",' '"not a date"' 5000 + line '"id":"bad-tokens",' "\"$timestamp\"" '"12"' + line '"id":"huge-tokens",' "\"$timestamp\"" 1e300 +} >"$ROOTS_HOME/.pi/profiles/home/agent/sessions/p/fork.jsonl" +line '"id":"moved",' "\"$timestamp\"" 1000 >"$ROOTS_HOME/moved/sessions/p/m.jsonl" + +roots=$(HOME="$ROOTS_HOME" PI_CODING_AGENT_DIR="$ROOTS_HOME/moved" XDG_CONFIG_HOME="$ROOTS_HOME/.config" XDG_CACHE_HOME="$ROOTS_HOME/.cache" \ + OLLAMA_API_KEY="" "$ROOT/bin/omarchy-agent-usage-ollama" --force) +[[ $(jq -r '.todayTotalTokens' <<<"$roots") == "1112" && $(jq -r '.totalPrompts' <<<"$roots") == "5" ]] || + fail "Ollama collector scans profiles and PI_CODING_AGENT_DIR past ignore rules, dedups forks, counts id-less lines, skips garbage rows" "$roots" +pass "Ollama collector scans every pi/omp root, dedups forks, and skips garbage rows" + +# A normal run reuses a scan for a few seconds only to absorb concurrent runs; +# --force always rescans. +line '"id":"later",' "\"$timestamp\"" 7 >>"$ROOTS_HOME/moved/sessions/p/m.jsonl" +again=$(HOME="$ROOTS_HOME" PI_CODING_AGENT_DIR="$ROOTS_HOME/moved" XDG_CONFIG_HOME="$ROOTS_HOME/.config" XDG_CACHE_HOME="$ROOTS_HOME/.cache" \ + OLLAMA_API_KEY="" "$ROOT/bin/omarchy-agent-usage-ollama") +forced=$(HOME="$ROOTS_HOME" PI_CODING_AGENT_DIR="$ROOTS_HOME/moved" XDG_CONFIG_HOME="$ROOTS_HOME/.config" XDG_CACHE_HOME="$ROOTS_HOME/.cache" \ + OLLAMA_API_KEY="" "$ROOT/bin/omarchy-agent-usage-ollama" --force) +[[ $(jq -r '.todayTotalTokens' <<<"$again") == "1112" && $(jq -r '.todayTotalTokens' <<<"$forced") == "1119" ]] || + fail "Ollama collector reuses a seconds-old scan on a normal run and rescans on --force" "$again $forced" +pass "Ollama collector reuses a seconds-old scan and rescans on --force" From b3652725a29dd24d1874ffea1879dc2a65f4c49a Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:09:34 -0400 Subject: [PATCH 079/546] Bring the MiniMax collector in line with the usage collector contract Key the limits cache by credential and host under cache_root(), zero elapsed kept windows instead of dropping them, keep the status empty when a failed check kept limits, report refused credentials as Sign-in expired, advise a retry only when the API couldn't be reached, reuse a probe for 15 seconds unless --force, and keep the last limits past an unexpected answer. Refuse redirects so the bearer never reaches another host, and skip media models' quota rows when choosing the plan's row. Co-authored-by: Josh LeBlanc Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-minimax | 275 +++++++++++++++++------ test/shell.d/agent-usage-minimax-test.sh | 201 ++++++++++++++--- 2 files changed, 374 insertions(+), 102 deletions(-) diff --git a/bin/omarchy-agent-usage-minimax b/bin/omarchy-agent-usage-minimax index 942041cea35..7311fc5f716 100755 --- a/bin/omarchy-agent-usage-minimax +++ b/bin/omarchy-agent-usage-minimax @@ -4,14 +4,21 @@ # omarchy:hidden=true """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. +The 5-hour session and weekly windows come from MiniMax's Token Plan API with +the key MiniMax's own `mmx` CLI or opencode saved. Each credential's last +good limits are cached, so a failed check keeps showing them, marked stale, +rather than dropping the section. 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 math import os import sys import tempfile @@ -20,17 +27,37 @@ import urllib.error import urllib.request from datetime import datetime, timedelta, timezone from pathlib import Path -from typing import Any +from typing import Any, Iterator AGENT_ID = "minimax" AGENT_NAME = "MiniMax" DEFAULT_BASE_URL = "https://api.minimax.io" +CN_BASE_URL = "https://api.minimaxi.com" AUTH_HELP = "Set MINIMAX_API_KEY, or sign in to MiniMax in opencode." +AUTH_HELP_BAD_KEY = "MiniMax rejected the saved credentials. Set MINIMAX_API_KEY, or sign in to MiniMax again with `mmx` or opencode." +# The plan the quota endpoint describes: only a Token Plan has these windows. TIER_LABEL = "Token Plan" +# 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 +# base_resp codes MiniMax answers a bad key with: 1004 is an authentication +# failure, 2049 an invalid API key. +AUTH_STATUS_CODES = {1004, 2049} +# The quota rows for MiniMax's media models ride on allowances of their own +# (video and image generation, speech, vision) and are not the coding plan. +MEDIA_MODEL_WORDS = ("video", "image", "audio", "vision") class MiniMaxError(Exception): - pass + def __init__(self, message: str, transport: bool = False, rejected: bool = False) -> None: + super().__init__(message) + self.transport = transport + self.rejected = rejected + + +def recent_dates() -> list[str]: + now = datetime.now() + return [(now - timedelta(days=offset)).strftime("%Y-%m-%d") for offset in range(6, -1, -1)] def empty_stats() -> dict[str, Any]: @@ -39,7 +66,7 @@ def empty_stats() -> dict[str, Any]: "todaySessions": 0, "todayTotalTokens": 0, "todayTokensByModel": {}, - "recentDays": [], + "recentDays": [{"date": day, "messageCount": 0} for day in recent_dates()], "totalPrompts": 0, "totalSessions": 0, "activeDays": 0, @@ -70,6 +97,12 @@ def base_record(**overrides: Any) -> dict[str, Any]: return record +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_json(path: Path) -> Any: try: return json.loads(path.read_text(encoding="utf-8")) @@ -85,14 +118,19 @@ def write_json(path: Path, payload: Any) -> None: try: with os.fdopen(fd, "w", encoding="utf-8") as handle: json.dump(payload, handle) + os.chmod(tmp, 0o644) 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 limits_cache_paths(api_key: str, base_url: str) -> tuple[Path, Path]: + # Keyed by the credential and the host it signs in to: another key, or the + # same key in the other region, never shows this one's allowance. + digest = hashlib.sha256(f"{base_url.rstrip('/')}\0{api_key}".encode("utf-8")).hexdigest()[:16] + root = cache_root() + return root / f"minimax-limits-{digest}.json", root / f"minimax-limits-{digest}.lock" def mmx_config_dir() -> Path: @@ -137,12 +175,12 @@ def config_base_url(config: dict[str, Any]) -> str: if isinstance(resource_url, str) and resource_url.strip(): return resource_url.strip() if oauth.get("region") == "cn": - return "https://api.minimaxi.com" + return CN_BASE_URL 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 CN_BASE_URL return DEFAULT_BASE_URL @@ -168,20 +206,25 @@ def credentials() -> tuple[str, str]: def number(value: Any) -> float | None: - try: - result = float(value) - return result if result == result else None - except Exception: + # Booleans are ints to Python; neither they nor NaN or infinity are numbers + # the quota API meant. + if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value): return None + return float(value) 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() + try: + end_ms = number(entry.get(end_key)) + if end_ms and end_ms > 0: + return datetime.fromtimestamp(end_ms / 1000, timezone.utc).isoformat() + # Without an end time MiniMax still says how long the window has left, + # which places its own reset rather than guessing one. + remains_ms = number(entry.get(remains_key)) + if remains_ms and remains_ms > 0: + return (datetime.now(timezone.utc) + timedelta(milliseconds=remains_ms)).isoformat() + except (OverflowError, OSError, ValueError): + pass return "" @@ -192,17 +235,23 @@ def used_fraction(remaining: Any) -> float | None: return max(0.0, min(1.0, 1.0 - percent / 100.0)) +def is_media_row(row: dict[str, Any]) -> bool: + name = str(row.get("model_name") or "").strip().lower() + return any(word in name for word in MEDIA_MODEL_WORDS) + + 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) + # "general" is the language/agent quota. Without it the first row that + # isn't a media model's stands for the plan. + rows = [row for row in rows if isinstance(row, dict)] + general = next((row for row in rows if row.get("model_name") == "general"), None) + if general is None: + general = next((row for row in rows if not is_media_row(row)), 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 = [] @@ -225,6 +274,16 @@ def quota_limits(payload: dict[str, Any]) -> list[dict[str, Any]]: return limits +class NoRedirect(urllib.request.HTTPRedirectHandler): + # A redirect would carry the bearer to whatever host it names. + def redirect_request(self, req, fp, code, msg, headers, newurl): # noqa: ANN001 + return None + + +def urlopen(request: urllib.request.Request, timeout: float) -> Any: + return urllib.request.build_opener(NoRedirect()).open(request, timeout=timeout) + + def fetch_quota(api_key: str, base_url: str) -> dict[str, Any]: endpoint = base_url.rstrip("/") + "/v1/token_plan/remains" request = urllib.request.Request( @@ -232,29 +291,37 @@ def fetch_quota(api_key: str, base_url: str) -> dict[str, Any]: headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}, ) try: - with urllib.request.urlopen(request, timeout=12) as response: + with 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") 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 + raise MiniMaxError(AUTH_HELP_BAD_KEY, rejected=True) from error + if 300 <= error.code < 400: + raise MiniMaxError("MiniMax's quota endpoint redirected elsewhere; the key was not sent on.") from error + raise MiniMaxError(f"MiniMax's quota endpoint returned HTTP {error.code}.") from error + except OSError as error: + # URLError, timeouts and resets alike: no answer from any server. + raise MiniMaxError("Couldn't reach the MiniMax API. Retrying shortly.", transport=True) from error + except ValueError 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") + 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) + if isinstance(base_response, dict): + code = number(base_response.get("status_code")) + if code not in (None, 0): + if int(code) in AUTH_STATUS_CODES: + raise MiniMaxError(AUTH_HELP_BAD_KEY, rejected=True) + message = str(base_response.get("status_msg") or "MiniMax rejected the quota request")[:200] + raise MiniMaxError(message) 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]]: +# A kept window whose reset has passed comes back at zero: the allowance it +# measured has been renewed, so it is untouched rather than unknown. One with +# no reset time, or one that will not parse, is kept as it is. +def usable_cached_limits(limits: Any) -> list[dict[str, Any]]: now = datetime.now(timezone.utc) kept = [] for limit in limits if isinstance(limits, list) else []: @@ -266,44 +333,118 @@ def open_windows(limits: Any) -> list[dict[str, Any]]: resets = None if resets and resets.tzinfo is None: resets = resets.replace(tzinfo=timezone.utc) - if not resets or resets > now: - kept.append(limit) + kept.append(dict(limit, percent=0.0, resetsAt="") if resets and resets <= now else 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]: - if not api_key: - return base_record(usageStatusText="MiniMax unavailable", authHelpText=AUTH_HELP) +def cached_answer(cache_file: Path) -> tuple[list[dict[str, Any]], int]: + cached = read_json(cache_file) + cached = cached if isinstance(cached, dict) else {} + fetched_ms = number(cached.get("fetchedAtMs")) + return usable_cached_limits(cached.get("limits")), max(0, int(fetched_ms or 0)) - 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=open_windows(cached.get("limits")), - 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) + +def recent_probe(fetched_ms: int) -> bool: + # A negative age is a stamp from the future: the clock moved back since it + # was written, so it proves nothing about freshness. + age = time.time() - fetched_ms / 1000 + return 0 <= age < PROBE_MIN_INTERVAL_SECONDS + + +@contextlib.contextmanager +def probe_lock(path: Path) -> Iterator[None]: + # One probe at a time per credential; a lock that can't be taken only costs + # a duplicate request, so the probe runs unlocked rather than not at all. try: - write_json(limits_cache(), {"fetchedAtMs": fetched_at, "limits": limits}) + handle = path.open("a") except OSError: - pass + yield + return + with handle: + try: + fcntl.flock(handle, fcntl.LOCK_EX) + except OSError: + pass + yield + + +def collect_limits(api_key: str, base_url: str, force: bool) -> dict[str, Any]: + """Fresh limits when a probe is due, else the last known ones. + + `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. + """ + result = {"limits": [], "usageStatusText": "", "authHelpText": "", "retryAdvised": False, "live": False, "fetchedAtMs": 0} + cache_file, lock_file = limits_cache_paths(api_key, base_url) + + # --force is a person asking for fresh numbers, so it skips the reuse + # window; the interval absorbs repeated panel opens, not a refresh press. + fallback, fetched_ms = cached_answer(cache_file) + if fallback and not force and recent_probe(fetched_ms): + result.update(limits=fallback, live=True, fetchedAtMs=fetched_ms) + return result + + with probe_lock(lock_file): + if not force: + # Another run may have probed while this one waited for the lock. + fallback, fetched_ms = cached_answer(cache_file) + if fallback and recent_probe(fetched_ms): + result.update(limits=fallback, live=True, fetchedAtMs=fetched_ms) + return result + try: + limits = quota_limits(fetch_quota(api_key, base_url)) + except MiniMaxError as error: + failure = error + except Exception as error: + # An answer shaped like nothing expected is still a failed check: the + # kept limits stand. + failure = MiniMaxError(f"MiniMax's quota check failed ({type(error).__name__}).") + else: + result.update(limits=limits, live=True, fetchedAtMs=round(time.time() * 1000)) + try: + write_json(cache_file, {"fetchedAtMs": result["fetchedAtMs"], "limits": limits}) + except OSError as exc: + print(f"omarchy-agent-usage-minimax: could not write limits cache ({exc})", file=sys.stderr) + return result + + # A transport failure reached no server — typically the seconds after login + # before the network is up — so ask the shell to retry sooner. + result["retryAdvised"] = failure.transport + result["authHelpText"] = str(failure) + if fallback: + result.update(limits=fallback, fetchedAtMs=fetched_ms) + if failure.rejected: + result["usageStatusText"] = "Sign-in expired" + elif not fallback: + result["usageStatusText"] = "MiniMax limits unavailable" + return result + + +# 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, force: bool = False) -> dict[str, Any]: + if not api_key: + return base_record() + + limits = collect_limits(api_key, base_url, force) + record = base_record( + ready=True, + tierLabel=TIER_LABEL if limits["limits"] else "", + limits=limits["limits"], + limitsStale=bool(limits["limits"]) and not limits["live"], + limitsFetchedAt=limits["fetchedAtMs"], + usageStatusText=limits["usageStatusText"], + authHelpText=limits["authHelpText"], + ) + if limits["retryAdvised"]: + record["retryAdvised"] = True return record 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. + # --force re-probes past the reuse window. --limits-only takes the same path + # as a normal run: there is no local scan for it to reuse. parser.add_argument("--force", action="store_true") parser.add_argument("--limits-only", action="store_true") parser.add_argument("--base-url") @@ -311,9 +452,9 @@ def main() -> int: try: api_key, saved_base_url = credentials() - record = collect(api_key, args.base_url or saved_base_url) + record = collect(api_key, args.base_url or saved_base_url, args.force) except Exception as error: - record = base_record(usageStatusText="MiniMax limits unavailable", authHelpText="MiniMax usage scan failed") + 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 diff --git a/test/shell.d/agent-usage-minimax-test.sh b/test/shell.d/agent-usage-minimax-test.sh index 37ff0f769a5..2b9a46a2b4d 100755 --- a/test/shell.d/agent-usage-minimax-test.sh +++ b/test/shell.d/agent-usage-minimax-test.sh @@ -8,10 +8,12 @@ require_command python3 TEST_HOME=$(mktemp -d) trap 'rm -rf "$TEST_HOME"' EXIT +# A machine with no MiniMax key prints the full, zeroed record the panel +# skips, with no status to alarm anyone. 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 -c '{id, ready, tierLabel, limits}' <<<"$no_key") == '{"id":"minimax","ready":false,"tierLabel":"","limits":[]}' ]] || +[[ $(jq -c '{id, ready, tierLabel, limits, usageStatusText, authHelpText, limitsStale, limitsFetchedAt, days: (.recentDays | length), modelUsage}' <<<"$no_key") == '{"id":"minimax","ready":false,"tierLabel":"","limits":[],"usageStatusText":"","authHelpText":"","limitsStale":false,"limitsFetchedAt":0,"days":7,"modelUsage":{}}' ]] || fail "MiniMax collector prints a record the panel skips without credentials" "$no_key" pass "MiniMax collector prints a record the panel skips without credentials" @@ -29,6 +31,8 @@ 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") +os.environ.pop("MINIMAX_API_KEY", None) +os.environ.pop("MINIMAX_BASE_URL", None) loader = importlib.machinery.SourceFileLoader("minimax_collector", collector_path) spec = importlib.util.spec_from_loader(loader.name, loader) @@ -57,6 +61,14 @@ payload = { } limits = scanner.quota_limits(payload) +fetch_quota = scanner.fetch_quota + +# Without a "general" row the first row that isn't a media model's is the plan. +media_first = scanner.quota_limits({"model_remains": [ + {"model_name": "video", "current_interval_remaining_percent": 100}, + {"model_name": "Image-01", "current_interval_remaining_percent": 100}, + {"model_name": "MiniMax-M2", "current_interval_remaining_percent": 25}, +]}) config_dir = home / ".mmx" config_dir.mkdir(parents=True) @@ -79,45 +91,89 @@ 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 +import copy, time +from datetime import datetime, timedelta, timezone # 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") +probes = [] -def unreachable(api_key, base_url): - raise scanner.MiniMaxError("Could not reach the MiniMax API") +def answer(api_key, base_url): + probes.append(api_key) + return current -scanner.fetch_quota = unreachable -stale = scanner.collect("mmx_test", "https://example.invalid") +def failing(**kwargs): + def fetch(api_key, base_url): + probes.append(api_key) + raise scanner.MiniMaxError(kwargs.get("message", "failed"), transport=kwargs.get("transport", False), rejected=kwargs.get("rejected", False)) + return fetch -# A cached window that has since reset is dropped; one still open, or with -# no reset time, stays. -from datetime import datetime, timedelta, timezone +def run(fetch, key="mmx_test", force=False): + scanner.fetch_quota = fetch + probes.clear() + record = scanner.collect(key, "https://example.invalid", force) + return dict(record, probes=len(probes)) + +out = {"limits": limits, "mediaFirst": media_first, "savedKey": saved_key, "savedBaseUrl": saved_base_url, + "oauthKey": oauth_key, "oauthBaseUrl": oauth_base_url, "opencodeKey": opencode_key} + +out["record"] = run(answer, force=True) +# Within the probe interval a repeat run reuses the answer; --force doesn't. +out["reused"] = run(failing(transport=True)) +out["forced"] = run(answer, force=True) + +cache = scanner.limits_cache_paths("mmx_test", "https://example.invalid")[0] +def age_cache(fetched_ms, entries=None): + data = json.loads(cache.read_text()) + data["fetchedAtMs"] = fetched_ms + if entries is not None: + data["limits"] = entries + cache.write_text(json.dumps(data)) + +age_cache(1) +out["stale"] = run(failing(message="Couldn't reach the MiniMax API. Retrying shortly.", transport=True)) +out["rejected"] = run(failing(message=scanner.AUTH_HELP_BAD_KEY, rejected=True)) +out["server"] = run(failing(message="MiniMax's quota endpoint returned HTTP 500.")) + +def broken(api_key, base_url): + probes.append(api_key) + raise KeyError("surprise") +out["broken"] = run(broken) + +# Another key never shows this one's numbers. +out["otherKey"] = run(failing(message="MiniMax's quota endpoint returned HTTP 500."), key="other") +out["offline"] = run(failing(message="Couldn't reach the MiniMax API.", transport=True), key="other") + +# A cache stamped in the future (the clock went back) is no proof of a recent +# probe, but it is still the last good answer to fall back on. +age_cache(round((time.time() + 86400) * 1000)) +out["skewed"] = run(failing(transport=True)) + +# A kept window that has since reset comes back at zero; one still open, or +# with no reset time, stays as it was. now = datetime.now(timezone.utc) -scanner.write_json(scanner.limits_cache(), {"fetchedAtMs": 1, "limits": [ +age_cache(1, [ {"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, - "savedBaseUrl": saved_base_url, - "oauthKey": oauth_key, - "oauthBaseUrl": oauth_base_url, - "opencodeKey": opencode_key, - "record": record, - "stale": stale, - "reset": reset, -})) +]) +out["reset"] = run(failing(transport=True)) + +# A bad key MiniMax names in base_resp rather than with a status code. +import io +class Response(io.BytesIO): + def __enter__(self): return self + def __exit__(self, *a): return False +scanner.urlopen = lambda request, timeout: Response(json.dumps({"base_resp": {"status_code": 1004, "status_msg": "login fail"}}).encode()) +try: + fetch_quota("k", "https://example.invalid") + out["baseRespRejected"] = False +except scanner.MiniMaxError as error: + out["baseRespRejected"] = error.rejected + +print(json.dumps(out)) PY ) @@ -129,6 +185,10 @@ pass "MiniMax remaining percentages become Omarchy used fractions" fail "MiniMax quota reset timestamps are normalized" "$result" pass "MiniMax quota reset timestamps are normalized" +[[ $(jq -c '.mediaFirst | map(.percent)' <<<"$result") == '[0.75]' ]] || + fail "MiniMax collector skips media quota rows when choosing the plan's row" "$result" +pass "MiniMax collector skips media quota rows when choosing the plan's row" + [[ $(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" @@ -137,16 +197,87 @@ 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, limitsStale}' <<<"$result") == '{"id":"minimax","ready":true,"hasLocalStats":false,"hasPromptStats":false,"scope":"account","tierLabel":"Token Plan","limitsStale":false}' ]] || +[[ $(jq -c '.record | {id, ready, tierLabel, limitsStale, usageStatusText, authHelpText, probes}' <<<"$result") == '{"id":"minimax","ready":true,"tierLabel":"Token Plan","limitsStale":false,"usageStatusText":"","authHelpText":"","probes":1}' ]] || 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"}' ]] || +[[ $(jq -c '[.reused.probes, .reused.limitsStale, .forced.probes]' <<<"$result") == '[0,false,1]' ]] || + fail "MiniMax collector reuses a probe younger than the interval and re-probes on --force" "$result" +pass "MiniMax collector reuses a recent probe and re-probes on --force" + +[[ $(jq -c '.stale | {ready, tierLabel, limitsStale, percent: .limits[0].percent, usageStatusText, retryAdvised, limitsFetchedAt}' <<<"$result") == '{"ready":true,"tierLabel":"Token Plan","limitsStale":true,"percent":0.9,"usageStatusText":"","retryAdvised":true,"limitsFetchedAt":1}' ]] || 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" +[[ $(jq -r '.stale.authHelpText' <<<"$result") == *"reach"* ]] || + fail "MiniMax collector says what failed when it keeps the last limits" "$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" +[[ $(jq -c '.rejected | {usageStatusText, limitsStale, n: (.limits | length), retryAdvised}' <<<"$result") == '{"usageStatusText":"Sign-in expired","limitsStale":true,"n":2,"retryAdvised":null}' ]] || + fail "MiniMax collector reports refused credentials as an expired sign-in" "$result" +[[ $(jq -r '.baseRespRejected' <<<"$result") == "true" ]] || + fail "MiniMax collector reads an authentication failure in base_resp as a refused key" "$result" +pass "MiniMax collector reports refused credentials as Sign-in expired" + +[[ $(jq -c '[.server.usageStatusText, .server.limitsStale, .server.retryAdvised, .broken.usageStatusText, .broken.limitsStale]' <<<"$result") == '["",true,null,"",true]' ]] || + fail "MiniMax collector keeps the last limits past a server error or an unexpected answer" "$result" +pass "MiniMax collector keeps the last limits past a server error or an unexpected answer" + +[[ $(jq -c '.otherKey | {usageStatusText, n: (.limits | length), tierLabel, retryAdvised}' <<<"$result") == '{"usageStatusText":"MiniMax limits unavailable","n":0,"tierLabel":"","retryAdvised":null}' ]] || + fail "MiniMax collector keys its limits cache by credential and reports a failure with nothing kept" "$result" +[[ $(jq -c '.offline | {usageStatusText, retryAdvised}' <<<"$result") == '{"usageStatusText":"MiniMax limits unavailable","retryAdvised":true}' ]] || + fail "MiniMax collector advises a retry only when the API can't be reached" "$result" +pass "MiniMax collector keeps each key's limits apart and advises a retry only for transport failures" + +[[ $(jq -c '.skewed | {probes, limitsStale, n: (.limits | length)}' <<<"$result") == '{"probes":1,"limitsStale":true,"n":2}' ]] || + fail "MiniMax collector re-probes past a future-stamped cache and still falls back on it" "$result" +pass "MiniMax collector neither trusts nor loses a future-stamped cache" + +[[ $(jq -c '.reset.limits | map({label, percent, resetsAt: (.resetsAt != "")})' <<<"$result") == '[{"label":"Session (5-hour)","percent":0.0,"resetsAt":false},{"label":"Weekly (7-day)","percent":0.5,"resetsAt":true},{"label":"Other","percent":0.2,"resetsAt":false}]' ]] || + fail "MiniMax collector zeroes kept windows that have reset rather than dropping them" "$result" +pass "MiniMax collector zeroes kept windows that have reset" + +# The bearer never follows a redirect to another host. +redirect=$(python3 - "$ROOT/bin/omarchy-agent-usage-minimax" <<'PY' +import importlib.machinery, importlib.util, json, sys, threading +from http.server import BaseHTTPRequestHandler, HTTPServer + +loader = importlib.machinery.SourceFileLoader("minimax_collector", sys.argv[1]) +spec = importlib.util.spec_from_loader(loader.name, loader) +scanner = importlib.util.module_from_spec(spec) +loader.exec_module(scanner) + +seen = [] + +class Elsewhere(BaseHTTPRequestHandler): + def do_GET(self): + seen.append(self.headers.get("Authorization")) + self.send_response(200) + self.end_headers() + self.wfile.write(b'{"model_remains":[]}') + def log_message(self, *args): + pass + +elsewhere = HTTPServer(("127.0.0.1", 0), Elsewhere) + +class Redirecting(BaseHTTPRequestHandler): + def do_GET(self): + self.send_response(302) + self.send_header("Location", f"http://127.0.0.1:{elsewhere.server_port}/steal") + self.end_headers() + def log_message(self, *args): + pass + +origin = HTTPServer(("127.0.0.1", 0), Redirecting) +for server in (origin, elsewhere): + threading.Thread(target=server.serve_forever, daemon=True).start() + +try: + scanner.fetch_quota("secret", f"http://127.0.0.1:{origin.server_port}") + error = "" +except scanner.MiniMaxError as failure: + error = str(failure) +print(json.dumps({"forwarded": seen, "error": error})) +PY +) +[[ $(jq -c '.forwarded' <<<"$redirect") == '[]' && $(jq -r '.error' <<<"$redirect") == *"redirected"* ]] || + fail "MiniMax collector refuses to follow a redirect with the bearer" "$redirect" +pass "MiniMax collector refuses to follow a redirect with the bearer" From ad7306811b441998cbaec1cc38b0aa851dc08438 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:14:07 -0400 Subject: [PATCH 080/546] Count MiniMax Code, pi, omp, and opencode usage in the MiniMax collector MiniMax Code's own sessions under MINIMAX_DATA_DIR, MAVIS_DATA_DIR, or ~/.minimax count user messages as prompts and assistant usage as tokens, and pi/omp and opencode sessions on a MiniMax provider count the way the other provider collectors count theirs. The scan is cached per day under flock, and only when every source was read. Without a key, local usage still shows and the limits wait for one; keys pi saved are read too. Co-authored-by: Josh LeBlanc Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-minimax | 483 +++++++++++++++++- .../agent-usage-minimax-scanner-test.sh | 109 ++++ test/shell.d/agent-usage-minimax-test.sh | 15 +- 3 files changed, 579 insertions(+), 28 deletions(-) create mode 100644 test/shell.d/agent-usage-minimax-scanner-test.sh diff --git a/bin/omarchy-agent-usage-minimax b/bin/omarchy-agent-usage-minimax index 7311fc5f716..7c9bd19327d 100755 --- a/bin/omarchy-agent-usage-minimax +++ b/bin/omarchy-agent-usage-minimax @@ -2,24 +2,30 @@ # 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 usage into a display-ready JSON record. The 5-hour session and weekly windows come from MiniMax's Token Plan API with -the key MiniMax's own `mmx` CLI or opencode saved. Each credential's last +the key MiniMax's own `mmx` CLI, opencode, or pi saved. Each credential's last good limits are cached, so a failed check keeps showing them, marked stale, -rather than dropping the section. The agents panel only ever reads the JSON -this prints. +rather than dropping the section. + +Local stats come from MiniMax Code's own session files, and from pi/omp and +opencode sessions that ran on a MiniMax provider. The agents panel only ever +reads the JSON this prints. """ from __future__ import annotations import argparse +import base64 import contextlib import fcntl import hashlib import json import math import os +import sqlite3 +import subprocess import sys import tempfile import time @@ -33,13 +39,27 @@ AGENT_ID = "minimax" AGENT_NAME = "MiniMax" DEFAULT_BASE_URL = "https://api.minimax.io" CN_BASE_URL = "https://api.minimaxi.com" -AUTH_HELP = "Set MINIMAX_API_KEY, or sign in to MiniMax in opencode." +AUTH_HELP = "Set MINIMAX_API_KEY, or sign in to MiniMax in opencode, to show Token Plan limits." AUTH_HELP_BAD_KEY = "MiniMax rejected the saved credentials. Set MINIMAX_API_KEY, or sign in to MiniMax again with `mmx` or opencode." # The plan the quota endpoint describes: only a Token Plan has these windows. TIER_LABEL = "Token Plan" +# The pi/omp and opencode provider names MiniMax bills, by region. Exact +# names: a custom gateway in front of MiniMax is not this account. +PROVIDERS = ("minimax", "minimax-cn", "minimax-coding-plan", "minimax-cn-coding-plan") +# 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). --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 +# Bounds for numbers read from session files. A timestamp outside them, or a +# token count above the cap, is corruption rather than usage. +EARLIEST_SECONDS = 946_684_800 # 2000-01-01 +LATEST_SECONDS = 4_102_444_800 # 2100-01-01 +MAX_TOKENS = 10_000_000_000 # base_resp codes MiniMax answers a bad key with: 1004 is an authentication # failure, 2049 an invalid API key. AUTH_STATUS_CODES = {1004, 2049} @@ -83,8 +103,7 @@ def base_record(**overrides: Any) -> dict[str, Any]: "updatedAt": datetime.now(timezone.utc).isoformat(), "ready": False, "hasLocalStats": False, - "hasPromptStats": False, - "scope": "account", + "hasPromptStats": True, "tierLabel": "", "usageStatusText": "", "authHelpText": "", @@ -184,25 +203,44 @@ def config_base_url(config: dict[str, Any]) -> str: 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") +def opencode_data_dir() -> Path: + return Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local/share")) / "opencode" + + +def pi_agent_dir() -> Path: + configured = os.environ.get("PI_CODING_AGENT_DIR", "").strip() + return Path(os.path.expanduser(configured)) if configured else Path.home() / ".pi" / "agent" + + +def saved_key(path: Path, providers: tuple[str, ...]) -> tuple[str, str]: + """(key, provider) for the first provider `path` holds a literal key for.""" + parsed = read_json(path) if not isinstance(parsed, dict): - return "" - for provider in ("minimax", "minimax-coding", "minimax-token-plan"): + return "", "" + for provider in providers: 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 "" + # A pi `$VAR` or `!command` entry is resolved by pi alone. + if isinstance(value, str) and value.strip() and not value.strip().startswith(("!", "$")): + return value.strip(), provider + 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 + api_key = os.environ.get("MINIMAX_API_KEY", "").strip() or config_key(config) + base_url = config_base_url(config) + if not api_key: + # The key opencode or pi signed in to MiniMax with; a China-region + # provider's key only answers on the China host. + for path in (opencode_data_dir() / "auth.json", pi_agent_dir() / "auth.json", Path.home() / ".omp" / "agent" / "auth.json"): + api_key, provider = saved_key(path, ("minimax", "minimax-coding", "minimax-token-plan") + PROVIDERS[1:]) + if api_key: + if provider.startswith("minimax-cn"): + base_url = CN_BASE_URL + break + return api_key, os.environ.get("MINIMAX_BASE_URL", "").strip() or base_url def number(value: Any) -> float | None: @@ -420,14 +458,396 @@ def collect_limits(api_key: str, base_url: str, force: bool) -> dict[str, Any]: return result -# 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, force: bool = False) -> dict[str, Any]: +# -------------------------------------------------------------------------- +# Local session stats +# -------------------------------------------------------------------------- + +def token_count(value: Any) -> int: + # Booleans are ints to Python, and a NaN or negative count would poison + # every sum it reaches, so anything but a plain finite count reads as none. + if isinstance(value, bool) or not isinstance(value, (int, float)): + return 0 + if not math.isfinite(value) or value < 0 or value > MAX_TOKENS: + return 0 + return int(value) + + +def local_day(value: Any) -> str | None: + """The local date of an untrusted timestamp, or None when it isn't one.""" + if isinstance(value, bool): + return None + if isinstance(value, (int, float)): + if not math.isfinite(value): + return None + # MiniMax Code, pi, and opencode timestamps are milliseconds. + seconds = value / 1000 if value > 10_000_000_000 else value + if not EARLIEST_SECONDS <= seconds < LATEST_SECONDS: + return None + return datetime.fromtimestamp(seconds).strftime("%Y-%m-%d") + if not isinstance(value, str) or not value.strip(): + return None + try: + parsed = datetime.fromisoformat(value.strip().replace("Z", "+00:00")) + except ValueError: + return None + if parsed.tzinfo is not None: + parsed = parsed.astimezone() + if not 2000 <= parsed.year < 2100: + return None + return parsed.strftime("%Y-%m-%d") + + +def empty_bucket() -> dict[str, int]: + return { + "inputTokens": 0, + "outputTokens": 0, + "cacheReadInputTokens": 0, + "cacheCreationInputTokens": 0, + } + + +class Stats: + """Token and activity totals shared by every local source. + + Sessions and days are sets, so a day active in two sources counts once + toward activeDays. + """ + + def __init__(self) -> None: + self.today = datetime.now().strftime("%Y-%m-%d") + self.recent_dates = recent_dates() + self.recent = {day: 0 for day in self.recent_dates} + self.today_prompts = 0 + self.today_tokens = 0 + self.today_sessions: set[str] = set() + self.today_by_model: dict[str, int] = {} + self.total_prompts = 0 + self.sessions: set[str] = set() + self.active_dates: set[str] = set() + self.model_usage: dict[str, dict[str, int]] = {} + + def touch(self, session: str, day: str) -> None: + self.sessions.add(session) + self.active_dates.add(day) + if day == self.today: + self.today_sessions.add(session) + + def prompt(self, session: str, day: str) -> None: + self.touch(session, day) + self.total_prompts += 1 + if day == self.today: + self.today_prompts += 1 + + def tokens(self, session: str, day: str, model: str, inp: int, out: int, cread: int, cwrite: int) -> None: + self.touch(session, day) + total = inp + out + cread + cwrite + bucket = self.model_usage.setdefault(model, empty_bucket()) + bucket["inputTokens"] += inp + bucket["outputTokens"] += out + bucket["cacheReadInputTokens"] += cread + bucket["cacheCreationInputTokens"] += cwrite + if day in self.recent: + self.recent[day] += total + if day == self.today: + self.today_tokens += total + self.today_by_model[model] = self.today_by_model.get(model, 0) + total + + def record(self) -> dict[str, Any]: + return { + "todayPrompts": self.today_prompts, + "todaySessions": len(self.today_sessions), + "todayTotalTokens": self.today_tokens, + "todayTokensByModel": self.today_by_model, + "recentDays": [{"date": day, "messageCount": self.recent[day]} for day in self.recent_dates], + "totalPrompts": self.total_prompts, + "totalSessions": len(self.sessions), + "activeDays": len(self.active_dates), + "activeDates": sorted(self.active_dates), + "modelUsage": self.model_usage, + } + + +def pi_style_tokens(usage: dict[str, Any]) -> tuple[int, int, int, int]: + inp = token_count(usage.get("input")) + out = token_count(usage.get("output")) + cread = token_count(usage.get("cacheRead")) + cwrite = token_count(usage.get("cacheWrite")) + if not (inp or out): + # Only a total: MiniMax's input excludes the cache tokens, which keep + # their own buckets. + inp = max(0, token_count(usage.get("totalTokens")) - cread - cwrite) + return inp, out, cread, cwrite + + +def mcode_data_dir() -> Path: + for name in ("MINIMAX_DATA_DIR", "MAVIS_DATA_DIR"): + value = os.environ.get(name, "").strip() + if value: + return Path(os.path.abspath(os.path.expanduser(value))) + return Path.home() / ".minimax" + + +def scan_mcode(stats: Stats) -> bool: + """Add MiniMax Code's own sessions: user messages are prompts, assistant usage is tokens.""" + root = mcode_data_dir() / "v2" / "sessions" + if not root.is_dir(): + return True + complete = True + for path in sorted(root.rglob("messages.jsonl")): + # One directory per session; its messages carry no session id of their own. + session = "mcode:" + str(path.parent) + try: + with path.open("r", encoding="utf-8", errors="replace") as handle: + for raw in handle: + try: + entry = json.loads(raw) + except ValueError: + continue + message = entry.get("message") if isinstance(entry, dict) else None + if not isinstance(message, dict): + continue + day = local_day(message.get("timestamp")) or local_day(entry.get("timestamp")) + if day is None: + continue + if message.get("role") == "user": + stats.prompt(session, day) + continue + usage = message.get("usage") + if message.get("role") != "assistant" or not isinstance(usage, dict): + continue + inp, out, cread, cwrite = pi_style_tokens(usage) + if not (inp or out or cread or cwrite): + continue + model = str(message.get("model") or entry.get("model") or "minimax").strip() or "minimax" + stats.tokens(session, day, model, inp, out, cread, cwrite) + except OSError: + complete = False + return complete + + +# `omp --profile=` (and pi's equivalent) relocates the whole agent tree +# under /profiles//, and PI_CODING_AGENT_DIR moves pi's, so usage +# driven through either would otherwise go uncounted. +def pi_session_roots() -> list[Path]: + roots = [pi_agent_dir() / "sessions", Path.home() / ".omp" / "agent" / "sessions"] + for base in (Path.home() / ".pi", Path.home() / ".omp"): + try: + # Sorted so the scan order does not depend on directory order. + roots.extend(sorted(child / "agent" / "sessions" for child in (base / "profiles").iterdir() if child.is_dir())) + except OSError: + # No profiles directory, or it is unreadable: the default roots stand. + pass + unique: list[Path] = [] + for root in roots: + if root not in unique: + unique.append(root) + return unique + + +def matching_lines(root: Path, pattern: str) -> tuple[list[tuple[str, str]], bool]: + """(path, line) for each session line matching `pattern`, and whether every file was read.""" + try: + proc = subprocess.run( + # --no-ignore: these are data files, not a source tree. When $HOME is + # itself a git checkout (a common dotfiles setup), its ignore rules + # would otherwise hide every session file and count zero usage. + ["rg", "--json", "--no-ignore", "--hidden", "--glob", "*.jsonl", "-e", pattern, str(root)], + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + text=True, + errors="replace", + check=False, + ) + except OSError: + return [], False + lines = [] + for raw in proc.stdout.splitlines(): + try: + event = json.loads(raw) + if event.get("type") != "match": + continue + data = event["data"] + text = data["lines"].get("text") + if text is None: + # rg hands back lines that are not valid UTF-8 as base64. + text = base64.b64decode(data["lines"]["bytes"]).decode("utf-8", errors="replace") + lines.append((data["path"].get("text") or "", text)) + except Exception: + continue + # 1 is "no match"; 2 means some file could not be read. + return lines, proc.returncode in (0, 1) + + +def scan_pi(stats: Stats) -> bool: + """Add pi/omp assistant turns that ran on a MiniMax provider.""" + complete = True + # A forked session copies its parent's messages, ids and timestamps alike, + # so (id, timestamp) counts each once across files. A line with an id but + # no timestamp dedups within its file only; a line with no id always counts. + seen: set[tuple[str, ...]] = set() + pattern = r'"provider"\s*:\s*"(' + "|".join(PROVIDERS) + r')"' + for root in pi_session_roots(): + if not root.is_dir(): + continue + lines, read_all = matching_lines(root, pattern) + complete = complete and read_all + for path, line in lines: + try: + entry = json.loads(line) + except ValueError: + continue + if not isinstance(entry, dict) or entry.get("type") != "message": + continue + message = entry.get("message") + if not isinstance(message, dict) or message.get("role") != "assistant": + continue + if str(message.get("provider") or "") not in PROVIDERS: + continue + usage = message.get("usage") + if not isinstance(usage, dict): + continue + inp, out, cread, cwrite = pi_style_tokens(usage) + if not (inp or out or cread or cwrite): + continue + day = local_day(entry.get("timestamp")) or local_day(message.get("timestamp")) + if day is None: + continue + message_id = entry.get("id") + timestamp = entry.get("timestamp") + if isinstance(message_id, str) and message_id: + key = ("entry", message_id, str(timestamp)) if timestamp else ("path", path, message_id) + if key in seen: + continue + seen.add(key) + # pi counts each assistant turn as the prompt that asked for it. + stats.prompt(path, day) + stats.tokens(path, day, str(message.get("model") or "minimax"), inp, out, cread, cwrite) + return complete + + +def scan_opencode(stats: Stats) -> bool: + """Add opencode messages on a MiniMax provider. + + Read-only: opencode may be writing right now. + """ + db = opencode_data_dir() / "opencode.db" + if not db.is_file(): + return True + placeholders = ",".join("?" for _ in PROVIDERS) + 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 json_extract conditions are the authority; the LIKE gates are pure + # acceleration that skip rows whose JSON cannot name a MiniMax provider. + rows = conn.execute( + "SELECT session_id, data FROM message" + " WHERE data LIKE '%\"role\"%:%\"assistant\"%'" + " AND data LIKE '%\"providerID\"%:%\"minimax%'" + " AND CASE WHEN json_valid(data) THEN json_extract(data, '$.role') END = 'assistant'" + f" AND CASE WHEN json_valid(data) THEN json_extract(data, '$.providerID') END IN ({placeholders})", + PROVIDERS, + ) + for session_id, raw in rows: + try: + entry = json.loads(raw) + except (TypeError, ValueError): + continue + if not isinstance(entry, dict) or str(entry.get("providerID") or "") not in PROVIDERS: + continue + tokens = entry.get("tokens") if isinstance(entry.get("tokens"), dict) else {} + cache = tokens.get("cache") if isinstance(tokens.get("cache"), dict) else {} + inp = token_count(tokens.get("input")) + # opencode keeps thinking tokens out of output; both are generated. + out = token_count(tokens.get("output")) + token_count(tokens.get("reasoning")) + cread = token_count(cache.get("read")) + cwrite = token_count(cache.get("write")) + if not (inp or out or cread or cwrite): + continue + created = entry.get("time").get("created") if isinstance(entry.get("time"), dict) else None + day = local_day(created) + if day is None: + continue + session = "opencode:" + str(session_id) + model = str(entry.get("modelID") or "").rstrip("/").split("/")[-1] or "minimax" + stats.prompt(session, day) + stats.tokens(session, day, model, inp, out, cread, cwrite) + except sqlite3.Error: + return False + finally: + conn.close() + return True + + +def scan_sessions() -> tuple[dict[str, Any], bool]: + stats = Stats() + # Every source runs, so one unreadable source still leaves the others. + complete = [scan_mcode(stats), scan_pi(stats), scan_opencode(stats)] + return stats.record(), all(complete) + + +def cached_local_stats(max_age_seconds: float) -> dict[str, Any]: + """Local stats, with the cache as a pure optimization.""" + sources = [mcode_data_dir(), opencode_data_dir(), *pi_session_roots()] + digest = hashlib.sha1("\0".join(str(path) for path in sources).encode("utf-8")).hexdigest()[:16] + try: + root = cache_root() + except OSError: + return scan_sessions()[0] + cache_file, lock_file = root / f"minimax-scan-{digest}.json", root / f"minimax-scan-{digest}.lock" + + # A scan only answers for the day it ran on: its today* counts would be + # yesterday's after midnight. + def reusable() -> dict[str, Any] | None: + cached = read_json(cache_file) + if not isinstance(cached, dict) or cached.get("scanDate") != datetime.now().strftime("%Y-%m-%d"): + return None + scanned_ms = number(cached.get("scannedAtMs")) or 0 + age = time.time() - scanned_ms / 1000 + stats = cached.get("stats") + return stats if 0 <= age < max_age_seconds and isinstance(stats, dict) else None + + stats = reusable() + if stats is not None: + return stats + with probe_lock(lock_file): + stats = reusable() + if stats is not None: + return stats + stats, complete = scan_sessions() + # A scan that missed a source would stand in for a real one until the + # next rescan; only a complete one is kept. + if complete: + try: + write_json(cache_file, {"scanDate": datetime.now().strftime("%Y-%m-%d"), "scannedAtMs": round(time.time() * 1000), "stats": stats}) + except OSError as exc: + print(f"omarchy-agent-usage-minimax: could not write scan cache ({exc})", file=sys.stderr) + return stats + + +# -------------------------------------------------------------------------- +# Record +# -------------------------------------------------------------------------- + +# Nobody using MiniMax here: an empty record, which the panel skips. With a +# key, a failed check keeps the last good limits, marked stale; without one, +# local usage still shows and the limits wait for a key. +def collect(api_key: str, base_url: str, force: bool = False, stats: dict[str, Any] | None = None) -> dict[str, Any]: + record = base_record() + if stats: + record.update(stats) + has_stats = bool(record["modelUsage"]) or record["totalPrompts"] > 0 + record["hasLocalStats"] = has_stats + if not api_key: - return base_record() + if has_stats: + record.update(ready=True, usageStatusText="Waiting for auth", authHelpText=AUTH_HELP) + return record limits = collect_limits(api_key, base_url, force) - record = base_record( + record.update( ready=True, tierLabel=TIER_LABEL if limits["limits"] else "", limits=limits["limits"], @@ -442,17 +862,26 @@ def collect(api_key: str, base_url: str, force: bool = False) -> dict[str, Any]: def main() -> int: - parser = argparse.ArgumentParser(description="Print the MiniMax Token Plan usage record as JSON") - # --force re-probes past the reuse window. --limits-only takes the same path - # as a normal run: there is no local scan for it to reuse. + parser = argparse.ArgumentParser(description="Print the MiniMax 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") parser.add_argument("--base-url") args = parser.parse_args() + try: + scan_age = 0 if args.force else (LIMITS_ONLY_REUSE_SECONDS if args.limits_only else SCAN_REUSE_SECONDS) + stats = cached_local_stats(scan_age) + except Exception as error: + # Local stats are an addition; their failure must not cost the limits. + print(f"omarchy-agent-usage-minimax: local scan failed ({type(error).__name__})", file=sys.stderr) + stats = None + try: api_key, saved_base_url = credentials() - record = collect(api_key, args.base_url or saved_base_url, args.force) + record = collect(api_key, args.base_url or saved_base_url, args.force, stats) 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/test/shell.d/agent-usage-minimax-scanner-test.sh b/test/shell.d/agent-usage-minimax-scanner-test.sh new file mode 100644 index 00000000000..2a71d8d50df --- /dev/null +++ b/test/shell.d/agent-usage-minimax-scanner-test.sh @@ -0,0 +1,109 @@ +#!/bin/bash + +source "$(dirname "$0")/base-test.sh" + +require_command jq +require_command python3 + +TEST_HOME=$(mktemp -d) +trap 'rm -rf "$TEST_HOME"' EXIT + +today="$(date +%Y-%m-%d)" +timestamp="${today}T12:00:00Z" +now_ms=$(python3 -c 'import time; print(round(time.time() * 1000))') + +# MiniMax Code's own sessions: user messages are prompts, assistant usage is +# tokens. A row with a garbage timestamp or token count drops on its own. +mkdir -p "$TEST_HOME/.minimax/v2/sessions/s1" "$TEST_HOME/.minimax/v2/sessions/s2" +cat >"$TEST_HOME/.minimax/v2/sessions/s1/messages.jsonl" <"$TEST_HOME/.minimax/v2/sessions/s2/messages.jsonl" <"$TEST_HOME/.pi/agent/sessions/p/a.jsonl" <"$TEST_HOME/.omp/profiles/work/agent/sessions/p/b.jsonl" <"$TEST_HOME/moved/v2/sessions/s/messages.jsonl" +moved=$(run env MINIMAX_DATA_DIR="$TEST_HOME/moved" MAVIS_DATA_DIR="$TEST_HOME/.minimax" "$ROOT/bin/omarchy-agent-usage-minimax" --force) +[[ $(jq -r '.todayTotalTokens' <<<"$moved") == "1165" ]] || + fail "MiniMax collector reads MiniMax Code's sessions from MINIMAX_DATA_DIR" "$moved" +pass "MiniMax collector reads MiniMax Code's sessions from MINIMAX_DATA_DIR" + +# --limits-only reuses today's scan; --force rescans. +printf '{"message":{"role":"assistant","timestamp":%s,"model":"MiniMax-M2","usage":{"input":1}}}\n' "$now_ms" >>"$TEST_HOME/.minimax/v2/sessions/s1/messages.jsonl" +reused=$(run "$ROOT/bin/omarchy-agent-usage-minimax" --limits-only) +forced=$(run "$ROOT/bin/omarchy-agent-usage-minimax" --force) +[[ $(jq -r '.todayTotalTokens' <<<"$reused") == "1180" && $(jq -r '.todayTotalTokens' <<<"$forced") == "1181" ]] || + fail "MiniMax collector reuses a recent scan for --limits-only and rescans on --force" "$reused $forced" +pass "MiniMax collector reuses a recent scan for --limits-only and rescans on --force" + +# A scan that couldn't read a source isn't kept for the next run. +if (( EUID != 0 )); then + rm -f "$TEST_HOME/.cache/omarchy/agent-usage/"minimax-scan-*.json + chmod 000 "$TEST_HOME/.minimax/v2/sessions/s2/messages.jsonl" + run "$ROOT/bin/omarchy-agent-usage-minimax" --force >/dev/null + chmod 644 "$TEST_HOME/.minimax/v2/sessions/s2/messages.jsonl" + partial=$(run "$ROOT/bin/omarchy-agent-usage-minimax" --limits-only) + [[ $(jq -r '.totalPrompts' <<<"$partial") == "5" ]] || + fail "MiniMax collector does not cache a scan that missed a file" "$partial" + pass "MiniMax collector does not cache a scan that missed a file" +else + skip "MiniMax collector does not cache a scan that missed a file (root reads anything)" +fi diff --git a/test/shell.d/agent-usage-minimax-test.sh b/test/shell.d/agent-usage-minimax-test.sh index 2b9a46a2b4d..3c521c08a9e 100755 --- a/test/shell.d/agent-usage-minimax-test.sh +++ b/test/shell.d/agent-usage-minimax-test.sh @@ -91,6 +91,14 @@ opencode.parent.mkdir(parents=True) opencode.write_text(json.dumps({"minimax": {"type": "api", "key": "mmx_opencode"}})) opencode_key, opencode_base_url = scanner.credentials() +# A China-region key pi saved answers on the China host. +opencode.unlink() +pi_auth = home / ".pi/agent/auth.json" +pi_auth.parent.mkdir(parents=True) +pi_auth.write_text(json.dumps({"minimax-cn": {"type": "api_key", "key": "mmx_pi_cn"}})) +pi_key, pi_base_url = scanner.credentials() +pi_auth.unlink() + import copy, time from datetime import datetime, timedelta, timezone @@ -117,7 +125,8 @@ def run(fetch, key="mmx_test", force=False): return dict(record, probes=len(probes)) out = {"limits": limits, "mediaFirst": media_first, "savedKey": saved_key, "savedBaseUrl": saved_base_url, - "oauthKey": oauth_key, "oauthBaseUrl": oauth_base_url, "opencodeKey": opencode_key} + "oauthKey": oauth_key, "oauthBaseUrl": oauth_base_url, "opencodeKey": opencode_key, + "piKey": pi_key, "piBaseUrl": pi_base_url} out["record"] = run(answer, force=True) # Within the probe interval a repeat run reuses the answer; --force doesn't. @@ -193,6 +202,10 @@ pass "MiniMax collector skips media quota rows when choosing the plan's row" fail "MiniMax credentials use mmx and opencode logins" "$result" pass "MiniMax credentials use mmx and opencode logins" +[[ $(jq -r '.piKey + ":" + .piBaseUrl' <<<"$result") == "mmx_pi_cn:https://api.minimaxi.com" ]] || + fail "MiniMax credentials use a China-region key pi saved, on the China host" "$result" +pass "MiniMax credentials use a China-region key pi saved, on the China host" + [[ $(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" From 79fb091c1ea025555493fbaf4fee5231cf6b64be Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:14:07 -0400 Subject: [PATCH 081/546] Give MiniMax its mark in the agents panel Co-authored-by: Josh LeBlanc Co-Authored-By: Claude Opus 5.5 --- shell/plugins/agents/assets/minimax-light.svg | 1 + shell/plugins/agents/assets/minimax.svg | 1 + 2 files changed, 2 insertions(+) create mode 100644 shell/plugins/agents/assets/minimax-light.svg create mode 100644 shell/plugins/agents/assets/minimax.svg diff --git a/shell/plugins/agents/assets/minimax-light.svg b/shell/plugins/agents/assets/minimax-light.svg new file mode 100644 index 00000000000..290d7c2411f --- /dev/null +++ b/shell/plugins/agents/assets/minimax-light.svg @@ -0,0 +1 @@ + diff --git a/shell/plugins/agents/assets/minimax.svg b/shell/plugins/agents/assets/minimax.svg new file mode 100644 index 00000000000..74a90280532 --- /dev/null +++ b/shell/plugins/agents/assets/minimax.svg @@ -0,0 +1 @@ + From 87420c794ea281515cdc7f0fd744b757252f6aa3 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:01:27 -0400 Subject: [PATCH 082/546] Bring OpenCode's limits in line with the collector contract A machine without a key shows local stats with no status, a refused or lapsed credential reads as Sign-in expired, and a key with no Go plan behind it stays silent and isn't asked again for an hour. The limits cache is read whatever its file time, a timeout advises a retry while a garbled answer doesn't, and a signed-in machine is ready. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-opencode | 95 +++++++++++++----- .../agent-usage-opencode-limits-test.sh | 99 ++++++++++++++++--- .../agent-usage-opencode-scanner-test.sh | 2 +- 3 files changed, 154 insertions(+), 42 deletions(-) diff --git a/bin/omarchy-agent-usage-opencode b/bin/omarchy-agent-usage-opencode index e668da83e3a..0b8b9ec80e3 100755 --- a/bin/omarchy-agent-usage-opencode +++ b/bin/omarchy-agent-usage-opencode @@ -40,12 +40,16 @@ API_BASE_URL = "https://opencode.ai" USAGE_PATH = "/zen/go/v1/usage" CONSOLE_USAGE_PATH = "/inference/go/v1/usage" CONSOLE_ORG_HEADER = "x-opencode-org-id" -AUTH_HELP = "Run `opencode providers login`, set OPENCODE_API_KEY, or sign in with pi to read subscription limits." CONSOLE_EXPIRED_HELP = "OpenCode's Console session has expired. Run `opencode auth login` to refresh it." # A panel opened and closed repeatedly must not become one request per flick. PROBE_MIN_INTERVAL_SECONDS = 15 +# A key with no Go plan behind it (a Zen pay-as-you-go account) gets the same +# answer every time, so it is asked again hourly rather than on every run; +# --force asks at once, for someone who just subscribed. +NO_PLAN_RECHECK_SECONDS = 3600 + # The subscription has two provider ids: "opencode" for the plain Zen # gateway and "opencode-go" for the Go plan. Both spend the same account and # both appear in pi sessions and in opencode's own database. @@ -660,21 +664,27 @@ def probe_limits(credential): }) try: with urllib.request.urlopen(request, timeout=10) as response: - payload = json.loads(response.read().decode("utf-8", errors="replace")) + body = response.read() except urllib.error.HTTPError as error: # A refused credential may just be the wrong one for this account, so the # next way in still gets its turn. if error.code == 401: return {"ok": False, "refused": True, "helpText": "OpenCode rejected the saved key. Reconnect OpenCode Go with /connect in the TUI, or set OPENCODE_API_KEY."} if error.code == 403: - return {"ok": False, "refused": True, "helpText": "The saved key has no OpenCode Go subscription behind it. Local OpenCode stats are still shown."} + # A key without a Go plan behind it: Zen pay-as-you-go has no meters, + # which is no trouble at all. + return {"ok": False, "noPlan": True} return {"ok": False, "helpText": f"OpenCode's usage endpoint returned status {error.code}. Local OpenCode 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. + # A transport failure reached no server at all: no route, no DNS, a + # timeout. 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 OpenCode's usage endpoint. Retrying shortly. Local OpenCode stats are still shown."} + try: + payload = json.loads(body.decode("utf-8", errors="replace")) + except ValueError: + payload = None if not isinstance(payload, dict) or not payload: return {"ok": False, "helpText": "OpenCode's usage endpoint returned no data. Local OpenCode stats are still shown."} limits = parse_usage(payload) @@ -699,6 +709,14 @@ def limit_window_open(entry, now): return resets_at > now +def read_json(path): + try: + value = json.loads(path.read_text(encoding="utf-8")) + except Exception: + return {} + return value if isinstance(value, dict) else {} + + def open_cached_limits(cached): entries = cached.get("limits") if isinstance(cached, dict) else None if not isinstance(entries, list): @@ -717,38 +735,57 @@ def limits_cache_path(token): def collect_limits(force): - result = {"limits": [], "tierLabel": "", "usageStatusText": "Local usage only", "authHelpText": "", "retryAdvised": False, "stale": False, "fetchedAtMs": 0} + result = {"limits": [], "tierLabel": "", "usageStatusText": "", "authHelpText": "", "retryAdvised": False, "stale": False, "fetchedAtMs": 0, "signedIn": False} + # No key at all is a Zen user paying as they go, or someone who only runs + # opencode locally: local stats stand alone, with nothing to fix. candidates = credential_candidates() if not candidates: - result["authHelpText"] = AUTH_HELP return result + result["signedIn"] = True # Limits only ever show for a Go plan, so a tier follows them. def answer(limits, fetched_at, stale): - result.update(limits=limits, tierLabel="Go" if limits else "", usageStatusText="", stale=stale, fetchedAtMs=fetched_at if limits else 0) + result.update(limits=limits, tierLabel="Go" if limits else "", stale=stale, fetchedAtMs=fetched_at if limits else 0) return result + def remember(cache, payload): + try: + write_json(cache, payload) + except Exception as exc: + print(f"omarchy-agent-usage-opencode: could not write limits cache ({exc})", file=sys.stderr) + refused = None + no_plan = False for candidate in candidates: cache = limits_cache_path(candidate.token) - cached = read_fresh_json(cache, float("inf")) or {} + # Read whatever the file time says, so a clock that went back keeps it. + cached = read_json(cache) fallback = open_cached_limits(cached) fetched_at = number(cached.get("fetchedAtMs")) + age = time.time() - fetched_at / 1000 + + # A key already known to have no Go plan isn't asked again for a while. + if cached.get("noPlan") is True and not force and 0 <= age < NO_PLAN_RECHECK_SECONDS: + no_plan = True + continue # A panel opened and closed repeatedly must not become one request per # flick, so a recent answer is reused; --force skips the wait. - if fallback and not force and 0 <= time.time() - fetched_at / 1000 < PROBE_MIN_INTERVAL_SECONDS: + if fallback and not force and 0 <= age < PROBE_MIN_INTERVAL_SECONDS: return answer(fallback, fetched_at, False) probe = probe_limits(candidate) if probe.get("ok"): fetched_at = round(time.time() * 1000) - try: - write_json(cache, {"fetchedAtMs": fetched_at, "limits": probe["limits"]}) - except Exception: - pass + remember(cache, {"fetchedAtMs": fetched_at, "limits": probe["limits"]}) return answer(probe["limits"], fetched_at, False) + if probe.get("noPlan"): + # The windows of a plan that's gone mean nothing now, so they go too. + remember(cache, {"fetchedAtMs": round(time.time() * 1000), "limits": [], "noPlan": True}) + no_plan = True + continue + if probe.get("refused"): # Keep the first refusal that still has open windows, so a later # credential without a cache can't blank them. A lapsed Console @@ -761,24 +798,28 @@ def collect_limits(force): continue # Last good numbers, marked stale, beat an empty limits section while - # the probe is down. + # the probe is down. A key that already said it has no plan makes an + # outage on another one no news. result["retryAdvised"] = bool(probe.get("transport")) answer(fallback, fetched_at, True) - if not fallback: + if fallback: + result["authHelpText"] = probe["helpText"] + elif not no_plan: result["usageStatusText"] = "OpenCode limits unavailable" result["authHelpText"] = probe["helpText"] return result + # A working key without a Go plan is the answer: the other ways in only + # failed to say the same. + if no_plan: + return result + # Every way in was refused: an account problem, not an outage. Keep any - # still-open numbers, but say what went wrong. + # still-open numbers, but say a new sign-in is needed. probe, fallback, fetched_at, expired = refused answer(fallback, fetched_at, True) - if expired: - result["usageStatusText"] = "OpenCode Console session expired" - result["authHelpText"] = CONSOLE_EXPIRED_HELP - else: - result["usageStatusText"] = "Sign-in rejected" if fallback else "OpenCode limits unavailable" - result["authHelpText"] = probe["helpText"] + result["usageStatusText"] = "Sign-in expired" + result["authHelpText"] = CONSOLE_EXPIRED_HELP if expired else probe["helpText"] return result @@ -797,14 +838,14 @@ def main(): stats = cached_local_stats(max_age) limits = collect_limits(args.force) - # Nobody uses OpenCode's subscription here: no usage and no limits leave - # an empty tier, a record the panel skips. + # Nobody uses OpenCode here: no usage and no credential leave a record the + # panel skips. record = { "schemaVersion": 1, "id": AGENT_ID, "name": AGENT_NAME, "updatedAt": datetime.now(timezone.utc).isoformat(), - "ready": number(stats.get("totalPrompts")) > 0 or bool(limits["limits"]), + "ready": number(stats.get("totalPrompts")) > 0 or bool(limits["limits"]) or limits["signedIn"], "hasLocalStats": True, "tierLabel": limits["tierLabel"], "usageStatusText": limits["usageStatusText"], diff --git a/test/shell.d/agent-usage-opencode-limits-test.sh b/test/shell.d/agent-usage-opencode-limits-test.sh index b861b20ae5f..656583ba3be 100644 --- a/test/shell.d/agent-usage-opencode-limits-test.sh +++ b/test/shell.d/agent-usage-opencode-limits-test.sh @@ -19,7 +19,9 @@ import io import json import os import pathlib +import sqlite3 import sys +import time import urllib.error loader = importlib.machinery.SourceFileLoader("collector", os.environ["COLLECTOR"]) @@ -111,17 +113,45 @@ clear_cache() collector.urllib.request.urlopen = answer(error=urllib.error.HTTPError("https://x", 401, "Unauthorized", {}, None)) result = collector.collect_limits(True) check( - "OpenCode collector asks for a fresh key on 401", - result["limits"] == [] and result["usageStatusText"] == "OpenCode limits unavailable" and "rejected the saved key" in result["authHelpText"], + "OpenCode collector asks for a new sign-in on 401", + result["limits"] == [] and result["usageStatusText"] == "Sign-in expired" and "rejected the saved key" in result["authHelpText"] + and result["retryAdvised"] is False, json.dumps(result), ) +# A Zen pay-as-you-go key has no Go plan behind it. That is no trouble, and +# the answer is remembered so the next runs don't ask again. clear_cache() -collector.urllib.request.urlopen = answer(error=urllib.error.HTTPError("https://x", 403, "Forbidden", {}, None)) +asked = [] +def no_plan(request, timeout=None): + asked.append(request.full_url) + raise urllib.error.HTTPError(request.full_url, 403, "Forbidden", {}, None) +collector.urllib.request.urlopen = no_plan +result = collector.collect_limits(True) +check( + "OpenCode collector stays silent for a key without a Go plan", + result["limits"] == [] and result["usageStatusText"] == "" and result["authHelpText"] == "" and result["tierLabel"] == "" + and result["stale"] is False and len(asked) == 1, + json.dumps(result), +) +result = collector.collect_limits(False) +check( + "OpenCode collector doesn't ask again about a key without a Go plan", + result["usageStatusText"] == "" and result["limits"] == [] and len(asked) == 1, + json.dumps(result), +) +collector.urllib.request.urlopen = answer({"usage": {"weekly": {"percent": 10, "resetsAt": "2999-01-01T00:00:00.000Z"}}}) result = collector.collect_limits(True) check( - "OpenCode collector names a missing Go subscription on 403", - "no OpenCode Go subscription" in result["authHelpText"], + "OpenCode collector --force asks again after a Go plan appears", + [(w["label"], w["percent"]) for w in result["limits"]] == [("Weekly (7-day)", 0.1)] and result["tierLabel"] == "Go", + json.dumps(result), +) +collector.urllib.request.urlopen = no_plan +result = collector.collect_limits(True) +check( + "OpenCode collector drops a cancelled plan's windows", + result["limits"] == [] and result["usageStatusText"] == "" and result["tierLabel"] == "", json.dumps(result), ) @@ -129,8 +159,9 @@ clear_cache() collector.urllib.request.urlopen = answer(error=urllib.error.HTTPError("https://x", 500, "Server Error", {}, None)) result = collector.collect_limits(True) check( - "OpenCode collector shows a non-auth status code", - "returned status 500" in result["authHelpText"], + "OpenCode collector shows a non-auth status code without advising a retry", + "returned status 500" in result["authHelpText"] and result["usageStatusText"] == "OpenCode limits unavailable" + and result["retryAdvised"] is False, json.dumps(result), ) @@ -143,12 +174,31 @@ check( json.dumps(result), ) +clear_cache() +collector.urllib.request.urlopen = answer(error=TimeoutError("timed out")) +result = collector.collect_limits(True) +check( + "OpenCode collector advises a retry on a timeout", + result["retryAdvised"] is True and "Couldn't reach" in result["authHelpText"], + json.dumps(result), +) + +# A server that answered with garbage was reached; retrying soon won't help. +clear_cache() +collector.urllib.request.urlopen = lambda request, timeout=None: io.BytesIO(b"oops") +result = collector.collect_limits(True) +check( + "OpenCode collector reads an unparseable answer as no data, not an outage", + result["retryAdvised"] is False and "returned no data" in result["authHelpText"], + json.dumps(result), +) + clear_key() -collector.urllib.request.urlopen = answer({"usage": {"rolling": {"percent": 1}}}) +collector.urllib.request.urlopen = answer(error=RuntimeError("no key, no request")) result = collector.collect_limits(True) check( - "OpenCode collector falls back to local stats without a key", - result["limits"] == [] and result["usageStatusText"] == "Local usage only" and result["authHelpText"] != "", + "OpenCode collector shows local stats alone without a key", + result["limits"] == [] and result["usageStatusText"] == "" and result["authHelpText"] == "" and result["signedIn"] is False, json.dumps(result), ) @@ -169,7 +219,30 @@ check( "OpenCode collector keeps cached windows, stale, after a failed probe", [(w["label"], w["percent"], w["resetsAt"]) for w in result["limits"]] == [("Rolling (5h)", 0.0, ""), ("Monthly", 0.4, "2999-01-01T00:00:00.000Z")] and result["stale"] is True - and result["fetchedAtMs"] == 1700000000000 and result["tierLabel"] == "Go" and result["usageStatusText"] == "", + and result["fetchedAtMs"] == 1700000000000 and result["tierLabel"] == "Go" and result["usageStatusText"] == "" + and result["retryAdvised"] is False, + json.dumps(result), +) + +# A cache written before the clock went back is still the last good answer. +future = time.time() + 86400 +os.utime(collector.limits_cache_path("zen_test"), (future, future)) +collector.urllib.request.urlopen = answer(error=urllib.error.URLError("no route")) +result = collector.collect_limits(True) +check( + "OpenCode collector keeps a cache whose file time is in the future", + [w["label"] for w in result["limits"]] == ["Rolling (5h)", "Monthly"] and result["stale"] is True and result["retryAdvised"] is True + and result["usageStatusText"] == "", + json.dumps(result), +) + +# A refused key keeps its last numbers, stale, and asks for a new sign-in. +collector.urllib.request.urlopen = answer(error=urllib.error.HTTPError("https://x", 401, "Unauthorized", {}, None)) +result = collector.collect_limits(True) +check( + "OpenCode collector keeps the last limits when the key is refused", + [w["label"] for w in result["limits"]] == ["Rolling (5h)", "Monthly"] and result["stale"] is True + and result["usageStatusText"] == "Sign-in expired", json.dumps(result), ) @@ -227,8 +300,6 @@ check("OpenCode collector reads pi's credential as a last resort", [c.token for # OpenCode V2 keeps its credentials in the database: a Go key first, and the # Console session last, asked on Console's endpoint with its workspace. -import sqlite3 -import time clear_key() v2_home = os.path.join(home, "v2-data") os.makedirs(os.path.join(v2_home, "opencode"), exist_ok=True) @@ -275,7 +346,7 @@ collector.urllib.request.urlopen = answer(error=urllib.error.HTTPError("https:// result = collector.collect_limits(True) check( "OpenCode collector names an expired Console session", - result["usageStatusText"] == "OpenCode Console session expired" and result["limits"] == [], + result["usageStatusText"] == "Sign-in expired" and "Console session has expired" in result["authHelpText"] and result["limits"] == [], json.dumps(result), ) PY diff --git a/test/shell.d/agent-usage-opencode-scanner-test.sh b/test/shell.d/agent-usage-opencode-scanner-test.sh index abcf48b38d4..bab2be2718e 100644 --- a/test/shell.d/agent-usage-opencode-scanner-test.sh +++ b/test/shell.d/agent-usage-opencode-scanner-test.sh @@ -269,6 +269,6 @@ EMPTY_HOME=$(mktemp -d) trap 'rm -rf "$PI_HOME" "$NOID_HOME" "$OPENCODE_HOME" "$CACHE_HOME" "$SESSION_HOME" "$V2_HOME" "$ODD_HOME" "$EMPTY_HOME"' EXIT result=$(HOME="$EMPTY_HOME" XDG_CACHE_HOME="$EMPTY_HOME/.cache" XDG_DATA_HOME="$EMPTY_HOME/.local/share" \ OPENCODE_API_KEY= PATH="$EMPTY_HOME/bin:$PATH" "$ROOT/bin/omarchy-agent-usage-opencode") -[[ $(jq -c '{ready, tierLabel, limits, totalPrompts, limitsStale}' <<<"$result") == '{"ready":false,"tierLabel":"","limits":[],"totalPrompts":0,"limitsStale":false}' ]] || +[[ $(jq -c '{ready, tierLabel, limits, totalPrompts, limitsStale, limitsFetchedAt, usageStatusText, authHelpText, recentDays: (.recentDays | length), activeDates, modelUsage}' <<<"$result") == '{"ready":false,"tierLabel":"","limits":[],"totalPrompts":0,"limitsStale":false,"limitsFetchedAt":0,"usageStatusText":"","authHelpText":"","recentDays":7,"activeDates":[],"modelUsage":{}}' ]] || fail "OpenCode collector gives an empty record where nobody uses it" "$result" pass "OpenCode collector gives an empty record where nobody uses it" From a70ec7a9c57f5e5d4e69a90f84d3d17261e57048 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:01:44 -0400 Subject: [PATCH 083/546] Read every pi session root and untrusted fields safely in OpenCode pi and omp profiles and PI_CODING_AGENT_DIR are scanned with rg --no-ignore, a fork's copy of a message counts once by id and timestamp, a usage block with only a total still counts, and either spelling of the token fields reads. A bad timestamp or token count skips its row or field, and a scan that couldn't read the sessions isn't cached. OPENCODE_GO_API_KEY and a plain-string auth.json key are read too. Co-authored-by: Alvaro Minarro Gil Co-authored-by: ivanvan08 Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-opencode | 213 ++++++++++++------ .../agent-usage-opencode-limits-test.sh | 11 +- .../agent-usage-opencode-scanner-test.sh | 63 +++++- 3 files changed, 218 insertions(+), 69 deletions(-) diff --git a/bin/omarchy-agent-usage-opencode b/bin/omarchy-agent-usage-opencode index 0b8b9ec80e3..a8e938812b5 100755 --- a/bin/omarchy-agent-usage-opencode +++ b/bin/omarchy-agent-usage-opencode @@ -16,6 +16,7 @@ import argparse import fcntl import hashlib import json +import math import os import shutil import sqlite3 @@ -70,20 +71,44 @@ def obj(value): return value if isinstance(value, dict) else {} +# Counts and times read from disk or the network are untrusted: a boolean, +# NaN, negative, or absurdly large value reads as nothing, never as a number +# that would swamp the totals. The ceiling sits well above any epoch +# millisecond time. +NUMBER_CEILING = 10 ** 15 + + def number(value): + if isinstance(value, bool): + return 0 try: - return int(value or 0) - except Exception: + n = float(value or 0) + except (TypeError, ValueError, OverflowError): + return 0 + if not math.isfinite(n) or n < 0 or n >= NUMBER_CEILING: return 0 + return round(n) + + +# pi, omp, and the harnesses that copy their format spell usage camelCase or +# snake_case; the first spelling present wins. +def usage_tokens(usage, *keys): + for key in keys: + if key in usage: + return number(usage.get(key)) + return 0 # Usage percentages arrive as 0-100; the record contract wants 0-1. An # unreadable value returns -1 so the caller can drop the window. def fraction(value): + if isinstance(value, bool): + return -1.0 try: - return float(value) / 100.0 + percent = float(value) / 100.0 except (TypeError, ValueError): return -1.0 + return percent if math.isfinite(percent) else -1.0 def model_name(raw): @@ -119,25 +144,35 @@ def find_command(name): return shutil.which(name, path=ENV.get("PATH")) +# The local day a message was spent on, or None when its timestamp can't be +# trusted, so the caller skips that row rather than the whole scan. A message +# with no timestamp at all is today's, as it always was. def local_day(value): if value is None: return datetime.now().strftime("%Y-%m-%d") + if isinstance(value, str) and value.strip().isdigit(): + value = int(value.strip()) + if isinstance(value, bool): + return None if isinstance(value, (int, float)): + seconds = number(value) + if seconds <= 0: + return None # pi message timestamps are milliseconds; opencode writes milliseconds too. - if value > 10_000_000_000: - value = value / 1000 - return datetime.fromtimestamp(value).strftime("%Y-%m-%d") - text = str(value) + if seconds > 10_000_000_000: + seconds = seconds / 1000 + try: + return datetime.fromtimestamp(seconds).strftime("%Y-%m-%d") + except (OverflowError, OSError, ValueError): + return None + text = str(value).strip() 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") + dt = datetime.fromisoformat(text.replace("Z", "+00:00")) + except ValueError: + return None + if dt.tzinfo is not None: + dt = dt.astimezone() + return dt.strftime("%Y-%m-%d") now = datetime.now() @@ -184,29 +219,52 @@ def add_usage(day, session_key, model, input_tokens, output_tokens, cache_read, today_tokens_by_model[model] = today_tokens_by_model.get(model, 0) + total +# pi's and omp's session trees, each profile's too: `--profile=` moves +# the whole agent tree under /profiles//, and PI_CODING_AGENT_DIR +# moves pi's. +def pi_session_roots(): + roots = [] + if os.environ.get("PI_CODING_AGENT_DIR"): + roots.append(Path(os.environ["PI_CODING_AGENT_DIR"]).expanduser() / "sessions") + for base in (Path.home() / ".pi", Path.home() / ".omp"): + roots.append(base / "agent" / "sessions") + try: + # Sorted so the scan order does not depend on directory order. + roots.extend(sorted(child / "agent" / "sessions" for child in (base / "profiles").iterdir() if child.is_dir())) + except OSError: + # No profiles directory, or it is unreadable: the default root stands. + pass + return list(dict.fromkeys(roots)) + + def scan_pi_sessions(): # pi and omp can spend the subscription without opencode's CLI ever # running, so its own database cannot be the only source. Their JSONL # transcripts carry the provider, model, and token usage per message. - roots = [ - Path.home() / ".pi" / "agent" / "sessions", - Path.home() / ".omp" / "agent" / "sessions", - ] - rg = find_command("rg") or "rg" + # Returns whether every root was read in full, so a partial scan is never + # cached as the whole story. + roots = [root for root in pi_session_roots() if root.is_dir()] + if not roots: + return True + rg = find_command("rg") + if not rg: + print("omarchy-agent-usage-opencode: rg is missing; pi and omp sessions were not scanned", file=sys.stderr) + return False + complete = True for root in roots: - if not root.exists(): - continue try: proc = subprocess.Popen( - [rg, "--json", "-e", r'"provider"\s*:\s*"opencode(-go)?"', str(root)], + # --no-ignore: a home that is itself a git checkout would otherwise + # hide every session file behind its ignore rules. + [rg, "--json", "--no-ignore", "-e", r'"provider"\s*:\s*"opencode(-go)?"', str(root)], stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, text=True, errors="replace", env=ENV, ) - except FileNotFoundError: - return + except OSError: + return False assert proc.stdout is not None for raw in proc.stdout: @@ -219,43 +277,60 @@ def scan_pi_sessions(): entry = json.loads(line) except Exception: continue - - if not isinstance(entry, dict) or entry.get("type") != "message": - continue - # Only an id can say two lines are the same message; without one, each - # line is its own. - message_id = str(entry.get("id") or "") - if message_id: - message_key = path + ":" + message_id - if message_key in seen_pi_messages: - continue - seen_pi_messages.add(message_key) - message = obj(entry.get("message")) - if message.get("role") != "assistant": - continue - if str(message.get("provider") or "") not in PROVIDERS: - continue - - usage = obj(message.get("usage")) - if not usage: - continue - input_tokens = number(usage.get("input")) - # pi folds thinking into its own field; opencode counts it as generated. - output_tokens = number(usage.get("output")) + number(usage.get("reasoning")) - cache_read = number(usage.get("cacheRead")) - cache_write = number(usage.get("cacheWrite")) - if not (input_tokens or output_tokens or cache_read or cache_write): + try: + scan_pi_entry(entry, path) + except Exception: + # One odd line costs that line, never the scan. continue - day = local_day(entry.get("timestamp") or message.get("timestamp")) - # A model id may carry a thinking variant suffix ("deepseek-v4.1-flash:high"). - model = model_name(message.get("model")) - add_usage(day, path, model, input_tokens, output_tokens, cache_read, cache_write) - try: - proc.wait(timeout=1) + # rg exits 1 for no match and 2 when it couldn't read something. + if proc.wait(timeout=1) not in (0, 1): + complete = False except Exception: proc.kill() + complete = False + return complete + + +def scan_pi_entry(entry, path): + if not isinstance(entry, dict) or entry.get("type") != "message": + return + message = obj(entry.get("message")) + if message.get("role") != "assistant": + return + if str(message.get("provider") or "") not in PROVIDERS: + return + # A fork copies its parent's entries with their id and timestamp, so the + # pair names a message across files; an id alone only within its file. + # Without an id, each line is its own. + message_id = entry.get("id") + if message_id: + timestamp = entry.get("timestamp") + message_key = (str(message_id), str(timestamp)) if timestamp else (path, str(message_id)) + if message_key in seen_pi_messages: + return + seen_pi_messages.add(message_key) + + usage = obj(message.get("usage")) + if not usage: + return + input_tokens = usage_tokens(usage, "input", "inputTokens", "input_tokens") + # pi folds thinking into its own field; opencode counts it as generated. + output_tokens = usage_tokens(usage, "output", "outputTokens", "output_tokens") + usage_tokens(usage, "reasoning", "reasoningTokens", "reasoning_tokens") + cache_read = usage_tokens(usage, "cacheRead", "cacheReadTokens", "cache_read", "cache_read_input_tokens") + cache_write = usage_tokens(usage, "cacheWrite", "cacheWriteTokens", "cache_write", "cache_creation_input_tokens") + if not (input_tokens or output_tokens or cache_read or cache_write): + # Some writers keep only the total; it still spent the subscription. + input_tokens = usage_tokens(usage, "totalTokens", "total_tokens") + if not input_tokens: + return + + day = local_day(entry.get("timestamp") or message.get("timestamp")) + if day is None: + return + # A model id may carry a thinking variant suffix ("deepseek-v4.1-flash:high"). + add_usage(day, path, model_name(message.get("model")), input_tokens, output_tokens, cache_read, cache_write) def data_home(): @@ -314,6 +389,8 @@ def store_row(message_id, session_id, entry, row_type, created, seen): return seen.add(message_id) day = local_day(obj(entry.get("time")).get("created") or created) + if day is None: + return add_usage(day, "opencode:" + str(session_id), model_name(entry.get("modelID") or model), input_tokens, output_tokens, cache_read, cache_write) @@ -406,9 +483,10 @@ def cache_root(): def scan_cache_paths(): db = opencode_db() - # The digest covers every data path the scan reads: the opencode DB and, - # via Path.home(), the pi/omp session roots. - digest = hashlib.sha1((str(Path.home()) + "\n" + str(db)).encode("utf-8")).hexdigest()[:16] + # The digest covers every data path the scan reads: the opencode DB and + # the pi/omp session roots. + sources = [str(db)] + [str(root) for root in pi_session_roots()] + digest = hashlib.sha1("\n".join(sources).encode("utf-8")).hexdigest()[:16] root = cache_root() return root / f"opencode-scan-{digest}.json", root / f"opencode-scan-{digest}.lock" @@ -493,8 +571,8 @@ def local_stats(): def run_local_scans(): - scan_pi_sessions() - complete = scan_opencode_db() + complete = scan_pi_sessions() + complete = scan_opencode_db() and complete return local_stats(), complete @@ -551,10 +629,13 @@ def read_auth_key(path): return "" # Only the subscription's own provider ids: a custom gateway named like # "opencode-local" holds a key for somewhere else, never to be sent here. + # An entry is an {type, key} object, or the key as a plain string. entries = [parsed[provider] for provider in PROVIDERS if provider in parsed] for entry in entries: - if isinstance(entry, dict) and entry.get("key"): - return str(entry["key"]).strip() + if isinstance(entry, str) and entry.strip(): + return entry.strip() + if isinstance(entry, dict) and isinstance(entry.get("key"), str) and entry["key"].strip(): + return entry["key"].strip() return "" @@ -605,7 +686,7 @@ def credential_candidates(): # one the endpoint accepts isn't knowable from here, so a refused one only # hands over to the next. store_keys, console = read_store_credentials(opencode_db()) - tokens = [str(os.environ.get("OPENCODE_API_KEY", "")).strip()] + tokens = [str(os.environ.get(name, "")).strip() for name in ("OPENCODE_API_KEY", "OPENCODE_GO_API_KEY")] tokens += store_keys tokens.append(read_auth_key(data_home() / "opencode" / "auth.json")) tokens.append(read_auth_key(Path.home() / ".pi" / "agent" / "auth.json")) diff --git a/test/shell.d/agent-usage-opencode-limits-test.sh b/test/shell.d/agent-usage-opencode-limits-test.sh index 656583ba3be..837b0a9fb60 100644 --- a/test/shell.d/agent-usage-opencode-limits-test.sh +++ b/test/shell.d/agent-usage-opencode-limits-test.sh @@ -4,7 +4,7 @@ source "$(dirname "$0")/base-test.sh" require_command python3 -unset OPENCODE_API_KEY OPENCODE_DB +unset OPENCODE_API_KEY OPENCODE_GO_API_KEY OPENCODE_DB PI_CODING_AGENT_DIR # collect_limits reaches the Zen gateway, so the reader that interprets its # answer is exercised on its own: the collector loads as a module, and a fake @@ -290,6 +290,15 @@ with open(os.path.join(opencode_home, "opencode", "auth.json"), "w") as handle: json.dump({"opencode-local": {"type": "api_key", "key": "gateway_secret"}}, handle) check("OpenCode collector never sends a look-alike provider's key", "gateway_secret" not in [c.token for c in collector.credential_candidates()]) +# An auth.json entry may be the key itself, and the Go key has an +# environment variable of its own beside OPENCODE_API_KEY. +with open(os.path.join(opencode_home, "opencode", "auth.json"), "w") as handle: + json.dump({"opencode-go": " plain_key "}, handle) +check("OpenCode collector reads a plain-string auth.json key", [c.token for c in collector.credential_candidates()] == ["plain_key"]) +os.environ["OPENCODE_GO_API_KEY"] = "go_env" +check("OpenCode collector reads OPENCODE_GO_API_KEY first", [c.token for c in collector.credential_candidates()] == ["go_env", "plain_key"]) +os.environ.pop("OPENCODE_GO_API_KEY") + clear_key() pi_home = os.path.join(home, "pi") os.makedirs(os.path.join(pi_home, ".pi", "agent"), exist_ok=True) diff --git a/test/shell.d/agent-usage-opencode-scanner-test.sh b/test/shell.d/agent-usage-opencode-scanner-test.sh index bab2be2718e..794f93f896c 100644 --- a/test/shell.d/agent-usage-opencode-scanner-test.sh +++ b/test/shell.d/agent-usage-opencode-scanner-test.sh @@ -5,7 +5,7 @@ source "$(dirname "$0")/base-test.sh" require_command jq require_command python3 -unset OPENCODE_API_KEY OPENCODE_DB +unset OPENCODE_API_KEY OPENCODE_GO_API_KEY OPENCODE_DB PI_CODING_AGENT_DIR timestamp="$(date +%Y-%m-%d)T12:00:00Z" @@ -263,10 +263,69 @@ result=$(HOME="$ODD_HOME" XDG_CACHE_HOME="$ODD_HOME/.cache" XDG_DATA_HOME="$ODD_ fail "OpenCode collector survives malformed nested fields" "$result" pass "OpenCode collector survives malformed nested fields" +# Profiles keep their own session trees, PI_CODING_AGENT_DIR moves pi's, and +# a home that is a git checkout ignoring everything still gets scanned. A +# fork's copy of its parent's message counts once, a usage block with only a +# total still counts, and either spelling of the token fields reads. +ROOTS_HOME=$(mktemp -d) +trap 'rm -rf "$PI_HOME" "$NOID_HOME" "$OPENCODE_HOME" "$CACHE_HOME" "$SESSION_HOME" "$V2_HOME" "$ODD_HOME" "$ROOTS_HOME"' EXIT +mkdir -p "$ROOTS_HOME/.git" "$ROOTS_HOME/.pi/agent/sessions/p" "$ROOTS_HOME/.omp/profiles/work/agent/sessions/p" \ + "$ROOTS_HOME/.pi/profiles/home/agent/sessions/p" "$ROOTS_HOME/moved/sessions/p" +printf '*\n' >"$ROOTS_HOME/.gitignore" +cat >"$ROOTS_HOME/.pi/agent/sessions/p/parent.jsonl" <"$ROOTS_HOME/.pi/agent/sessions/p/fork.jsonl" <"$ROOTS_HOME/.omp/profiles/work/agent/sessions/p/omp.jsonl" <"$ROOTS_HOME/.pi/profiles/home/agent/sessions/p/pi.jsonl" <"$ROOTS_HOME/moved/sessions/p/pi.jsonl" </dev/null 2>&1 +[[ -z $(find "$ROOTS_HOME/.cache" -name 'opencode-scan-*.json' 2>/dev/null) ]] || + fail "OpenCode collector doesn't cache a scan that couldn't read the sessions" +pass "OpenCode collector doesn't cache a scan that couldn't read the sessions" + +# A timestamp or token count that can't be trusted costs its row or field, +# never the scan. +BAD_HOME=$(mktemp -d) +trap 'rm -rf "$PI_HOME" "$NOID_HOME" "$OPENCODE_HOME" "$CACHE_HOME" "$SESSION_HOME" "$V2_HOME" "$ODD_HOME" "$ROOTS_HOME" "$BAD_HOME"' EXIT +mkdir -p "$BAD_HOME/.pi/agent/sessions/p" +cat >"$BAD_HOME/.pi/agent/sessions/p/pi.jsonl" < Date: Wed, 14 Oct 2026 23:14:20 -0400 Subject: [PATCH 084/546] Round OpenCode's probe times down so an immediate rerun reuses them Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-opencode | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/bin/omarchy-agent-usage-opencode b/bin/omarchy-agent-usage-opencode index a8e938812b5..8427fffb4da 100755 --- a/bin/omarchy-agent-usage-opencode +++ b/bin/omarchy-agent-usage-opencode @@ -857,13 +857,15 @@ def collect_limits(force): probe = probe_limits(candidate) if probe.get("ok"): - fetched_at = round(time.time() * 1000) + # Rounded down: rounded up, a run in the same millisecond would find the + # answer in the future and ask again. + fetched_at = int(time.time() * 1000) remember(cache, {"fetchedAtMs": fetched_at, "limits": probe["limits"]}) return answer(probe["limits"], fetched_at, False) if probe.get("noPlan"): # The windows of a plan that's gone mean nothing now, so they go too. - remember(cache, {"fetchedAtMs": round(time.time() * 1000), "limits": [], "noPlan": True}) + remember(cache, {"fetchedAtMs": int(time.time() * 1000), "limits": [], "noPlan": True}) no_plan = True continue From e619df752315ee28c6b5e54e948d86097f6d0a9d Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:15:08 -0400 Subject: [PATCH 085/546] Bring DeepSeek in line with the collector contract The local scan is cached per day and rebuilt under flock, with the usual reuse windows and --force, and a scan that couldn't read every source isn't cached. The balance is asked at most every 15 seconds and kept per key, so a failed check shows the last one stale instead of dropping it; a refused key reads as Sign-in expired, and only an unreachable API, timeouts included, advises a retry. pi and omp profiles and PI_CODING_AGENT_DIR are scanned with rg --no-ignore, a fork's copy of a message counts once, a line without an id always counts, and untrusted timestamps and counts skip their row or field. Local usage alone makes DeepSeek ready. Co-authored-by: Alvaro Minarro Gil Co-authored-by: ivanvan08 Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-deepseek | 561 +++++++++++++----- .../agent-usage-deepseek-balance-test.sh | 219 +++++++ .../agent-usage-deepseek-scanner-test.sh | 242 ++++---- 3 files changed, 740 insertions(+), 282 deletions(-) create mode 100755 test/shell.d/agent-usage-deepseek-balance-test.sh diff --git a/bin/omarchy-agent-usage-deepseek b/bin/omarchy-agent-usage-deepseek index fedc514e346..bee20bdb0e9 100755 --- a/bin/omarchy-agent-usage-deepseek +++ b/bin/omarchy-agent-usage-deepseek @@ -20,11 +20,17 @@ like the codex collector treats pi sessions. from __future__ import annotations import argparse +import fcntl +import hashlib import json +import math 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 @@ -36,16 +42,53 @@ AGENT_ID = "deepseek" AGENT_NAME = "DeepSeek" API_BASE_URL = "https://api.deepseek.com" +# A panel opened and closed repeatedly must not become one request per flick. +PROBE_MIN_INTERVAL_SECONDS = 15 + +# 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). --limits-only promises only a fresh balance, so it may reuse a scan +# for up to 15 minutes. +SCAN_REUSE_SECONDS = 20 +LIMITS_ONLY_REUSE_SECONDS = 900 class DeepSeekError(Exception): - pass + def __init__(self, message: str, transport: bool = False, refused: bool = False) -> None: + super().__init__(message) + self.transport = transport + self.refused = refused + + +# Counts and times read from disk or the network are untrusted: a boolean, +# NaN, negative, or absurdly large value reads as nothing, never as a number +# that would swamp the totals. The ceiling sits well above any epoch +# millisecond time. +NUMBER_CEILING = 10 ** 15 def number(value: Any) -> int: + if isinstance(value, bool): + return 0 try: - return int(value or 0) - except Exception: + n = float(value or 0) + except (TypeError, ValueError, OverflowError): + return 0 + if not math.isfinite(n) or n < 0 or n >= NUMBER_CEILING: return 0 + return round(n) + + +def obj(value: Any) -> dict[str, Any]: + return value if isinstance(value, dict) else {} + + +# pi, omp, and the harnesses that copy their format spell usage camelCase or +# snake_case; the first spelling present wins. +def usage_tokens(usage: dict[str, Any], *keys: str) -> int: + for key in keys: + if key in usage: + return number(usage.get(key)) + return 0 def model_name(raw: Any) -> str: @@ -53,25 +96,34 @@ def model_name(raw: Any) -> str: return value if value else AGENT_ID -def local_day(value: Any) -> str: +# The local day a message was spent on, or None when its timestamp can't be +# trusted, so the caller skips that row rather than the whole scan. A message +# with no timestamp at all is today's. +def local_day(value: Any) -> str | None: if value is None: return datetime.now().strftime("%Y-%m-%d") + if isinstance(value, str) and value.strip().isdigit(): + value = int(value.strip()) + if isinstance(value, bool): + return None if isinstance(value, (int, float)): + seconds = number(value) + if seconds <= 0: + return None # 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) + if seconds > 10_000_000_000: + seconds = seconds / 1000 + try: + return datetime.fromtimestamp(seconds).strftime("%Y-%m-%d") + except (OverflowError, OSError, ValueError): + return None 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") + dt = datetime.fromisoformat(str(value).strip().replace("Z", "+00:00")) + except ValueError: + return None + if dt.tzinfo is not None: + dt = dt.astimezone() + return dt.strftime("%Y-%m-%d") def runtime_env() -> dict[str, str]: @@ -103,42 +155,6 @@ def empty_bucket() -> dict[str, int]: } -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": 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": "", - "usageStatusText": "", - "authHelpText": "", - "limits": [], - } - record.update(empty_stats()) - record.update(overrides) - return record - - class Stats: """Token and activity totals shared by every local source. @@ -197,32 +213,56 @@ class Stats: } +# ------------------------------------------------------------- pi and omp + +# pi's and omp's session trees, each profile's too: `--profile=` moves +# the whole agent tree under /profiles//, and PI_CODING_AGENT_DIR +# moves pi's. def pi_session_roots() -> list[Path]: - return [ - Path.home() / ".pi" / "agent" / "sessions", - Path.home() / ".omp" / "agent" / "sessions", - ] + roots: list[Path] = [] + if os.environ.get("PI_CODING_AGENT_DIR"): + roots.append(Path(os.environ["PI_CODING_AGENT_DIR"]).expanduser() / "sessions") + for base in (Path.home() / ".pi", Path.home() / ".omp"): + roots.append(base / "agent" / "sessions") + try: + # Sorted so the scan order does not depend on directory order. + roots.extend(sorted(child / "agent" / "sessions" for child in (base / "profiles").iterdir() if child.is_dir())) + except OSError: + # No profiles directory, or it is unreadable: the default root stands. + pass + return list(dict.fromkeys(roots)) -def scan_pi_sessions(stats: Stats) -> None: - """Add pi/omp session messages that ran on the deepseek provider.""" - seen_messages: set[str] = set() +def scan_pi_sessions(stats: Stats) -> bool: + """Add pi/omp session messages that ran on the deepseek provider. - rg = find_command("rg") or "rg" - for root in pi_session_roots(): - if not root.exists(): - continue + Returns whether every root was read in full, so a partial scan is never + cached as the whole story. + """ + roots = [root for root in pi_session_roots() if root.is_dir()] + if not roots: + return True + rg = find_command("rg") + if not rg: + print("omarchy-agent-usage-deepseek: rg is missing; pi and omp sessions were not scanned", file=sys.stderr) + return False + + seen_messages: set[tuple[str, str]] = set() + complete = True + for root in roots: try: proc = subprocess.Popen( - [rg, "--json", "-e", r'"provider"\s*:\s*"deepseek"', str(root)], + # --no-ignore: a home that is itself a git checkout would otherwise + # hide every session file behind its ignore rules. + [rg, "--json", "--no-ignore", "-e", r'"provider"\s*:\s*"deepseek"', str(root)], stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, text=True, errors="replace", env=ENV, ) - except FileNotFoundError: - continue + except OSError: + return False assert proc.stdout is not None for raw in proc.stdout: @@ -233,64 +273,80 @@ def scan_pi_sessions(stats: Stats) -> None: line = event.get("data", {}).get("lines", {}).get("text", "") path = event.get("data", {}).get("path", {}).get("text", "pi-session") entry = json.loads(line) + add_pi_entry(stats, entry, path, seen_messages) except Exception: + # One odd line costs that line, never the scan. 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) + # rg exits 1 for no match and 2 when it couldn't read something. + if proc.wait(timeout=1) not in (0, 1): + complete = False except Exception: proc.kill() + complete = False + return complete +def add_pi_entry(stats: Stats, entry: Any, path: str, seen_messages: set[tuple[str, str]]) -> None: + if not isinstance(entry, dict) or entry.get("type") != "message": + return + message = obj(entry.get("message")) + if message.get("role") != "assistant": + return + if str(message.get("provider") or "") != "deepseek": + return + # A fork copies its parent's entries with their id and timestamp, so the + # pair names a message across files; an id alone only within its file. + # Without an id, each line is its own. + message_id = entry.get("id") + if message_id: + timestamp = entry.get("timestamp") + key = (str(message_id), str(timestamp)) if timestamp else (path, str(message_id)) + if key in seen_messages: + return + seen_messages.add(key) + + usage = obj(message.get("usage")) + if not usage: + return + input_tokens = usage_tokens(usage, "input", "inputTokens", "input_tokens") + output_tokens = usage_tokens(usage, "output", "outputTokens", "output_tokens") + cache_read = usage_tokens(usage, "cacheRead", "cacheReadTokens", "cache_read", "cache_read_input_tokens") + cache_write = usage_tokens(usage, "cacheWrite", "cacheWriteTokens", "cache_write", "cache_creation_input_tokens") + if not (input_tokens or output_tokens or cache_read or cache_write): + # Some writers keep only the total; it was still spent. + input_tokens = usage_tokens(usage, "totalTokens", "total_tokens") + if not input_tokens: + return + + day = local_day(entry.get("timestamp") or message.get("timestamp")) + if day is None: + return + stats.add(path, day, model_name(message.get("model")), input_tokens, output_tokens, cache_read, cache_write) + + +# --------------------------------------------------------------- opencode + 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: +def scan_opencode_sessions(stats: Stats) -> bool: """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. + database error keeps what it already read, but says it is incomplete. """ db = opencode_db_path() if not db.is_file(): - return - - import sqlite3 + return True try: conn = sqlite3.connect(db.resolve().as_uri() + "?mode=ro", uri=True, timeout=2) except sqlite3.Error: - return + return False try: conn.execute("PRAGMA query_only = ON") # The json_extract conditions are the authority; the LIKE gates are pure @@ -309,8 +365,8 @@ def scan_opencode_sessions(stats: Stats) -> None: continue if str(entry.get("providerID") or "") != "deepseek": continue - tokens = entry.get("tokens") or {} - cache = tokens.get("cache") or {} + tokens = obj(entry.get("tokens")) + cache = obj(tokens.get("cache")) 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")) @@ -318,15 +374,116 @@ def scan_opencode_sessions(stats: Stats) -> None: 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")) + day = local_day(obj(entry.get("time")).get("created")) + if day is None: + continue 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 + return False finally: conn.close() + return True + + +# ------------------------------------------------------------- scan cache + +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_json(path: Path) -> dict[str, Any]: + try: + value = json.loads(path.read_text(encoding="utf-8")) + except Exception: + return {} + return value if isinstance(value, dict) else {} + + +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: + # A negative age means the mtime is in the future: the clock moved + # backwards since the write, so the scan's freshness cannot be trusted. + age = time.time() - path.stat().st_mtime + if 0 <= age < max_age_seconds: + return read_json(path) + except OSError: + return None + return None + + +def write_json(path: Path, payload: dict[str, Any]) -> None: + # A temp name unique to this writer: 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 + + +def scan_cache_paths() -> tuple[Path, Path]: + # The digest covers every data path the scan reads. + sources = [str(opencode_db_path())] + [str(root) for root in pi_session_roots()] + digest = hashlib.sha1("\n".join(sources).encode("utf-8")).hexdigest()[:16] + root = cache_root() + return root / f"deepseek-scan-{digest}.json", root / f"deepseek-scan-{digest}.lock" + + +def run_local_scans() -> tuple[dict[str, Any], bool]: + stats = Stats() + complete = scan_pi_sessions(stats) + complete = scan_opencode_sessions(stats) and complete + return stats.record(), complete + + +# today* fields only mean "today" on the day they were scanned, so a scan +# from another local date is a miss whatever its file time says. +def read_cached_stats(cache_file: Path, max_age: float) -> dict[str, Any] | None: + cached = read_fresh_json(cache_file, max_age) + if not cached or cached.get("schemaVersion") != 1 or cached.get("scanDate") != datetime.now().strftime("%Y-%m-%d"): + return None + stats = cached.get("stats") + if not isinstance(stats, dict) or not all(key in stats for key in ("todayPrompts", "todayTotalTokens", "recentDays", "activeDates", "modelUsage")): + return None + return stats + + +def cached_local_stats(max_age: float) -> dict[str, Any]: + # The cache is an optimization: any failure of its own degrades to a + # direct scan. + try: + cache_file, lock_file = scan_cache_paths() + cached = read_cached_stats(cache_file, max_age) + if cached is not None: + return cached + with lock_file.open("w") as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + cached = read_cached_stats(cache_file, max_age) + if cached is not None: + return cached + stats, complete = run_local_scans() + # An interrupted scan still serves this run, but caching it would hide + # the missing usage from every reader until the cache expires. + if complete: + write_json(cache_file, {"schemaVersion": 1, "scanDate": datetime.now().strftime("%Y-%m-%d"), "stats": stats}) + return stats + except Exception as exc: + print(f"omarchy-agent-usage-deepseek: cache unavailable ({exc}); scanning directly", file=sys.stderr) + stats, _ = run_local_scans() + return stats # -------------------------------------------------------------- credentials @@ -335,6 +492,8 @@ def read_key_from_dict(parsed: Any, provider: str) -> str: if not isinstance(parsed, dict): return "" entry = parsed.get(provider) + if isinstance(entry, str): + return entry.strip() if not isinstance(entry, dict): return "" return str(entry.get("key") or entry.get("apiKey") or entry.get("api_key") or "").strip() @@ -343,19 +502,25 @@ def read_key_from_dict(parsed: Any, provider: str) -> str: 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): + except (OSError, ValueError): return "" def api_key() -> str: + # The environment first, then the agents that keep a DeepSeek key, read + # only: pi, then opencode. 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") + for read in ( + lambda: read_json_key(Path.home() / ".pi" / "agent" / "auth.json", "deepseek"), + lambda: read_json_key(data_home / "opencode" / "auth.json", "deepseek"), + ): + key = read() + if key: + return key + return "" def config_path() -> Path: @@ -367,12 +532,22 @@ 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): + except (OSError, ValueError): return {} # ----------------------------------------------------------------- balance +def money(value: Any) -> Decimal: + if isinstance(value, bool): + return Decimal("0") + try: + amount = Decimal(str(value or "0")) + except InvalidOperation: + return Decimal("0") + return amount if amount.is_finite() else Decimal("0") + + 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" @@ -382,27 +557,26 @@ def fetch_balance(api_base_url: str, key: str) -> dict[str, Any]: ) try: with urllib.request.urlopen(request, timeout=15) as response: - payload = json.load(response) + body = response.read() 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 error.code in (401, 403): + raise DeepSeekError("DeepSeek rejected the saved API key. Set a current one in DEEPSEEK_API_KEY or the agent that holds it.", refused=True) + raise DeepSeekError(f"DeepSeek's balance endpoint returned HTTP {error.code}.") + except Exception as error: + # No answer at all: no route, no DNS, a timeout. Worth a retry soon. + raise DeepSeekError("Couldn't reach the DeepSeek API. Retrying shortly.", transport=True) from error + try: + payload = json.loads(body.decode("utf-8", errors="replace")) + except ValueError 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") + 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") + raise DeepSeekError("DeepSeek's 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")), @@ -411,65 +585,142 @@ def fetch_balance(api_base_url: str, key: str) -> dict[str, Any]: } -def balance_record(key: str, api_base_url: str) -> dict[str, Any]: +def balance_record(remaining: Decimal, currency: 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 = money(config.get("fundedAmount")) + if 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"), + "currency": str(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()) +# One cache per key and endpoint, so a machine signed in to another account +# never shows the last one's balance. +def balance_cache_path(key: str, api_base_url: str) -> Path: + digest = hashlib.sha256((api_base_url.rstrip("/") + "\n" + key).encode("utf-8")).hexdigest()[:16] + return cache_root() / f"deepseek-limits-{digest}.json" + + +def collect_balance(key: str, api_base_url: str, force: bool) -> dict[str, Any]: + result: dict[str, Any] = {"balance": None, "stale": False, "fetchedAtMs": 0, "usageStatusText": "", "authHelpText": "", "retryAdvised": False} + cache = None + cached: dict[str, Any] = {} + try: + cache = balance_cache_path(key, api_base_url) + cached = read_json(cache) + except OSError as exc: + print(f"omarchy-agent-usage-deepseek: could not read the balance cache ({exc})", file=sys.stderr) + kept = obj(cached.get("balance")) + fetched_at = number(cached.get("fetchedAtMs")) + + # A panel opened and closed repeatedly must not become one request per + # flick, so a recent answer is reused; --force skips the wait. + if kept and not force and 0 <= time.time() - fetched_at / 1000 < PROBE_MIN_INTERVAL_SECONDS: + result.update(balance=balance_record(money(kept.get("remaining")), str(kept.get("currency") or "")), fetchedAtMs=fetched_at) + return result + + try: + live = fetch_balance(api_base_url, key) + except DeepSeekError as error: + result["retryAdvised"] = error.transport + result["authHelpText"] = str(error) + if kept: + # The last good balance, marked stale with when it was read, beats an + # empty section while DeepSeek can't be asked. + result.update(balance=balance_record(money(kept.get("remaining")), str(kept.get("currency") or "")), stale=True, fetchedAtMs=fetched_at) + if error.refused: + result["usageStatusText"] = "Sign-in expired" + elif not kept: + result["usageStatusText"] = "Balance unavailable" + return result + + # Rounded down: rounded up, a run in the same millisecond would find the + # answer in the future and ask again. + fetched_at = int(time.time() * 1000) + if cache is not None: + try: + write_json(cache, {"fetchedAtMs": fetched_at, "limits": [], "balance": {"remaining": str(live["remaining"]), "currency": live["currency"]}}) + except OSError as exc: + print(f"omarchy-agent-usage-deepseek: could not write the balance cache ({exc})", file=sys.stderr) + result.update(balance=balance_record(live["remaining"], live["currency"]), fetchedAtMs=fetched_at) + return result + + +# ------------------------------------------------------------------ record + +def base_record(stats: dict[str, Any] | None = None) -> dict[str, Any]: + record: dict[str, Any] = { + "schemaVersion": 1, + "id": AGENT_ID, + "name": AGENT_NAME, + "updatedAt": datetime.now(timezone.utc).isoformat(), + "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": "", + "usageStatusText": "", + "authHelpText": "", + "limits": [], + "limitsStale": False, + "limitsFetchedAt": 0, + } + record.update(stats if stats is not None else Stats().record()) + return record + + +def collect(api_base_url: str, max_age: float, force: bool) -> dict[str, Any]: + record = base_record(cached_local_stats(max_age)) + record["ready"] = number(record.get("totalPrompts")) > 0 # 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: - record["usageStatusText"] = "Balance unavailable" - record["authHelpText"] = str(error) + if not key: + return record + record.update(ready=True, tierLabel="Prepaid") + balance = collect_balance(key, api_base_url, force) + if balance["balance"] is not None: + record["balance"] = balance["balance"] + record.update( + usageStatusText=balance["usageStatusText"], + authHelpText=balance["authHelpText"], + limitsStale=balance["stale"], + limitsFetchedAt=balance["fetchedAtMs"], + ) + if balance["retryAdvised"]: + record["retryAdvised"] = True 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. + # --force rescans everything and asks DeepSeek past the reuse window. + # --limits-only (the panel's quick refresh) may reuse a scan for 15 + # minutes, but asks for the balance like a normal run. 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() + max_age = 0 if args.force else (LIMITS_ONLY_REUSE_SECONDS if args.limits_only else SCAN_REUSE_SECONDS) try: - record = scan(args.api_base_url) + record = collect(args.api_base_url, max_age, args.force) 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) + # A bug, not a state the panel should explain: an empty record keeps + # the panel steady, and stderr says what broke. + record = base_record() + print(f"omarchy-agent-usage-deepseek: {type(error).__name__}: {error}", file=sys.stderr) print(json.dumps(record, separators=(",", ":"))) return 0 diff --git a/test/shell.d/agent-usage-deepseek-balance-test.sh b/test/shell.d/agent-usage-deepseek-balance-test.sh new file mode 100755 index 00000000000..ce367270627 --- /dev/null +++ b/test/shell.d/agent-usage-deepseek-balance-test.sh @@ -0,0 +1,219 @@ +#!/bin/bash + +source "$(dirname "$0")/base-test.sh" + +require_command python3 + +unset DEEPSEEK_API_KEY DEEPSEEK_API_BASE_URL PI_CODING_AGENT_DIR + +# The balance comes from DeepSeek's API, so the reader that interprets its +# answer is exercised on its own: the collector loads as a module, and a fake +# urlopen stands in for the response. +TEST_HOME=$(mktemp -d) +trap 'rm -rf "$TEST_HOME"' EXIT + +COLLECTOR="$ROOT/bin/omarchy-agent-usage-deepseek" TEST_HOME="$TEST_HOME" python3 - <<'PY' +import importlib.machinery +import importlib.util +import io +import json +import os +import sys +import time +import urllib.error +from pathlib import Path + +home = os.environ["TEST_HOME"] +# Never the real sign-ins: every store lives under the test home. +os.environ["HOME"] = os.path.join(home, "home") +os.environ["XDG_CACHE_HOME"] = os.path.join(home, "cache") +os.environ["XDG_DATA_HOME"] = os.path.join(home, "data") +os.environ["XDG_CONFIG_HOME"] = os.path.join(home, "config") +os.makedirs(os.environ["HOME"], exist_ok=True) + +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) + +BASE = "https://api.deepseek.com" +cache_dir = Path(os.environ["XDG_CACHE_HOME"]) / "omarchy" / "agent-usage" + + +def check(description, condition, detail=""): + if condition: + print("ok - " + description) + return + print("not ok - " + description, file=sys.stderr) + if detail: + print(detail, file=sys.stderr) + sys.exit(1) + + +asked = [] + + +def answer(total="42.50", currency="CNY", error=None, body=None): + def open_(request, timeout=None): + asked.append((request.full_url, request.get_header("Authorization"))) + if error is not None: + raise error + if body is not None: + return io.BytesIO(body) + return io.BytesIO(json.dumps({ + "is_available": True, + "balance_infos": [{"currency": currency, "total_balance": total, "granted_balance": "2.50", "topped_up_balance": "40.00"}], + }).encode()) + return open_ + + +def clear_cache(): + for cache in cache_dir.glob("deepseek-limits-*.json"): + cache.unlink() + + +def record(force=True): + return collector.collect(BASE, 0, force) + + +# Unused: no key and no local usage is a record the panel skips. +collector.urllib.request.urlopen = answer(error=RuntimeError("no key, no request")) +result = record() +check( + "DeepSeek collector prints an empty record on an unused machine", + result["ready"] is False and result["tierLabel"] == "" and "balance" not in result and result["limits"] == [] + and result["usageStatusText"] == "" and result["authHelpText"] == "" and "retryAdvised" not in result and asked == [], + json.dumps(result), +) + +os.environ["DEEPSEEK_API_KEY"] = "key-123" +collector.urllib.request.urlopen = answer() +result = record() +check( + "DeepSeek collector reads the live balance", + result["ready"] is True and result["tierLabel"] == "Prepaid" and result["balance"]["remaining"] == 42.5 + and result["balance"]["currency"] == "CNY" and result["balance"]["estimated"] is False + and result["limitsStale"] is False and result["limitsFetchedAt"] > 0 and result["usageStatusText"] == "" + and asked[-1] == ("https://api.deepseek.com/user/balance", "Bearer key-123"), + json.dumps({"result": result, "asked": asked}), +) +fetched_at = result["limitsFetchedAt"] + +# A panel that opens and shuts repeatedly must not ask every time. +asked.clear() +collector.urllib.request.urlopen = answer(error=RuntimeError("probe should have been skipped")) +result = record(force=False) +check( + "DeepSeek collector reuses a balance it just read", + asked == [] and result["balance"]["remaining"] == 42.5 and result["limitsStale"] is False and result["limitsFetchedAt"] == fetched_at, + json.dumps(result), +) + +collector.urllib.request.urlopen = answer(total="40.00") +result = record(force=True) +check( + "DeepSeek collector --force asks past the reuse window", + len(asked) == 1 and result["balance"]["remaining"] == 40.0, + json.dumps(result), +) + +# The estimated ledger comes from fundedAmount, applied to a kept balance too. +config = Path(os.environ["XDG_CONFIG_HOME"]) / "omarchy" / "agents" +config.mkdir(parents=True, exist_ok=True) +(config / "deepseek.json").write_text(json.dumps({"fundedAmount": 50})) +collector.urllib.request.urlopen = answer(total="8.15", currency="USD") +result = record() +check( + "DeepSeek collector derives the estimated ledger from fundedAmount", + result["balance"] == {"remaining": 8.15, "funded": 50.0, "spent": 41.85, "currency": "USD", "estimated": True}, + json.dumps(result), +) +fetched_at = result["limitsFetchedAt"] + +# A failed check keeps the last good balance, marked stale with when it was +# read; only an unreachable service asks for a quick retry. +collector.urllib.request.urlopen = answer(error=urllib.error.URLError("no route")) +result = record() +check( + "DeepSeek collector keeps the last balance, stale, when DeepSeek can't be reached", + result["balance"]["remaining"] == 8.15 and result["balance"]["estimated"] is True and result["limitsStale"] is True + and result["limitsFetchedAt"] == fetched_at and result["usageStatusText"] == "" and result.get("retryAdvised") is True + and "Couldn't reach" in result["authHelpText"], + json.dumps(result), +) + +collector.urllib.request.urlopen = answer(error=TimeoutError("timed out")) +result = record() +check( + "DeepSeek collector advises a retry on a timeout", + result.get("retryAdvised") is True and result["limitsStale"] is True and "Couldn't reach" in result["authHelpText"], + json.dumps(result), +) + +collector.urllib.request.urlopen = answer(error=urllib.error.HTTPError(BASE, 500, "Server Error", {}, None)) +result = record() +check( + "DeepSeek collector doesn't advise a retry on a server error", + "retryAdvised" not in result and result["limitsStale"] is True and "HTTP 500" in result["authHelpText"], + json.dumps(result), +) + +collector.urllib.request.urlopen = answer(body=b"oops") +result = record() +check( + "DeepSeek collector reads a garbled answer as invalid, not unreachable", + "retryAdvised" not in result and "invalid balance response" in result["authHelpText"] and result["limitsStale"] is True, + json.dumps(result), +) + +# A cache written before the clock went back is still the last good answer. +future = time.time() + 86400 +cache = collector.balance_cache_path("key-123", BASE) +os.utime(cache, (future, future)) +collector.urllib.request.urlopen = answer(error=urllib.error.URLError("no route")) +result = record() +check( + "DeepSeek collector keeps a balance cache whose file time is in the future", + result["balance"]["remaining"] == 8.15 and result["limitsStale"] is True, + json.dumps(result), +) + +# A refused key keeps its last balance, stale, and asks for a new sign-in. +collector.urllib.request.urlopen = answer(error=urllib.error.HTTPError(BASE, 401, "Unauthorized", {}, None)) +result = record() +check( + "DeepSeek collector asks for a new key when DeepSeek refuses it", + result["usageStatusText"] == "Sign-in expired" and result["balance"]["remaining"] == 8.15 and result["limitsStale"] is True + and "retryAdvised" not in result, + json.dumps(result), +) + +# With nothing kept, a failed check says the balance is unavailable. +clear_cache() +collector.urllib.request.urlopen = answer(error=urllib.error.HTTPError(BASE, 500, "Server Error", {}, None)) +result = record() +check( + "DeepSeek collector says the balance is unavailable with nothing kept", + result["usageStatusText"] == "Balance unavailable" and "balance" not in result and result["limitsStale"] is False + and result["limitsFetchedAt"] == 0 and result["ready"] is True and result["tierLabel"] == "Prepaid", + json.dumps(result), +) + +# Another key never shows this one's balance. +collector.urllib.request.urlopen = answer(total="1.00") +record() +os.environ["DEEPSEEK_API_KEY"] = "key-other" +collector.urllib.request.urlopen = answer(error=urllib.error.URLError("no route")) +result = record() +check( + "DeepSeek collector keeps a balance per key", + "balance" not in result and result["usageStatusText"] == "Balance unavailable" and result.get("retryAdvised") is True, + json.dumps(result), +) + +# Untrusted amounts read as nothing rather than as a number. +check( + "DeepSeek collector reads untrusted amounts as zero", + [float(collector.money(value)) for value in ("NaN", "Infinity", True, "abc", None, "3.5")] == [0.0, 0.0, 0.0, 0.0, 0.0, 3.5], +) +PY diff --git a/test/shell.d/agent-usage-deepseek-scanner-test.sh b/test/shell.d/agent-usage-deepseek-scanner-test.sh index 91226a78aa1..e37dc248f1e 100755 --- a/test/shell.d/agent-usage-deepseek-scanner-test.sh +++ b/test/shell.d/agent-usage-deepseek-scanner-test.sh @@ -5,18 +5,33 @@ source "$(dirname "$0")/base-test.sh" require_command jq 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 +unset DEEPSEEK_API_KEY DEEPSEEK_API_BASE_URL PI_CODING_AGENT_DIR timestamp="$(date +%Y-%m-%d)T12:00:00Z" +now_ms="$(date +%s)000" + +HOMES=() +trap 'rm -rf "${HOMES[@]}"' EXIT + +new_home() { + local home + home=$(mktemp -d) + HOMES+=("$home") + printf '%s\n' "$home" +} + +# The collector run as it is installed, with every store under one home. +collect() { + local home="$1" + shift + HOME="$home" XDG_CACHE_HOME="$home/.cache" XDG_DATA_HOME="$home/.local/share" XDG_CONFIG_HOME="$home/.config" \ + "$ROOT/bin/omarchy-agent-usage-deepseek" "$@" 2>/dev/null +} # ---------------------------------------------------------------- pi and omp -PI_HOME=$(mktemp -d) -trap 'rm -rf "$TEST_HOME" "$PI_HOME"' EXIT +PI_HOME=$(new_home) 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) +result=$(collect "$PI_HOME") [[ $(jq -r '.todayTotalTokens' <<<"$result") == "49" ]] || fail "DeepSeek collector counts usage from pi and omp sessions" "$result" @@ -43,24 +57,69 @@ pass "DeepSeek collector filters pi and omp sessions to DeepSeek providers" 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) + ":" + .tierLabel + ":" + (.hasPromptStats|tostring) + ":" + (has("balance")|tostring)' <<<"$result") == "deepseek:false::true:false" ]] || +# Local usage alone puts DeepSeek in use; without a key there is no plan and +# no balance section. +[[ $(jq -c '{id, ready, tierLabel, hasPromptStats, balance: has("balance"), usageStatusText, authHelpText}' <<<"$result") == '{"id":"deepseek","ready":true,"tierLabel":"","hasPromptStats":true,"balance":false,"usageStatusText":"","authHelpText":""}' ]] || 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}' ]] || +# Nothing used and no key: a record the panel skips, with the full empty +# stats block. +EMPTY_HOME=$(new_home) +result=$(collect "$EMPTY_HOME") +[[ $(jq -c '{ready, tierLabel, limits, limitsStale, limitsFetchedAt, usageStatusText, authHelpText, totalPrompts, activeDays, activeDates, modelUsage, recentDays: (.recentDays | length), balance: has("balance")}' <<<"$result") == '{"ready":false,"tierLabel":"","limits":[],"limitsStale":false,"limitsFetchedAt":0,"usageStatusText":"","authHelpText":"","totalPrompts":0,"activeDays":0,"activeDates":[],"modelUsage":{},"recentDays":7,"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" +# Profiles keep their own session trees, PI_CODING_AGENT_DIR moves pi's, and +# a home that is a git checkout ignoring everything still gets scanned. A +# fork's copy of its parent's message counts once, a line without an id +# always counts, a usage block with only a total still counts, and either +# spelling of the token fields reads. +ROOTS_HOME=$(new_home) +mkdir -p "$ROOTS_HOME/.git" "$ROOTS_HOME/.pi/agent/sessions/p" "$ROOTS_HOME/.omp/profiles/work/agent/sessions/p" "$ROOTS_HOME/moved/sessions/p" +printf '*\n' >"$ROOTS_HOME/.gitignore" +cat >"$ROOTS_HOME/.pi/agent/sessions/p/parent.jsonl" <"$ROOTS_HOME/.pi/agent/sessions/p/fork.jsonl" <"$ROOTS_HOME/.omp/profiles/work/agent/sessions/p/omp.jsonl" <"$ROOTS_HOME/moved/sessions/p/pi.jsonl" <"$BAD_HOME/.pi/agent/sessions/p/pi.jsonl" </dev/null) +result=$(collect "$OPENCODE_HOME") [[ $(jq -r '.todayTotalTokens' <<<"$result") == "157" ]] || fail "DeepSeek collector counts opencode messages on a DeepSeek provider" "$result" @@ -110,116 +168,46 @@ 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) +result=$(collect "$OPENCODE_HOME" --force) [[ $(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] +# --------------------------------------------------------------- scan cache -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" +# A normal run reuses a scan for a few seconds, --limits-only for longer, +# and --force rescans. +CACHE_HOME=$(new_home) +mkdir -p "$CACHE_HOME/.pi/agent/sessions/p" +cat >"$CACHE_HOME/.pi/agent/sessions/p/pi.jsonl" <>"$CACHE_HOME/.pi/agent/sessions/p/pi.jsonl" </dev/null 2>&1 +[[ -z $(find "$CACHE_HOME/.cache" -name 'deepseek-scan-*.json' 2>/dev/null) ]] || + fail "DeepSeek collector doesn't cache a scan that couldn't read the sessions" +pass "DeepSeek collector doesn't cache a scan that couldn't read the sessions" From 07fbe4ba1bd17fc2c0377828461e7ab28e557d02 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:15:45 -0400 Subject: [PATCH 086/546] Count DeepSeek Harness sessions and read its key in DeepSeek dsh keeps a session log per session under $DSH_HOME/sessions, Zstandard- compressed by default. Its DeepSeek routes count, folded the way dsh's own token meter folds them: a settled answer replaces the streamed sample of its step, a retry adds an attempt, and a fork's inherited history stays its parent's. The key is read, never written, from dsh's managed .credentials.yaml and the .env beside it. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-deepseek | 237 +++++++++++++++++- .../agent-usage-deepseek-balance-test.sh | 2 +- .../agent-usage-deepseek-scanner-test.sh | 132 +++++++++- 3 files changed, 365 insertions(+), 6 deletions(-) diff --git a/bin/omarchy-agent-usage-deepseek b/bin/omarchy-agent-usage-deepseek index bee20bdb0e9..06838afa155 100755 --- a/bin/omarchy-agent-usage-deepseek +++ b/bin/omarchy-agent-usage-deepseek @@ -5,8 +5,9 @@ """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 +deepseek provider, from opencode sessions that ran on a deepseek provider, and +from DeepSeek Harness (dsh) session logs on its DeepSeek routes. 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 @@ -25,6 +26,7 @@ import hashlib import json import math import os +import re import shutil import sqlite3 import subprocess @@ -38,6 +40,11 @@ from decimal import Decimal, InvalidOperation from pathlib import Path from typing import Any +try: + from compression import zstd +except ImportError: + zstd = None + AGENT_ID = "deepseek" AGENT_NAME = "DeepSeek" API_BASE_URL = "https://api.deepseek.com" @@ -52,6 +59,16 @@ PROBE_MIN_INTERVAL_SECONDS = 15 SCAN_REUSE_SECONDS = 20 LIMITS_ONLY_REUSE_SECONDS = 900 +# DeepSeek Harness bills DeepSeek through these two provider routes: an API +# key, and a DeepSeek account sign-in. Its other routes (pi-ai's providers) +# belong to whoever bills them. +DSH_PROVIDERS = ("deepseek-official", "deepseek-account") + +# One canonical session log generation: session.jsonl, session.v4.jsonl, and +# their Zstandard twins. +DSH_LOG_NAME = re.compile(r"session(?:\.v([1-9][0-9]*))?\.jsonl(\.zstd)?") + + class DeepSeekError(Exception): def __init__(self, message: str, transport: bool = False, refused: bool = False) -> None: super().__init__(message) @@ -388,6 +405,164 @@ def scan_opencode_sessions(stats: Stats) -> bool: return True +# -------------------------------------------------------- DeepSeek Harness +# +# dsh keeps one append-only event log per session under +# $DSH_HOME/sessions///, Zstandard-compressed by default. +# Each billed model attempt settles as an assistant/message (or a log-only +# assistant/attempt) carrying the provider's usage; the request/header before +# it names the provider route and model. + + +def dsh_home() -> Path: + # An empty or blank DSH_HOME reads as unset, as dsh itself treats it. + configured = os.environ.get("DSH_HOME", "").strip() + return Path(configured).expanduser() if configured else Path.home() / ".dsh" + + +def dsh_session_logs(root: Path) -> list[Path]: + # The newest generation in each session directory is the current log; + # older ones are kept only as migration history. + newest: dict[Path, tuple[int, Path]] = {} + for path in root.rglob("session*.jsonl*"): + match = DSH_LOG_NAME.fullmatch(path.name) + if not match or not path.is_file(): + continue + version = int(match.group(1) or 0) + current = newest.get(path.parent) + if current is None or version > current[0]: + newest[path.parent] = (version, path) + return [path for _, path in sorted(newest.values(), key=lambda item: str(item[1]))] + + +def dsh_log_lines(path: Path) -> tuple[list[bytes], bool]: + """The complete records of one log, and whether all of it could be read. + + A torn final line or frame from a crash mid-write is dropped, as dsh's own + reader does, without making the scan incomplete. + """ + data = path.read_bytes() + if path.name.endswith(".zstd"): + if zstd is None: + print("omarchy-agent-usage-deepseek: this Python has no Zstandard support; DeepSeek Harness sessions were not scanned", file=sys.stderr) + return [], False + chunks = [] + while data: + decompressor = zstd.ZstdDecompressor() + try: + chunks.append(decompressor.decompress(data)) + except zstd.ZstdError: + break + if not decompressor.eof: + break + data = decompressor.unused_data + data = b"".join(chunks) + lines = data.split(b"\n") + # Whatever follows the last newline is an unfinished record. + return lines[:-1], True + + +def dsh_usage_sample(event: dict[str, Any]) -> dict[str, Any] | None: + data = obj(event.get("data")) + if event.get("type") == "assistant/message" and isinstance(data.get("usage"), dict): + return data["usage"] + stream = data.get("stream") + if not isinstance(stream, list): + return None + for record in reversed(stream): + chunk = obj(obj(record).get("chunk")) + if obj(record).get("type") == "chunk" and chunk.get("type") == "usage": + return obj(chunk.get("usage")) + return None + + +def scan_dsh_log(stats: Stats, path: Path) -> bool: + lines, complete = dsh_log_lines(path) + if not lines: + return complete + try: + header = obj(json.loads(lines[0])) + except ValueError: + return complete + if header.get("type") != "session": + return complete + events = [] + for line in lines[1:]: + try: + event = json.loads(line) + except ValueError: + continue + if isinstance(event, dict): + events.append(event) + + # A fork's log opens with a copy of its parent's history; the last + # `session/end-seed` marked inherited is where its own spending begins. + start = 0 + if header.get("isSeeded") is True: + for index, event in enumerate(events): + if event.get("type") == "session/end-seed" and obj(event.get("data")).get("inherited") is True: + start = index + 1 + + session = "dsh:" + str(header.get("id") or path.parent.name) + provider, model = "", "" + # One slot per billed attempt. A settlement replaces the streaming sample + # of the same turn and step, until a retry starts a new attempt, the same + # fold dsh's own token meter runs. + slots: list[list[Any]] = [] + last: tuple[Any, Any] | None = None + for index, event in enumerate(events): + kind = event.get("type") + data = obj(event.get("data")) + if kind == "request/header": + config = obj(obj(data.get("header")).get("config")) + provider, model = str(config.get("provider") or ""), str(config.get("model") or "") + continue + if kind == "llm/retry-started": + if last == (data.get("turn"), data.get("step")): + last = None + continue + if kind not in ("assistant/message", "assistant/attempt"): + continue + usage = dsh_usage_sample(event) + if usage is None: + continue + slot = (data.get("turn"), data.get("step")) + sample = [ + number(usage.get("inputTokens")), + number(usage.get("outputTokens")), + number(usage.get("cacheReadTokens")), + number(usage.get("cacheWriteTokens")), + ] + entry = [index >= start and provider in DSH_PROVIDERS, local_day(event.get("time")), model, sample] + if last == slot and slots: + slots[-1] = entry + else: + slots.append(entry) + last = slot + + for counted, day, slot_model, (input_tokens, output_tokens, cache_read, cache_write) in slots: + if counted and day is not None and (input_tokens or output_tokens or cache_read or cache_write): + stats.add(session, day, model_name(slot_model), input_tokens, output_tokens, cache_read, cache_write) + return complete + + +def scan_dsh_sessions(stats: Stats) -> bool: + root = dsh_home() / "sessions" + if not root.is_dir(): + return True + complete = True + try: + logs = dsh_session_logs(root) + except OSError: + return False + for path in logs: + try: + complete = scan_dsh_log(stats, path) and complete + except OSError: + complete = False + return complete + + # ------------------------------------------------------------- scan cache def cache_root() -> Path: @@ -436,7 +611,7 @@ def write_json(path: Path, payload: dict[str, Any]) -> None: def scan_cache_paths() -> tuple[Path, Path]: # The digest covers every data path the scan reads. - sources = [str(opencode_db_path())] + [str(root) for root in pi_session_roots()] + sources = [str(opencode_db_path()), str(dsh_home())] + [str(root) for root in pi_session_roots()] digest = hashlib.sha1("\n".join(sources).encode("utf-8")).hexdigest()[:16] root = cache_root() return root / f"deepseek-scan-{digest}.json", root / f"deepseek-scan-{digest}.lock" @@ -446,6 +621,7 @@ def run_local_scans() -> tuple[dict[str, Any], bool]: stats = Stats() complete = scan_pi_sessions(stats) complete = scan_opencode_sessions(stats) and complete + complete = scan_dsh_sessions(stats) and complete return stats.record(), complete @@ -506,9 +682,60 @@ def read_json_key(path: Path, provider: str) -> str: return "" +def unquote(value: str) -> str: + value = value.strip() + if value[:1] in ("'", '"'): + end = value.find(value[0], 1) + return value[1:end] if end > 0 else "" + # An unquoted scalar ends where a comment begins. + return value.split(" #", 1)[0].strip() + + +def read_dsh_credentials_key(path: Path) -> str: + """DEEPSEEK_API_KEY from dsh's managed `.credentials.yaml`. + + Only the one plain scalar is needed, so this reads the line rather than + pulling in a YAML parser: `refs:` holds it in the current layout, and the + pre-release layout kept it at the top level. + """ + try: + text = path.read_text(encoding="utf-8") + except (OSError, UnicodeDecodeError): + return "" + section = "" + for line in text.splitlines(): + if not line.strip() or line.lstrip().startswith("#"): + continue + if not line[0].isspace(): + section = line.split(":", 1)[0].strip() + if section == "DEEPSEEK_API_KEY": + return unquote(line.split(":", 1)[1]) + continue + if section == "refs": + name, _, value = line.strip().partition(":") + if name.strip() == "DEEPSEEK_API_KEY": + return unquote(value) + return "" + + +def read_dotenv_key(path: Path) -> str: + try: + text = path.read_text(encoding="utf-8") + except (OSError, UnicodeDecodeError): + return "" + for line in text.splitlines(): + line = line.strip() + if line.startswith("export "): + line = line[len("export "):].strip() + name, sep, value = line.partition("=") + if sep and name.strip() == "DEEPSEEK_API_KEY": + return unquote(value) + return "" + + def api_key() -> str: # The environment first, then the agents that keep a DeepSeek key, read - # only: pi, then opencode. + # only: pi, opencode, and dsh's managed store with its `.env` fallback. env_key = str(os.environ.get("DEEPSEEK_API_KEY", "")).strip() if env_key: return env_key @@ -516,6 +743,8 @@ def api_key() -> str: for read in ( lambda: read_json_key(Path.home() / ".pi" / "agent" / "auth.json", "deepseek"), lambda: read_json_key(data_home / "opencode" / "auth.json", "deepseek"), + lambda: read_dsh_credentials_key(dsh_home() / ".credentials.yaml"), + lambda: read_dotenv_key(dsh_home() / ".env"), ): key = read() if key: diff --git a/test/shell.d/agent-usage-deepseek-balance-test.sh b/test/shell.d/agent-usage-deepseek-balance-test.sh index ce367270627..a2b0246ca88 100755 --- a/test/shell.d/agent-usage-deepseek-balance-test.sh +++ b/test/shell.d/agent-usage-deepseek-balance-test.sh @@ -4,7 +4,7 @@ source "$(dirname "$0")/base-test.sh" require_command python3 -unset DEEPSEEK_API_KEY DEEPSEEK_API_BASE_URL PI_CODING_AGENT_DIR +unset DEEPSEEK_API_KEY DEEPSEEK_API_BASE_URL DSH_HOME PI_CODING_AGENT_DIR # The balance comes from DeepSeek's API, so the reader that interprets its # answer is exercised on its own: the collector loads as a module, and a fake diff --git a/test/shell.d/agent-usage-deepseek-scanner-test.sh b/test/shell.d/agent-usage-deepseek-scanner-test.sh index e37dc248f1e..e27bd3387bd 100755 --- a/test/shell.d/agent-usage-deepseek-scanner-test.sh +++ b/test/shell.d/agent-usage-deepseek-scanner-test.sh @@ -6,7 +6,7 @@ require_command jq require_command python3 # A key in the environment would send the collector to DeepSeek. -unset DEEPSEEK_API_KEY DEEPSEEK_API_BASE_URL PI_CODING_AGENT_DIR +unset DEEPSEEK_API_KEY DEEPSEEK_API_BASE_URL DSH_HOME PI_CODING_AGENT_DIR timestamp="$(date +%Y-%m-%d)T12:00:00Z" now_ms="$(date +%s)000" @@ -174,6 +174,136 @@ result=$(collect "$OPENCODE_HOME" --force) 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" +# ---------------------------------------------------------- DeepSeek Harness + +# dsh keeps one event log per session under $DSH_HOME/sessions, Zstandard- +# compressed by default. Only its DeepSeek routes count; a retry's settled +# answer replaces the streamed sample of the same step; a fork's inherited +# history is its parent's; a torn final record is dropped; and only the +# newest generation of a session is read. +DSH_ROOT=$(new_home) +python3 - "$DSH_ROOT/dsh/sessions" "$now_ms" <<'PY' +import json +import sys +from pathlib import Path + +root = Path(sys.argv[1]) +now = int(sys.argv[2]) + +try: + from compression import zstd +except ImportError: + zstd = None + + +def header(session, seeded=False): + return {"type": "session", "version": 4, "id": session, "createdAt": now, "isSeeded": seeded, "delegationDepth": 0} + + +def route(provider, model, seq=1): + return {"type": "request/header", "seq": seq, "time": now, "data": {"header": {"config": {"provider": provider, "model": model}}, "reason": "initial"}} + + +def message(turn, step, usage, seq=2, kind="assistant/message"): + data = {"turn": turn, "step": step, "stream": []} + if kind == "assistant/message": + data["usage"] = usage + else: + data["stream"] = [{"type": "chunk", "chunk": {"type": "usage", "usage": usage}}] + return {"type": kind, "seq": seq, "time": now, "data": data} + + +def usage(input, output=0, read=0, write=0): + return {"inputTokens": input, "outputTokens": output, "cacheReadTokens": read, "cacheWriteTokens": write} + + +def write(path, records, compressed, torn=b""): + path.parent.mkdir(parents=True, exist_ok=True) + lines = [(json.dumps(record) + "\n").encode() for record in records] + if compressed: + path.write_bytes(b"".join(zstd.compress(line) for line in lines) + torn) + else: + path.write_bytes(b"".join(lines) + torn) + + +compressed = zstd is not None +suffix = ".jsonl.zstd" if compressed else ".jsonl" +project = root / "--home-me-code--" + +# A DeepSeek session: one step streamed a sample, then settled with the real +# count; a retry of step 2 adds a second billed attempt. +write(project / "s1" / ("session.v4" + suffix), [ + header("s1"), + route("deepseek-official", "deepseek-v4-pro"), + message(1, 1, usage(1, 1), kind="assistant/attempt"), + message(1, 1, usage(100, 10, 20, 0)), + message(1, 2, usage(5, 1), kind="assistant/attempt"), + {"type": "llm/retry-started", "seq": 9, "time": now, "data": {"turn": 1, "step": 2}}, + message(1, 2, usage(7, 3)), +], compressed, torn=b'{"type":"assistant/mess' if not compressed else b"\x28\xb5\x2f\xfd") +# An older generation of the same session is migration history only. +write(project / "s1" / "session.jsonl", [header("s1"), route("deepseek-official", "deepseek-v4-pro"), message(1, 1, usage(999999))], False) + +# A route billed by someone else. +write(project / "s2" / ("session.v4" + suffix), [ + header("s2"), + route("pi-ai:anthropic", "claude-opus-5"), + message(1, 1, usage(888888)), +], compressed) + +# A fork of s1: its copy of the parent's step is inherited, and only what +# follows its seed marker was spent in it. +write(project / "s3" / ("session.v4" + suffix), [ + header("s3", seeded=True), + route("deepseek-official", "deepseek-v4-pro"), + message(1, 1, usage(100, 10, 20, 0)), + {"type": "session/end-seed", "seq": 4, "time": now, "data": {"inherited": True}}, + route("deepseek-account", "deepseek-v4-flash", seq=5), + message(2, 1, usage(4000, 6, 0, 1)), +], compressed) + +# A root written with compression off reads the same, torn tail and all. +write(project / "s4" / "session.v4.jsonl", [ + header("s4"), + route("deepseek-official", "deepseek-chat"), + message(1, 1, usage(30000)), +], False, torn=b'{"type":"assistant/mess') +PY + +result=$(DSH_HOME="$DSH_ROOT/dsh" collect "$DSH_ROOT") +[[ $(jq -c '{todayTotalTokens, totalPrompts, totalSessions, modelUsage}' <<<"$result") == '{"todayTotalTokens":34153,"totalPrompts":5,"totalSessions":3,"modelUsage":{"deepseek-v4-pro":{"inputTokens":112,"outputTokens":14,"cacheReadInputTokens":20,"cacheCreationInputTokens":0},"deepseek-v4-flash":{"inputTokens":4000,"outputTokens":6,"cacheReadInputTokens":0,"cacheCreationInputTokens":1},"deepseek-chat":{"inputTokens":30000,"outputTokens":0,"cacheReadInputTokens":0,"cacheCreationInputTokens":0}}}' ]] || + fail "DeepSeek collector counts DeepSeek Harness sessions on its DeepSeek routes" "$result" +pass "DeepSeek collector counts DeepSeek Harness sessions on its DeepSeek routes" + +# dsh keeps its key in its managed .credentials.yaml, or a .env beside it. +# With no network to answer, the key still makes DeepSeek ready and prepaid. +mkdir -p "$DSH_ROOT/dsh" +cat >"$DSH_ROOT/dsh/.credentials.yaml" <<'EOF' +# managed by dsh +version: 1 +refs: + OTHER_KEY: nope + DEEPSEEK_API_KEY: "sk-from-dsh" # stored by the Models page +EOF +key=$(HOME="$DSH_ROOT" DSH_HOME="$DSH_ROOT/dsh" python3 - "$ROOT/bin/omarchy-agent-usage-deepseek" <<'PY' +import importlib.machinery +import importlib.util +import sys + +loader = importlib.machinery.SourceFileLoader("deepseek_collector", sys.argv[1]) +spec = importlib.util.spec_from_loader(loader.name, loader) +collector = importlib.util.module_from_spec(spec) +loader.exec_module(collector) +first = collector.api_key() +(collector.dsh_home() / ".credentials.yaml").unlink() +(collector.dsh_home() / ".env").write_text("export DEEPSEEK_API_KEY='sk-from-env-file'\n") +print(first + " " + collector.api_key()) +PY +) +[[ $key == "sk-from-dsh sk-from-env-file" ]] || + fail "DeepSeek collector reads dsh's stored key" "$key" +pass "DeepSeek collector reads dsh's stored key" + # --------------------------------------------------------------- scan cache # A normal run reuses a scan for a few seconds, --limits-only for longer, From e9c33348138b73f0e766532437196b8888b03245 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:19:17 -0400 Subject: [PATCH 087/546] Floor every collector's probe stamps so a rerun reuses them A stamp rounded up lands half a millisecond in the future, so the next run read it as a clock that went back, failed the 0 <= age guard, and probed again; the probe-reuse tests failed now and then because of it. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-agy | 2 +- bin/omarchy-agent-usage-claude | 2 +- bin/omarchy-agent-usage-codex | 2 +- bin/omarchy-agent-usage-copilot | 2 +- bin/omarchy-agent-usage-cursor | 2 +- bin/omarchy-agent-usage-grok | 2 +- bin/omarchy-agent-usage-grokbot | 2 +- bin/omarchy-agent-usage-kimchi | 2 +- bin/omarchy-agent-usage-minimax | 4 ++-- bin/omarchy-agent-usage-muse | 2 +- bin/omarchy-agent-usage-ollama | 4 ++-- bin/omarchy-agent-usage-zai | 2 +- docs/agent-usage-collectors.md | 2 +- 13 files changed, 15 insertions(+), 15 deletions(-) diff --git a/bin/omarchy-agent-usage-agy b/bin/omarchy-agent-usage-agy index 6847dfbde85..a3a45064604 100755 --- a/bin/omarchy-agent-usage-agy +++ b/bin/omarchy-agent-usage-agy @@ -399,7 +399,7 @@ def collect_limits(entries, force): def answered(entry, tier, limits): # A plan with no quota windows to report is healthy, with no meters. - fetched = round(time.time() * 1000) + fetched = int(time.time() * 1000) result.update(limits=limits, tierLabel=tier, live=True, fetchedAtMs=fetched) path = limits_cache(entry) if path: diff --git a/bin/omarchy-agent-usage-claude b/bin/omarchy-agent-usage-claude index c5d4362c5e1..75349ae73f7 100755 --- a/bin/omarchy-agent-usage-claude +++ b/bin/omarchy-agent-usage-claude @@ -1149,7 +1149,7 @@ def collect_limits(access_token: str, expires_at_ms: int, refresh_expires_at_ms: result["limits"] = probe["limits"] result["live"] = True result["authHelpText"] = "" - result["fetchedAtMs"] = round(time.time() * 1000) + result["fetchedAtMs"] = int(time.time() * 1000) write_json(probe_cache, {"fetchedAtMs": result["fetchedAtMs"], "limits": probe["limits"]}) return result diff --git a/bin/omarchy-agent-usage-codex b/bin/omarchy-agent-usage-codex index 993a559b13f..eaca5e30ef7 100755 --- a/bin/omarchy-agent-usage-codex +++ b/bin/omarchy-agent-usage-codex @@ -1275,7 +1275,7 @@ def collect_limits(home=None, force=False): if ok: rpc["live"] = True - rpc["fetchedAtMs"] = round(time.time() * 1000) + rpc["fetchedAtMs"] = int(time.time() * 1000) if same_login and rpc["limits"]: try: write_json(cache, { diff --git a/bin/omarchy-agent-usage-copilot b/bin/omarchy-agent-usage-copilot index 7bb02f07c80..0506710bafd 100755 --- a/bin/omarchy-agent-usage-copilot +++ b/bin/omarchy-agent-usage-copilot @@ -897,7 +897,7 @@ def fetch_remote_quota(force: bool = False) -> dict[str, Any]: quota = ask_quota(token) if quota.get("available"): - quota.update(stale=False, fetchedAtMs=round(time.time() * 1000)) + quota.update(stale=False, fetchedAtMs=int(time.time() * 1000)) try: write_json(cache, quota) except OSError: diff --git a/bin/omarchy-agent-usage-cursor b/bin/omarchy-agent-usage-cursor index 5a67bfc445d..e1c5cb9c38c 100755 --- a/bin/omarchy-agent-usage-cursor +++ b/bin/omarchy-agent-usage-cursor @@ -497,7 +497,7 @@ def open_limits(limits: Any) -> list[dict[str, Any]]: def now_ms() -> int: - return round(datetime.now(timezone.utc).timestamp() * 1000) + return int(datetime.now(timezone.utc).timestamp() * 1000) def scan(api_base_url: str, auth_path: Path, limits_only: bool, state_db: Path | None = None, force: bool = False) -> dict[str, Any]: diff --git a/bin/omarchy-agent-usage-grok b/bin/omarchy-agent-usage-grok index 9eac6e1414f..ca76605525f 100755 --- a/bin/omarchy-agent-usage-grok +++ b/bin/omarchy-agent-usage-grok @@ -315,7 +315,7 @@ def collect_limits(home, key, force): result["authHelpText"] = "Grok's usage endpoint answered without a usable limit. Local Grok stats are still shown." if limit: - result.update(limits=[limit], live=True, fetchedAtMs=round(time.time() * 1000), authHelpText="") + result.update(limits=[limit], live=True, fetchedAtMs=int(time.time() * 1000), authHelpText="") write_json(cache, {"fetchedAtMs": result["fetchedAtMs"], "limits": result["limits"]}) else: result["limits"] = fallback diff --git a/bin/omarchy-agent-usage-grokbot b/bin/omarchy-agent-usage-grokbot index 4e7ff530470..bc4b4146c3b 100755 --- a/bin/omarchy-agent-usage-grokbot +++ b/bin/omarchy-agent-usage-grokbot @@ -46,7 +46,7 @@ def cache_root(): def now_ms(): - return round(datetime.now(timezone.utc).timestamp() * 1000) + return int(datetime.now(timezone.utc).timestamp() * 1000) def finite(value): diff --git a/bin/omarchy-agent-usage-kimchi b/bin/omarchy-agent-usage-kimchi index 3704a4da8c4..4b59b4ecb77 100755 --- a/bin/omarchy-agent-usage-kimchi +++ b/bin/omarchy-agent-usage-kimchi @@ -477,7 +477,7 @@ def collect_billing(force): continue try: data = get_json(f"{base}/v1/{resource}", key) - parts[resource] = dict(parse(data if isinstance(data, dict) else {}), fetchedAtMs=round(time.time() * 1000)) + parts[resource] = dict(parse(data if isinstance(data, dict) else {}), fetchedAtMs=int(time.time() * 1000)) remember(cache, parts[resource]) continue except urllib.error.HTTPError as error: diff --git a/bin/omarchy-agent-usage-minimax b/bin/omarchy-agent-usage-minimax index 7c9bd19327d..bef6855ec92 100755 --- a/bin/omarchy-agent-usage-minimax +++ b/bin/omarchy-agent-usage-minimax @@ -438,7 +438,7 @@ def collect_limits(api_key: str, base_url: str, force: bool) -> dict[str, Any]: # kept limits stand. failure = MiniMaxError(f"MiniMax's quota check failed ({type(error).__name__}).") else: - result.update(limits=limits, live=True, fetchedAtMs=round(time.time() * 1000)) + result.update(limits=limits, live=True, fetchedAtMs=int(time.time() * 1000)) try: write_json(cache_file, {"fetchedAtMs": result["fetchedAtMs"], "limits": limits}) except OSError as exc: @@ -821,7 +821,7 @@ def cached_local_stats(max_age_seconds: float) -> dict[str, Any]: # next rescan; only a complete one is kept. if complete: try: - write_json(cache_file, {"scanDate": datetime.now().strftime("%Y-%m-%d"), "scannedAtMs": round(time.time() * 1000), "stats": stats}) + write_json(cache_file, {"scanDate": datetime.now().strftime("%Y-%m-%d"), "scannedAtMs": int(time.time() * 1000), "stats": stats}) except OSError as exc: print(f"omarchy-agent-usage-minimax: could not write scan cache ({exc})", file=sys.stderr) return stats diff --git a/bin/omarchy-agent-usage-muse b/bin/omarchy-agent-usage-muse index 7421187ff9b..a3c66465fb3 100755 --- a/bin/omarchy-agent-usage-muse +++ b/bin/omarchy-agent-usage-muse @@ -757,7 +757,7 @@ def collect_limits(force: bool) -> dict: if probe["ok"]: result["limits"] = probe["limits"] result["tierLabel"] = probe.get("tierLabel") or result["tierLabel"] - result["fetchedAtMs"] = round(time.time() * 1000) + result["fetchedAtMs"] = int(time.time() * 1000) if probe_cache is not None: try: write_json(probe_cache, {"fetchedAtMs": result["fetchedAtMs"], "limits": probe["limits"], "tierLabel": result["tierLabel"]}) diff --git a/bin/omarchy-agent-usage-ollama b/bin/omarchy-agent-usage-ollama index d7764f009fb..03eb4a735ba 100755 --- a/bin/omarchy-agent-usage-ollama +++ b/bin/omarchy-agent-usage-ollama @@ -404,7 +404,7 @@ def cached_scan(max_age_seconds: float, model_prefix: str) -> dict[str, Any]: if complete: write_json(cache_file, { "scanDate": local_date_string(), - "scannedAtMs": round(time.time() * 1000), + "scannedAtMs": int(time.time() * 1000), "modelPrefix": model_prefix, "stats": stats, }) @@ -525,7 +525,7 @@ def collect_limits(access_key: str, force: bool) -> dict[str, Any]: # A plan lookup that failed keeps the plan last named, rather than # blanking a label the account still has. plan = cached_plan if plan is None else plan - result.update(limits=probe["limits"], tierLabel=plan, live=True, fetchedAtMs=round(time.time() * 1000)) + result.update(limits=probe["limits"], tierLabel=plan, live=True, fetchedAtMs=int(time.time() * 1000)) write_json(probe_cache, {"fetchedAtMs": result["fetchedAtMs"], "limits": probe["limits"], "plan": plan}) return result diff --git a/bin/omarchy-agent-usage-zai b/bin/omarchy-agent-usage-zai index bd8744d1187..78f9a41ee80 100755 --- a/bin/omarchy-agent-usage-zai +++ b/bin/omarchy-agent-usage-zai @@ -493,7 +493,7 @@ def collect_limits(api_key: str, platform: str, force: bool) -> dict[str, Any]: probe = fetch_limits(api_key, platform) if probe["ok"]: - result.update(limits=probe["limits"], tierLabel=probe["tierLabel"], live=True, fetchedAtMs=round(time.time() * 1000)) + result.update(limits=probe["limits"], tierLabel=probe["tierLabel"], live=True, fetchedAtMs=int(time.time() * 1000)) try: write_json(cache_file, { "fetchedAtMs": result["fetchedAtMs"], diff --git a/docs/agent-usage-collectors.md b/docs/agent-usage-collectors.md index d7067e58395..c849822cff8 100644 --- a/docs/agent-usage-collectors.md +++ b/docs/agent-usage-collectors.md @@ -52,7 +52,7 @@ The detail goes in `authHelpText`: what failed and the one thing to do about it. ## Limits - Each credential or account caches its own last good answer at `cache_root() / "-limits-.json"`, `{fetchedAtMs, limits, ...}`, where `cache_root()` is `$XDG_CACHE_HOME/omarchy/agent-usage` and `` is a short hash of who the credential signs in as (a JWT subject, an email, an account id) or of the credential itself. Another sign-in never shows this one's numbers. -- The cache is read whatever its file time, so a clock that went back doesn't lose it, and every age check is `0 <= age < N`. +- The cache is read whatever its file time, so a clock that went back doesn't lose it, and every age check is `0 <= age < N`. Stamps are floored (`int(time.time() * 1000)`), never rounded: a stamp half a millisecond ahead reads as a clock that went back, and the next run probes again. - A failed check keeps the cached limits with `limitsStale: true` and their `limitsFetchedAt`. The panel shows their age on hover. - A kept window whose `resetsAt` has passed comes back at `percent: 0` with `resetsAt: ""`, never dropped: a reset allowance is untouched, not unknown. A window with no reset time is kept as it is. Reset times are only ever the service's own, never computed. - A meter for a balance or allowance is a real ratio of the service's numbers, never estimated from local usage unless its label says `(est.)`. From 16b5e621d5e27512e8ef9b8c2cf142c2ad93ec35 Mon Sep 17 00:00:00 2001 From: Joao Date: Wed, 14 Oct 2026 23:02:00 -0400 Subject: [PATCH 088/546] Add a Devin usage collector Local stats from the Devin CLI's sessions.db and quota from the user-status cache the CLI writes, with the Devin mark for the agents panel. Squashed from the collector, test and asset commits of #8124. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-devin | 534 ++++++++++++++++++ shell/plugins/agents/assets/devin-light.svg | 3 + shell/plugins/agents/assets/devin.svg | 3 + .../shell.d/agent-usage-devin-scanner-test.sh | 326 +++++++++++ 4 files changed, 866 insertions(+) create mode 100755 bin/omarchy-agent-usage-devin create mode 100644 shell/plugins/agents/assets/devin-light.svg create mode 100644 shell/plugins/agents/assets/devin.svg create mode 100644 test/shell.d/agent-usage-devin-scanner-test.sh diff --git a/bin/omarchy-agent-usage-devin b/bin/omarchy-agent-usage-devin new file mode 100755 index 00000000000..63d29ad0440 --- /dev/null +++ b/bin/omarchy-agent-usage-devin @@ -0,0 +1,534 @@ +#!/usr/bin/python3 +# omarchy:summary=Print the Devin CLI usage record as JSON +# omarchy:args=[--force] [--limits-only] +# omarchy:hidden=true +"""Collect Devin CLI usage into one display-ready JSON record. + +Token stats come from the Devin CLI session database at +~/.local/share/devin/cli/sessions.db. Each assistant message carries +metadata.metrics with input/output/cache token counts and the generation +model, so the collector aggregates those by day and by model the same way +the Claude and Codex collectors do. + +Rate limits (daily and weekly quota percentages with reset times) come from +the user-status cache that the Devin TUI writes at +~/.cache/devin/cli/user_status.*.bin on every startup. The TUI fetches +these from the server via a gRPC call (GetUserStatus); the collector reads +the cached result so it doesn't need to make its own network call. +""" + +import argparse +import base64 +import datetime as dt +import fcntl +import hashlib +import json +import os +import sqlite3 +import sys +import tempfile +import time +from pathlib import Path +from typing import Any + +AGENT_ID = "devin" +AGENT_NAME = "Devin" + +# 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, however low +# refreshIntervalSec is set. --limits-only promises only fresh limits, so it +# may reuse a scan for up to 15 minutes. Devin's limits come from the TUI's +# user-status cache, not a live API call, so they are always read fresh. +SCAN_REUSE_SECONDS = 20 +LIMITS_ONLY_REUSE_SECONDS = 900 + +# Devin's quota comes from a cache the TUI writes, never a live fetch, and +# cloud sessions keep spending it meanwhile. Numbers this recent still +# describe the account; older ones are flagged stale with their fetch time. +LIMITS_CURRENT_ENOUGH_SECONDS = 900 + + +def data_dir() -> Path: + return Path(os.environ.get("DEVIN_DATA_DIR") or (Path.home() / ".local" / "share" / "devin" / "cli")) + + +def db_path() -> Path: + return data_dir() / "sessions.db" + + +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 cache_paths() -> tuple[Path, Path]: + digest = hashlib.sha1(str(data_dir()).encode("utf-8")).hexdigest()[:16] + root = cache_root() + return root / f"devin-scan-{digest}.json", root / f"devin-scan-{digest}.lock" + + +def date_string(value: dt.date) -> str: + return value.strftime("%Y-%m-%d") + + +def today_string() -> str: + return date_string(dt.datetime.now().date()) + + +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 number(value: Any) -> int: + try: + return int(value or 0) + except (TypeError, ValueError): + return 0 + + +def model_name(raw: Any) -> str: + value = str(raw or "devin") + return value if value else "devin" + + +def empty_bucket() -> dict[str, int]: + return { + "inputTokens": 0, + "outputTokens": 0, + "cacheReadInputTokens": 0, + "cacheCreationInputTokens": 0, + } + + +# ---------------------------------------------------------------- local scan + + +def run_scan(conn: sqlite3.Connection) -> dict[str, Any]: + """Run the full SQLite scan and return the stats dict for caching.""" + today = dt.datetime.now().date() + recent_dates = recent_date_strings() + recent: dict[str, int] = {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() + + # json_extract raises on malformed JSON, aborting the whole query — and the + # direct-scan fallback — on a single bad row. CASE only evaluates the + # matching branch, so invalid chat_message rows are skipped instead; a bare + # AND would not order json_valid before json_extract. + for row in conn.execute( + """ + SELECT + date(mn.created_at, 'unixepoch', 'localtime') AS day, + json_extract(mn.chat_message, '$.metadata.metrics.input_tokens') AS input_tokens, + json_extract(mn.chat_message, '$.metadata.metrics.output_tokens') AS output_tokens, + json_extract(mn.chat_message, '$.metadata.metrics.cache_read_tokens') AS cache_read_tokens, + json_extract(mn.chat_message, '$.metadata.metrics.cache_creation_tokens') AS cache_creation_tokens, + json_extract(mn.chat_message, '$.metadata.generation_model') AS model + FROM message_nodes mn + WHERE CASE WHEN json_valid(mn.chat_message) + THEN json_extract(mn.chat_message, '$.metadata.metrics') IS NOT NULL + ELSE 0 END + """ + ): + inp = number(row["input_tokens"]) + out = number(row["output_tokens"]) + cache_read = number(row["cache_read_tokens"]) + cache_write = number(row["cache_creation_tokens"]) + total = inp + out + cache_read + cache_write + if total <= 0: + continue + + model = model_name(row["model"]) + day = row["day"] or local_date_string() + + bucket = model_usage.setdefault(model, empty_bucket()) + bucket["inputTokens"] += inp + bucket["outputTokens"] += out + bucket["cacheReadInputTokens"] += cache_read + bucket["cacheCreationInputTokens"] += cache_write + + if day: + active_dates.add(day) + if day in recent: + recent[day] += total + if day == date_string(today): + today_by_model[model] = today_by_model.get(model, 0) + total + + today_prompts = 0 + today_sessions = 0 + total_prompts = 0 + + for row in conn.execute( + """ + SELECT + date(timestamp, 'unixepoch', 'localtime') AS day, + COUNT(*) AS prompts, + COUNT(DISTINCT session_id) AS sessions + FROM prompt_history + GROUP BY day + """ + ): + day = row["day"] or local_date_string() + prompts = number(row["prompts"]) + total_prompts += prompts + if day == date_string(today): + today_prompts = prompts + today_sessions = number(row["sessions"]) + + # Per-day distinct counts can't be summed: a session spanning midnight + # belongs to two days. All-time sessions need their own distinct count. + total_sessions = number( + conn.execute("SELECT COUNT(DISTINCT session_id) FROM prompt_history").fetchone()[0] + ) + + return { + "todayPrompts": today_prompts, + "todaySessions": today_sessions, + "todayTotalTokens": sum(today_by_model.values()), + "todayTokensByModel": today_by_model, + "recentDays": [{"date": day, "messageCount": recent[day]} for day in recent_dates], + "totalPrompts": total_prompts, + "totalSessions": total_sessions, + "activeDays": len(active_dates), + "activeDates": sorted(active_dates), + "modelUsage": model_usage, + } + + +def read_fresh_json(path: Path, max_age_seconds: int) -> Any: + if max_age_seconds <= 0 or not path.exists(): + return None + try: + 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: + 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") + tmp.chmod(0o644) + tmp.replace(path) + except BaseException: + tmp.unlink(missing_ok=True) + raise + + +def read_cached_stats(cache_file: Path, max_age_seconds: int) -> 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, + # not merely old, whatever its mtime says. + if cached.get("scanDate") != today_string(): + 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_string(), "stats": stats}) + except Exception as exc: + print(f"omarchy-agent-usage-devin: could not write usage cache ({exc})", file=sys.stderr) + + +def cached_scan(conn: sqlite3.Connection, force: bool, limits_only: bool) -> dict[str, Any]: + """Local stats, with the cache as a pure optimization. + + The cache must never take the collector down: any cache-layer failure + (unwritable cache root, lock errors, disk full) degrades to a direct scan + and a warning on stderr. The JSON record is the contract; the cache is not. + """ + max_age = 0 if force else (LIMITS_ONLY_REUSE_SECONDS if limits_only else SCAN_REUSE_SECONDS) + try: + return _cached_scan(conn, max_age) + except sqlite3.Error: + raise # a database failure is not a cache-layer problem + except Exception as exc: + print(f"omarchy-agent-usage-devin: cache unavailable ({exc}); scanning directly", file=sys.stderr) + return run_scan(conn) + + +def _cached_scan(conn: sqlite3.Connection, max_age: int) -> dict[str, Any]: + cache_file, lock_file = cache_paths() + + cached = read_cached_stats(cache_file, max_age) + if cached is not None: + return cached + + with lock_file.open("w") as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + cached = read_cached_stats(cache_file, max_age) + if cached is not None: + return cached + stats = run_scan(conn) + write_cached_stats(cache_file, stats) + return stats + + +# ------------------------------------------------------------------- limits +# +# Devin does not expose rate limits through a CLI command or REST API. The +# TUI fetches GetUserStatus from the server on startup and caches the +# protobuf response at ~/.cache/devin/cli/user_status.*.bin. The collector +# reads that cache — the same data the /usage slash command renders — so it +# doesn't need credentials or a network call of its own. + + +def user_status_dir() -> Path: + return Path(os.environ.get("DEVIN_CACHE_DIR") or (Path.home() / ".cache" / "devin" / "cli")) + + +def credentials_mtime() -> float | None: + # The CLI rewrites credentials.toml on sign-in, so a user-status file fetched + # before that write describes a previous account. The fetch can land a beat + # before the file flush, hence the small skew allowance. + try: + return (data_dir().parent / "credentials.toml").stat().st_mtime - 30 + except OSError: + return None + + +def _decode_varint(data: bytes, pos: int) -> tuple[int, int]: + result = 0 + shift = 0 + while True: + byte = data[pos] + result |= (byte & 0x7F) << shift + pos += 1 + if not (byte & 0x80): + break + shift += 7 + return result, pos + + +def _find_field(data: bytes, target_fn: int) -> bytes | int | None: + """Return the value of field `target_fn` in a protobuf message, or None.""" + pos = 0 + while pos < len(data): + try: + tag, pos = _decode_varint(data, pos) + fn = tag >> 3 + wt = tag & 0x07 + if wt == 0: # varint + val, pos = _decode_varint(data, pos) + if fn == target_fn: + return val + elif wt == 1: # 64-bit + pos += 8 + elif wt == 2: # length-delimited + length, pos = _decode_varint(data, pos) + val = data[pos:pos + length] + pos += length + if fn == target_fn: + return val + elif wt == 5: # 32-bit + pos += 4 + else: + break + except IndexError: + break + return None + + +def _iso_from_unix(unix_ts: int | None) -> str: + if not unix_ts or unix_ts <= 0: + return "" + return dt.datetime.fromtimestamp(unix_ts, tz=dt.timezone.utc).isoformat() + + +# A cached percentage outlives the TUI session that measured it, but only +# until its window rolls over: once a window has reset, the figure describes +# a period that is over, and a stale 69% would misreport a fresh allowance. +# A window with no reset time, or one that won't parse, is kept — an +# unreadable timestamp is no reason to throw away a real number. +def _limit_window_open(entry: dict[str, Any], now: dt.datetime) -> bool: + raw = str(entry.get("resetsAt") or "") + if raw == "": + return True + try: + resets_at = dt.datetime.fromisoformat(raw.replace("Z", "+00:00")) + except Exception: + return True + if resets_at.tzinfo is None: + resets_at = resets_at.replace(tzinfo=dt.timezone.utc) + return resets_at > now + + +def parse_quota() -> dict[str, Any]: + """Read the user-status caches and return limits + tier label. + + The PlanStatus nested message carries daily/weekly quota percentages and + reset timestamps. Multiple cache files may exist (one per sign-in session), + and the CLI exposes no active-identity marker. The freshest fetch wins, but + files fetched before the last credentials write belong to a signed-out + account: after a switch with no fetch since, nothing qualifies and no quota + is reported rather than the previous account's. Without a credentials file + to date the sign-in, stale data beats none. + """ + candidates: list[tuple[float, dict[str, Any]]] = [] + for path in sorted(user_status_dir().glob("user_status.*.bin")): + try: + cache = json.loads(path.read_text(encoding="utf-8")) + fetched_at = number(cache.get("fetched_at_secs")) + payload_b64 = cache.get("payload") + if not payload_b64: + continue + payload = base64.b64decode(payload_b64) + plan_status = _find_field(payload, 13) + if not isinstance(plan_status, (bytes, bytearray)): + continue + plan_info = _find_field(plan_status, 1) + plan_name = None + if isinstance(plan_info, (bytes, bytearray)): + raw_name = _find_field(plan_info, 2) + if isinstance(raw_name, (bytes, bytearray)): + plan_name = raw_name.decode("utf-8", errors="replace") + daily_remaining_pct = _find_field(plan_status, 14) + weekly_remaining_pct = _find_field(plan_status, 15) + daily_reset = _find_field(plan_status, 17) + weekly_reset = _find_field(plan_status, 18) + limits: list[dict[str, Any]] = [] + # Both fields are *remaining* percent from GetUserStatus; the panel + # expects *used*, so invert: 99% remaining -> 1% used. + if isinstance(daily_remaining_pct, int) and daily_remaining_pct >= 0: + limits.append({ + "label": "Daily", + "percent": (100 - daily_remaining_pct) / 100.0, + "resetsAt": _iso_from_unix(daily_reset if isinstance(daily_reset, int) else None), + }) + if isinstance(weekly_remaining_pct, int) and weekly_remaining_pct >= 0: + limits.append({ + "label": "Weekly", + "percent": (100 - weekly_remaining_pct) / 100.0, + "resetsAt": _iso_from_unix(weekly_reset if isinstance(weekly_reset, int) else None), + }) + now = dt.datetime.now(dt.timezone.utc) + limits = [e for e in limits if _limit_window_open(e, now)] + candidates.append((fetched_at, { + "tierLabel": plan_name or "", + "limits": limits, + })) + except Exception: + continue + cutoff = credentials_mtime() + if cutoff is not None: + candidates = [c for c in candidates if c[0] >= cutoff] + if not candidates: + return {} + fetched_at, quota = max(candidates, key=lambda c: c[0]) + quota["limitsFetchedAt"] = fetched_at * 1000 + quota["limitsStale"] = time.time() - fetched_at > LIMITS_CURRENT_ENOUGH_SECONDS + return quota + + +# -------------------------------------------------------------------- record + + +def base_record(**overrides: Any) -> dict[str, Any]: + record: dict[str, Any] = { + "schemaVersion": 1, + "id": AGENT_ID, + "name": AGENT_NAME, + "updatedAt": dt.datetime.now(dt.timezone.utc).isoformat(), + "ready": False, + "hasLocalStats": False, + "tierLabel": "", + "usageStatusText": "", + "authHelpText": "", + "limits": [], + "limitsStale": False, + "limitsFetchedAt": 0, + } + record.update({ + "todayPrompts": 0, + "todaySessions": 0, + "todayTotalTokens": 0, + "todayTokensByModel": {}, + "recentDays": [{"date": day, "messageCount": 0} for day in recent_date_strings()], + "totalPrompts": 0, + "totalSessions": 0, + "activeDays": 0, + "activeDates": [], + "modelUsage": {}, + }) + record.update(overrides) + return record + + +def scan(force: bool = False, limits_only: bool = False) -> dict[str, Any]: + quota = parse_quota() + + path = db_path() + if not path.is_file(): + if quota.get("limits"): + # A live status cache means the CLI ran under this account; the database + # is absent or lives elsewhere, so this is "no local stats", not "not + # installed". + return base_record(ready=True, **quota) + return base_record(usageStatusText="Devin CLI not installed", authHelpText="Install it with: curl -fsSL https://cli.devin.ai/install.sh | bash", **quota) + + conn = None + try: + conn = sqlite3.connect(f"file:{path}?mode=ro", uri=True) + conn.row_factory = sqlite3.Row + conn.execute("PRAGMA query_only = ON") + stats = cached_scan(conn, force, limits_only) + ready = number(stats.get("totalPrompts")) > 0 or len(quota.get("limits", [])) > 0 + return base_record(ready=ready, hasLocalStats=True, **stats, **quota) + except Exception as error: + # A locked or damaged history database must not hide quota already read + # from the user-status cache; it only costs the local stats. + detail = str(error) if isinstance(error, sqlite3.Error) else f"{type(error).__name__}: {error}" + print(f"omarchy-agent-usage-devin: {detail}", file=sys.stderr) + return base_record( + ready=len(quota.get("limits", [])) > 0, + usageStatusText="Devin usage unavailable", + authHelpText=detail, + **quota, + ) + finally: + if conn is not None: + conn.close() + + +def main() -> int: + parser = argparse.ArgumentParser(description="Print the Devin CLI usage record as JSON") + parser.add_argument("--force", action="store_true") + parser.add_argument("--limits-only", action="store_true") + args = parser.parse_args() + + try: + record = scan(force=args.force, limits_only=args.limits_only) + except sqlite3.Error as error: + record = base_record(usageStatusText="Devin usage unavailable", authHelpText=str(error)) + print(f"omarchy-agent-usage-devin: {error}", file=sys.stderr) + except Exception as error: + record = base_record(usageStatusText="Devin usage unavailable", authHelpText=f"{type(error).__name__}: {error}") + print(f"omarchy-agent-usage-devin: {type(error).__name__}: {error}", 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/assets/devin-light.svg b/shell/plugins/agents/assets/devin-light.svg new file mode 100644 index 00000000000..17611b9db50 --- /dev/null +++ b/shell/plugins/agents/assets/devin-light.svg @@ -0,0 +1,3 @@ + + + diff --git a/shell/plugins/agents/assets/devin.svg b/shell/plugins/agents/assets/devin.svg new file mode 100644 index 00000000000..e2a93b3bcf5 --- /dev/null +++ b/shell/plugins/agents/assets/devin.svg @@ -0,0 +1,3 @@ + + + diff --git a/test/shell.d/agent-usage-devin-scanner-test.sh b/test/shell.d/agent-usage-devin-scanner-test.sh new file mode 100644 index 00000000000..925277f1e5b --- /dev/null +++ b/test/shell.d/agent-usage-devin-scanner-test.sh @@ -0,0 +1,326 @@ +#!/bin/bash + +source "$(dirname "$0")/base-test.sh" + +require_command jq +require_command python3 + +TEST_HOME=$(mktemp -d) +trap 'rm -rf "$TEST_HOME"' EXIT + +# Every collector run must use the test cache root: without this, scans would +# write to (or read from) the real XDG cache when it is set in the environment. +export XDG_CACHE_HOME="$TEST_HOME/.cache" + +DEVIN_DATA_DIR="$TEST_HOME/.local/share/devin/cli" +mkdir -p "$DEVIN_DATA_DIR" + +now=$(date +%s) +today=$(date +%Y-%m-%d) +yesterday=$(date -d yesterday +%Y-%m-%d 2>/dev/null || date -v-1d +%Y-%m-%d) + +# Build a mock sessions.db with the schema the collector expects. +python3 - "$DEVIN_DATA_DIR/sessions.db" "$now" "$today" "$yesterday" <<'PY' +import json, sqlite3, sys + +db_path, now, today, yesterday = sys.argv[1], int(sys.argv[2]), sys.argv[3], sys.argv[4] + +conn = sqlite3.connect(db_path) +conn.execute("CREATE TABLE message_nodes (row_id INTEGER PRIMARY KEY, session_id TEXT, node_id INTEGER, parent_node_id INTEGER, chat_message TEXT, created_at INTEGER, metadata TEXT)") +conn.execute("CREATE TABLE prompt_history (id INTEGER PRIMARY KEY, content TEXT, timestamp INTEGER, session_id TEXT, is_shell INTEGER)") + +def assistant_msg(model, inp, out, cache_read=0, cache_write=0): + return json.dumps({ + "metadata": { + "metrics": { + "input_tokens": inp, + "output_tokens": out, + "cache_read_tokens": cache_read, + "cache_creation_tokens": cache_write, + }, + "generation_model": model, + } + }) + +# Today: two turns with glm-5-2, one with devin-pro +conn.execute("INSERT INTO message_nodes (session_id, node_id, chat_message, created_at) VALUES (?, 1, ?, ?)", + ("s1", assistant_msg("glm-5-2", 100, 20, 60, 0), now)) +conn.execute("INSERT INTO message_nodes (session_id, node_id, chat_message, created_at) VALUES (?, 2, ?, ?)", + ("s1", assistant_msg("glm-5-2", 80, 10, 50, 0), now)) +conn.execute("INSERT INTO message_nodes (session_id, node_id, chat_message, created_at) VALUES (?, 3, ?, ?)", + ("s1", assistant_msg("devin-pro", 200, 40, 0, 10), now)) + +# Yesterday: one turn +conn.execute("INSERT INTO message_nodes (session_id, node_id, chat_message, created_at) VALUES (?, 4, ?, ?)", + ("s2", assistant_msg("glm-5-2", 50, 5, 0, 0), now - 86400)) + +# A row with no metrics should be skipped +conn.execute("INSERT INTO message_nodes (session_id, node_id, chat_message, created_at) VALUES (?, 5, ?, ?)", + ("s1", json.dumps({"metadata": {}}), now)) + +# A malformed row must not abort the whole scan +conn.execute("INSERT INTO message_nodes (session_id, node_id, chat_message, created_at) VALUES (?, 6, ?, ?)", + ("s1", "{not json", now)) + +# Prompt history: 3 prompts today across 2 sessions, 2 yesterday — one prompt +# belongs to today's session s1, so per-day session counts sum to 4 while the +# true distinct total is 3 +conn.execute("INSERT INTO prompt_history (content, timestamp, session_id, is_shell) VALUES (?, ?, 's1', 0)", ("hello", now)) +conn.execute("INSERT INTO prompt_history (content, timestamp, session_id, is_shell) VALUES (?, ?, 's1', 0)", ("world", now)) +conn.execute("INSERT INTO prompt_history (content, timestamp, session_id, is_shell) VALUES (?, ?, 's3', 0)", ("test", now)) +conn.execute("INSERT INTO prompt_history (content, timestamp, session_id, is_shell) VALUES (?, ?, 's2', 0)", ("yesterday", now - 86400)) +conn.execute("INSERT INTO prompt_history (content, timestamp, session_id, is_shell) VALUES (?, ?, 's1', 0)", ("late night", now - 86400)) + +conn.commit() +conn.close() +PY + +result=$(HOME="$TEST_HOME" DEVIN_DATA_DIR="$DEVIN_DATA_DIR" \ + "$ROOT/bin/omarchy-agent-usage-devin") + +# Today: (100+20+60) + (80+10+50) + (200+40+0+10) = 570 +[[ $(jq -r '.todayTotalTokens' <<<"$result") == "570" ]] || + fail "Devin collector counts today's tokens across models" "$result" +pass "Devin collector counts today's tokens across models" + +# 3 prompts today across 2 sessions +[[ $(jq -r '.todayPrompts' <<<"$result") == "3" ]] || + fail "Devin collector counts today's prompts" "$result" +[[ $(jq -r '.todaySessions' <<<"$result") == "2" ]] || + fail "Devin collector counts today's sessions" "$result" +pass "Devin collector counts today's prompts and sessions" + +# Model usage (all-time): glm-5-2 = (100+80+50) input, (20+10+5) output, (60+50+0) cache_read, 0 cache_write +# devin-pro = 200 input, 40 output, 0 cache_read, 10 cache_write +[[ $(jq -c '.modelUsage["glm-5-2"]' <<<"$result") == '{"cacheCreationInputTokens":0,"cacheReadInputTokens":110,"inputTokens":230,"outputTokens":35}' ]] || + fail "Devin collector aggregates glm-5-2 model usage" "$result" +[[ $(jq -c '.modelUsage["devin-pro"]' <<<"$result") == '{"cacheCreationInputTokens":10,"cacheReadInputTokens":0,"inputTokens":200,"outputTokens":40}' ]] || + fail "Devin collector aggregates devin-pro model usage" "$result" +pass "Devin collector aggregates model usage correctly" + +# Total prompts = 5 (3 today + 2 yesterday); total sessions = 3 distinct +# (s1, s2, s3) — s1 spans midnight and must not count twice +[[ $(jq -r '.totalPrompts' <<<"$result") == "5" ]] || + fail "Devin collector counts total prompts" "$result" +[[ $(jq -r '.totalSessions' <<<"$result") == "3" ]] || + fail "Devin collector counts a midnight-spanning session once" "$result" +pass "Devin collector counts total prompts and dedupes cross-day sessions" + +# Active days = 2 (today + yesterday) +[[ $(jq -r '.activeDays' <<<"$result") == "2" ]] || + fail "Devin collector counts active days" "$result" +pass "Devin collector counts active days" + +# Recent days: today should have 570, yesterday should have 55 (50+5) +[[ $(jq -r '[.recentDays[] | select(.date == "'$today'") | .messageCount] | .[0]' <<<"$result") == "570" ]] || + fail "Devin collector fills today's recent day" "$result" +[[ $(jq -r '[.recentDays[] | select(.date == "'$yesterday'") | .messageCount] | .[0]' <<<"$result") == "55" ]] || + fail "Devin collector fills yesterday's recent day" "$result" +pass "Devin collector fills the seven-day token series" + +# No limits +[[ $(jq -c '.limits' <<<"$result") == "[]" ]] || + fail "Devin collector reports no limits" "$result" +pass "Devin collector reports no limits" + +# Identifies itself +[[ $(jq -r '.id' <<<"$result") == "devin" ]] || + fail "Devin collector identifies itself" "$result" +[[ $(jq -r '.name' <<<"$result") == "Devin" ]] || + fail "Devin collector names itself" "$result" +pass "Devin collector identifies itself" + +# Missing database +rm "$DEVIN_DATA_DIR/sessions.db" +result=$(HOME="$TEST_HOME" DEVIN_DATA_DIR="$DEVIN_DATA_DIR" \ + "$ROOT/bin/omarchy-agent-usage-devin") +[[ $(jq -r '.usageStatusText' <<<"$result") == "Devin CLI not installed" ]] || + fail "Devin collector reports missing database" "$result" +[[ $(jq -r '.authHelpText' <<<"$result") != "" ]] || + fail "Devin collector shows install hint for missing database" "$result" +pass "Devin collector reports missing database gracefully" + +# Caching: a second run without --force should reuse the cached scan +# (rebuild the DB since we deleted it above) +mkdir -p "$DEVIN_DATA_DIR" +python3 - "$DEVIN_DATA_DIR/sessions.db" "$now" <<'PY' +import json, sqlite3, sys +db_path, now = sys.argv[1], int(sys.argv[2]) +conn = sqlite3.connect(db_path) +conn.execute("CREATE TABLE message_nodes (row_id INTEGER PRIMARY KEY, session_id TEXT, node_id INTEGER, chat_message TEXT, created_at INTEGER, metadata TEXT)") +conn.execute("CREATE TABLE prompt_history (id INTEGER PRIMARY KEY, content TEXT, timestamp INTEGER, session_id TEXT, is_shell INTEGER)") +conn.execute("INSERT INTO message_nodes (session_id, node_id, chat_message, created_at) VALUES ('s1', 1, ?, ?)", + (json.dumps({"metadata": {"metrics": {"input_tokens": 10, "output_tokens": 5, "cache_read_tokens": 0, "cache_creation_tokens": 0}, "generation_model": "test"}}), now)) +conn.execute("INSERT INTO prompt_history (content, timestamp, session_id, is_shell) VALUES ('hi', ?, 's1', 0)", (now,)) +conn.commit() +conn.close() +PY + +# Clear the scan cache the earlier scenarios wrote so this run starts cold +rm -f "$TEST_HOME/.cache/omarchy/agent-usage/devin-scan-"* + +# First run populates the cache +result1=$(HOME="$TEST_HOME" DEVIN_DATA_DIR="$DEVIN_DATA_DIR" \ + "$ROOT/bin/omarchy-agent-usage-devin") +[[ $(jq -r '.todayTotalTokens' <<<"$result1") == "15" ]] || + fail "Devin collector first scan populates cache" "$result1" + +# Add more data after the cache is written +python3 - "$DEVIN_DATA_DIR/sessions.db" "$now" <<'PY' +import json, sqlite3, sys +db_path, now = sys.argv[1], int(sys.argv[2]) +conn = sqlite3.connect(db_path) +conn.execute("INSERT INTO message_nodes (session_id, node_id, chat_message, created_at) VALUES ('s1', 2, ?, ?)", + (json.dumps({"metadata": {"metrics": {"input_tokens": 100, "output_tokens": 50, "cache_read_tokens": 0, "cache_creation_tokens": 0}, "generation_model": "test"}}), now)) +conn.commit() +conn.close() +PY + +# Second run without --force should still serve the cached value (15, not 165) +result2=$(HOME="$TEST_HOME" DEVIN_DATA_DIR="$DEVIN_DATA_DIR" \ + "$ROOT/bin/omarchy-agent-usage-devin") +[[ $(jq -r '.todayTotalTokens' <<<"$result2") == "15" ]] || + fail "Devin collector serves cached scan without --force" "$result2" +pass "Devin collector serves cached scan without --force" + +# --force bypasses the cache and sees the new data (165) +result3=$(HOME="$TEST_HOME" DEVIN_DATA_DIR="$DEVIN_DATA_DIR" \ + "$ROOT/bin/omarchy-agent-usage-devin" --force) +[[ $(jq -r '.todayTotalTokens' <<<"$result3") == "165" ]] || + fail "Devin collector --force rescans past the cache" "$result3" +pass "Devin collector --force rescans past the cache" + +# Quota / limits from user-status cache +# Build a minimal protobuf PlanStatus: field 1 (PlanInfo) with field 2 (plan name "Pro"), +# field 14 (daily_remaining_pct=69), field 15 (weekly_remaining_pct=53), +# field 17 (daily_reset_unix), field 18 (weekly_reset_unix) +python3 - "$TEST_HOME/.cache/devin/cli" "$now" <<'PY' +import base64, json, os, struct, sys + +cache_dir, now = sys.argv[1], int(sys.argv[2]) +os.makedirs(cache_dir, exist_ok=True) + +def varint(val): + out = b"" + while True: + byte = val & 0x7F + val >>= 7 + if val: + out += bytes([byte | 0x80]) + else: + out += bytes([byte]) + break + return out + +def field(fn, data): + if isinstance(data, int): + return varint((fn << 3) | 0) + varint(data) + return varint((fn << 3) | 2) + varint(len(data)) + data + +# PlanInfo: field 2 = "Pro" +plan_info = field(2, b"Pro") +# PlanStatus: field 1 = PlanInfo, field 14 = 69 (remaining), field 15 = 53 (remaining), +# field 17 = now+3600, field 18 = now+86400*4 +plan_status = field(1, plan_info) + field(14, 69) + field(15, 53) + field(17, now + 3600) + field(18, now + 86400 * 4) +# Top-level: field 13 = PlanStatus +payload = field(13, plan_status) + +cache = { + "version": 1, + "identity_digest": "test123", + "fetched_at_secs": now, + "payload": base64.b64encode(payload).decode("ascii"), +} +with open(os.path.join(cache_dir, "user_status.test123.bin"), "w") as f: + json.dump(cache, f) +PY + +result=$(HOME="$TEST_HOME" DEVIN_DATA_DIR="$DEVIN_DATA_DIR" \ + "$ROOT/bin/omarchy-agent-usage-devin" --force) +[[ $(jq -r '.tierLabel' <<<"$result") == "Pro" ]] || + fail "Devin collector reads plan name from user-status cache" "$result" +[[ $(jq -r '.limits[0].label' <<<"$result") == "Daily" ]] || + fail "Devin collector reports daily limit label" "$result" +[[ $(jq -r '.limits[0].percent' <<<"$result") == "0.31" ]] || + fail "Devin collector reports daily limit percent (31% used = 69% remaining)" "$result" +[[ $(jq -r '.limits[1].label' <<<"$result") == "Weekly" ]] || + fail "Devin collector reports weekly limit label" "$result" +[[ $(jq -r '.limits[1].percent' <<<"$result") == "0.47" ]] || + fail "Devin collector reports weekly limit percent (47% used = 53% remaining)" "$result" +[[ $(jq -r '.limitsStale' <<<"$result") == "false" ]] || + fail "Devin collector marks a just-fetched status cache as current" "$result" +[[ $(jq -r '.limitsFetchedAt' <<<"$result") == "$((now * 1000))" ]] || + fail "Devin collector reports when the status cache was fetched" "$result" +pass "Devin collector parses quota from user-status cache" + +# Quota without a session database: the status cache alone means the CLI has +# run, so the record is ready with limits rather than "not installed" +rm "$DEVIN_DATA_DIR/sessions.db" +result=$(HOME="$TEST_HOME" DEVIN_DATA_DIR="$DEVIN_DATA_DIR" \ + "$ROOT/bin/omarchy-agent-usage-devin" --force) +[[ $(jq -r '.ready' <<<"$result") == "true" ]] || + fail "Devin collector stays ready on quota without a session database" "$result" +[[ $(jq -r '.limits | length' <<<"$result") == "2" ]] || + fail "Devin collector keeps limits without a session database" "$result" +[[ $(jq -r '.usageStatusText' <<<"$result") == "" ]] || + fail "Devin collector drops the install hint when quota proves the CLI ran" "$result" +pass "Devin collector reports quota without a session database" + +# Status caches fetched under a previous sign-in are dropped: once +# credentials.toml is newer than every fetch, no quota is reported rather than +# the signed-out account's +creds="$TEST_HOME/.local/share/devin/credentials.toml" +echo 'api_server_url = "https://api.devin.ai"' >"$creds" +touch -d "@$((now + 7200))" "$creds" +result=$(HOME="$TEST_HOME" DEVIN_DATA_DIR="$DEVIN_DATA_DIR" \ + "$ROOT/bin/omarchy-agent-usage-devin" --force) +[[ $(jq -r '.tierLabel' <<<"$result") == "" ]] || + fail "Devin collector drops quota fetched before the last sign-in" "$result" +[[ $(jq -c '.limits' <<<"$result") == "[]" ]] || + fail "Devin collector drops limits fetched before the last sign-in" "$result" +pass "Devin collector drops status caches from before the last sign-in" + +# A credentials file older than the caches leaves the freshest fetch winning +touch -d "@$((now - 7200))" "$creds" +result=$(HOME="$TEST_HOME" DEVIN_DATA_DIR="$DEVIN_DATA_DIR" \ + "$ROOT/bin/omarchy-agent-usage-devin" --force) +[[ $(jq -r '.tierLabel' <<<"$result") == "Pro" ]] || + fail "Devin collector serves quota fetched under the current credentials" "$result" +pass "Devin collector serves quota fetched under the current credentials" + +# A status cache older than the freshness window keeps its limits but is +# flagged stale, with its fetch time reported so the panel can warn +python3 - "$TEST_HOME/.cache/devin/cli/user_status.test123.bin" "$now" <<'PY' +import json, sys +path = sys.argv[1] +with open(path) as f: + cache = json.load(f) +cache["fetched_at_secs"] = int(sys.argv[2]) - 3600 +with open(path, "w") as f: + json.dump(cache, f) +PY + +result=$(HOME="$TEST_HOME" DEVIN_DATA_DIR="$DEVIN_DATA_DIR" \ + "$ROOT/bin/omarchy-agent-usage-devin" --force) +[[ $(jq -r '.limitsStale' <<<"$result") == "true" ]] || + fail "Devin collector flags an old status cache as stale" "$result" +[[ $(jq -r '.limitsFetchedAt' <<<"$result") == "$(( (now - 3600) * 1000 ))" ]] || + fail "Devin collector reports the stale cache's fetch time" "$result" +[[ $(jq -r '.limits | length' <<<"$result") == "2" ]] || + fail "Devin collector keeps limits from a stale cache" "$result" +pass "Devin collector flags old quota as stale with its fetch time" + +# A damaged session database must not hide working quota: the scan fails, but +# the limits already read from the user-status cache still render, with the +# database error reported alongside +rm -f "$TEST_HOME/.cache/omarchy/agent-usage/devin-scan-"* +echo "not a database" > "$DEVIN_DATA_DIR/sessions.db" +result=$(HOME="$TEST_HOME" DEVIN_DATA_DIR="$DEVIN_DATA_DIR" \ + "$ROOT/bin/omarchy-agent-usage-devin" --force) +[[ $(jq -r '.usageStatusText' <<<"$result") == "Devin usage unavailable" ]] || + fail "Devin collector reports an unreadable session database" "$result" +[[ $(jq -r '.ready' <<<"$result") == "true" ]] || + fail "Devin collector stays ready on quota with an unreadable database" "$result" +[[ $(jq -r '.limits | length' <<<"$result") == "2" ]] || + fail "Devin collector keeps quota when the session database is unreadable" "$result" +pass "Devin collector keeps quota when the session database is unreadable" From 4d6b189d82633a84815c9011f76a896e5997ffb2 Mon Sep 17 00:00:00 2001 From: aserper Date: Wed, 14 Oct 2026 23:02:00 -0400 Subject: [PATCH 089/546] Add a Synthetic quota collector Quotas from Synthetic's /v2/quotas endpoint, with the Synthetic mark for the agents panel. The collector, test and asset parts of #7051. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-synthetic | 317 ++++++++++++++++++ .../plugins/agents/assets/synthetic-light.svg | 13 + shell/plugins/agents/assets/synthetic.svg | 13 + test/shell.d/agent-usage-synthetic-test.sh | 154 +++++++++ 4 files changed, 497 insertions(+) create mode 100755 bin/omarchy-agent-usage-synthetic create mode 100644 shell/plugins/agents/assets/synthetic-light.svg create mode 100644 shell/plugins/agents/assets/synthetic.svg create mode 100755 test/shell.d/agent-usage-synthetic-test.sh diff --git a/bin/omarchy-agent-usage-synthetic b/bin/omarchy-agent-usage-synthetic new file mode 100755 index 00000000000..b101706e37d --- /dev/null +++ b/bin/omarchy-agent-usage-synthetic @@ -0,0 +1,317 @@ +#!/usr/bin/python3 +# omarchy:summary=Print the Synthetic Labs usage record as JSON +# omarchy:args=[--force] [--limits-only] +# omarchy:hidden=true +"""Collect Synthetic Labs quotas into one display-ready JSON record. + +The account-global limits come from Synthetic's /v2/quotas endpoint. There +are no local transcript stats: the agents panel displays the weekly credit, +rolling request, and search allowances this collector prints. +""" + +from __future__ import annotations + +import argparse +import json +import math +import os +import sys +import urllib.error +import urllib.request +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +AGENT_ID = "synthetic" +AGENT_NAME = "Synthetic Labs" +API_BASE_URL = "https://api.synthetic.new" +AUTH_HELP = "Set SYNTHETIC_API_KEY or sign in to Synthetic in Pi or OpenCode." + + +class SyntheticError(Exception): + def __init__(self, message: str, retry_advised: bool = False): + super().__init__(message) + self.retry_advised = retry_advised + + +def utc_now() -> str: + return datetime.now(timezone.utc).isoformat() + + +def finite_number(value: Any) -> float | None: + try: + parsed = float(value) + return parsed if math.isfinite(parsed) else None + except (TypeError, ValueError): + return None + + +def number(value: Any, fallback: float = 0.0) -> float: + parsed = finite_number(value) + if parsed is None: + return fallback + return parsed + + +def ratio(used: Any, limit: Any) -> float | None: + used_value = finite_number(used) + total = finite_number(limit) + if used_value is None or total is None or total <= 0: + return None + return max(0.0, min(1.0, used_value / total)) + + +def remaining_ratio(remaining: Any, maximum: Any) -> float | None: + remaining_value = finite_number(remaining) + total = finite_number(maximum) + if remaining_value is None or total is None or total <= 0: + return None + return max(0.0, min(1.0, 1.0 - remaining_value / total)) + + +def reset_at(value: Any) -> str: + return str(value or "").strip() + + +def format_quantity(value: Any) -> str: + parsed = number(value) + if parsed.is_integer(): + return f"{int(parsed):,}" + return f"{parsed:,.2f}".rstrip("0").rstrip(".") + + +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": utc_now(), + "ready": False, + "hasLocalStats": False, + "hasPromptStats": False, + "scope": "account", + "tierLabel": "Subscription", + "usageStatusText": "", + "authHelpText": "", + "limits": [], + } + record.update(empty_stats()) + record.update(overrides) + return record + + +def stored_key(path: Path) -> str: + try: + payload = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + return "" + entry = payload.get("synthetic") if isinstance(payload, dict) else None + if not isinstance(entry, dict): + return "" + return str(entry.get("key") or "").strip() + + +def credentials() -> str: + explicit = str(os.environ.get("SYNTHETIC_API_KEY", "")).strip() + if explicit: + return explicit + # Pi and OpenCode use the same provider id and {"key": ...} auth shape. + pi_key = stored_key(Path.home() / ".pi" / "agent" / "auth.json") + if pi_key: + return pi_key + data_home = Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) + return stored_key(data_home / "opencode" / "auth.json") + + +def quota_limit(label: str, payload: Any) -> dict[str, Any] | None: + if not isinstance(payload, dict): + return None + used = payload.get("requests") + maximum = payload.get("limit") + percent = ratio(used, maximum) + if percent is None: + return None + return { + "label": label, + "title": label, + "percent": percent, + "detail": f"{format_quantity(used)} / {format_quantity(maximum)} used", + "resetsAt": reset_at(payload.get("renewsAt")), + } + + +def parse_limits(payload: dict[str, Any]) -> tuple[list[dict[str, Any]], str, str]: + limits: list[dict[str, Any]] = [] + + rolling = payload.get("rollingFiveHourLimit") + weekly = payload.get("weeklyTokenLimit") + has_modern_limit = False + limited = False + if isinstance(rolling, dict): + percent = remaining_ratio(rolling.get("remaining"), rolling.get("max")) + if percent is not None: + has_modern_limit = True + maximum = number(rolling.get("max")) + used = max(0.0, min(maximum, maximum - number(rolling.get("remaining")))) + limits.append({ + "label": "Requests used / 5h", + "title": "Requests used / 5h", + "percent": percent, + "detail": f"{format_quantity(used)} / {format_quantity(maximum)} used", + "resetsAt": reset_at(rolling.get("nextTickAt")), + "resetLabel": "Next tick in", + }) + limited = rolling.get("limited") is True + + weekly_remaining = None + weekly_percent = finite_number(weekly.get("percentRemaining")) if isinstance(weekly, dict) else None + if isinstance(weekly, dict) and weekly_percent is not None: + has_modern_limit = True + remaining_percent = max(0.0, min(100.0, weekly_percent)) + weekly_remaining = remaining_percent + remaining_credits = str(weekly.get("remainingCredits") or "").strip() + maximum_credits = str(weekly.get("maxCredits") or "").strip() + detail = ( + f"{remaining_credits} / {maximum_credits} left" + if remaining_credits and maximum_credits + else f"{format_quantity(remaining_percent)}% left" + ) + limits.append({ + "label": "Credits used / week", + "title": "Credits used / week", + "percent": 1.0 - remaining_percent / 100.0, + "detail": detail, + "resetsAt": reset_at(weekly.get("nextRegenAt")), + "resetLabel": "Next regen in", + }) + + # The endpoint documented only this bucket originally. Modern accounts + # still receive it at zero alongside the credit/request system, so use it + # only when neither modern allowance parsed successfully. + if not has_modern_limit: + subscription = quota_limit("Subscription requests", payload.get("subscription")) + if subscription: + limits.insert(0, subscription) + + search = payload.get("search") + if isinstance(search, dict): + hourly = quota_limit("Search used / hour", search.get("hourly")) + if hourly: + limits.append(hourly) + + free_tools = quota_limit("Free tool calls", payload.get("freeToolCalls")) + if free_tools: + limits.append(free_tools) + + if limited: + return limits, "Rolling limit reached", "Synthetic reports that the rolling five-hour limit is currently active." + if weekly_remaining is not None and weekly_remaining <= 0: + return limits, "Weekly credits exhausted", "Synthetic reports no weekly credits remaining." + return limits, "", "" + + +class SyntheticClient: + def __init__(self, api_key: str, base_url: str = API_BASE_URL): + self.api_key = api_key + self.base_url = base_url.rstrip("/") + + def quotas(self) -> dict[str, Any]: + request = urllib.request.Request( + self.base_url + "/v2/quotas", + headers={ + "Authorization": "Bearer " + self.api_key, + "Accept": "application/json", + "User-Agent": "omarchy-agents-synthetic/1.0", + }, + ) + try: + with urllib.request.urlopen(request, timeout=15) as response: + payload = json.load(response) + except urllib.error.HTTPError as error: + if error.code == 401: + raise SyntheticError("Synthetic rejected the API key") from error + if error.code == 403: + raise SyntheticError("The Synthetic API key cannot read quota data") from error + if error.code == 429: + raise SyntheticError("Synthetic is rate limiting quota checks") from error + raise SyntheticError(f"Synthetic quota API returned HTTP {error.code}") from error + except urllib.error.URLError as error: + raise SyntheticError("Could not reach the Synthetic quota API", retry_advised=True) from error + except TimeoutError as error: + raise SyntheticError("Synthetic quota request timed out", retry_advised=True) from error + except json.JSONDecodeError as error: + raise SyntheticError("Synthetic returned an invalid quota response") from error + if not isinstance(payload, dict): + raise SyntheticError("Synthetic returned an invalid quota response") + return payload + + +def scan(base_url: str) -> dict[str, Any]: + api_key = credentials() + if not api_key: + return base_record(usageStatusText="Waiting for auth", authHelpText=AUTH_HELP) + + payload = SyntheticClient(api_key, base_url).quotas() + limits, status, help_text = parse_limits(payload) + if not limits: + return base_record( + usageStatusText="Synthetic quota unavailable", + authHelpText="The quota endpoint returned no supported limits.", + ) + tier_label = "Subscription" + weekly = payload.get("weeklyTokenLimit") + if isinstance(weekly, dict) and str(weekly.get("remainingCredits") or "").strip(): + tier_label = str(weekly["remainingCredits"]).strip() + " weekly credits left" + return base_record( + ready=True, + tierLabel=tier_label, + usageStatusText=status, + authHelpText=help_text, + limits=limits, + ) + + +def main() -> int: + parser = argparse.ArgumentParser(description="Print Synthetic quota status as JSON") + # The shared updater passes these flags to every collector. Synthetic's + # quota request is inexpensive and does not count against subscription limits. + parser.add_argument("--force", action="store_true") + parser.add_argument("--limits-only", action="store_true") + parser.parse_args() + + try: + record = scan(os.environ.get("SYNTHETIC_API_BASE_URL", API_BASE_URL)) + except SyntheticError as error: + record = base_record( + usageStatusText="Synthetic unavailable", + authHelpText=str(error), + retryAdvised=error.retry_advised, + ) + except Exception as error: + record = base_record( + usageStatusText="Synthetic unavailable", + authHelpText="Synthetic quota scan failed.", + ) + print(f"omarchy-agent-usage-synthetic: {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/assets/synthetic-light.svg b/shell/plugins/agents/assets/synthetic-light.svg new file mode 100644 index 00000000000..5a8ae34ab38 --- /dev/null +++ b/shell/plugins/agents/assets/synthetic-light.svg @@ -0,0 +1,13 @@ + + + + + + + + + + + + + \ No newline at end of file diff --git a/shell/plugins/agents/assets/synthetic.svg b/shell/plugins/agents/assets/synthetic.svg new file mode 100644 index 00000000000..8a4eb744b7d --- /dev/null +++ b/shell/plugins/agents/assets/synthetic.svg @@ -0,0 +1,13 @@ + + + + + + + + + + + + + \ No newline at end of file diff --git a/test/shell.d/agent-usage-synthetic-test.sh b/test/shell.d/agent-usage-synthetic-test.sh new file mode 100755 index 00000000000..eeecb267a10 --- /dev/null +++ b/test/shell.d/agent-usage-synthetic-test.sh @@ -0,0 +1,154 @@ +#!/bin/bash + +set -euo pipefail + +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" + +require_command jq +require_command python3 + +TEST_HOME=$(mktemp -d) +trap 'rm -rf "$TEST_HOME"' EXIT + +collector="$ROOT/bin/omarchy-agent-usage-synthetic" + +no_auth=$(HOME="$TEST_HOME" XDG_DATA_HOME="$TEST_HOME/.local/share" SYNTHETIC_API_KEY="" "$collector") +[[ $(jq -r '.id + ":" + (.ready | tostring) + ":" + (.hasLocalStats | tostring)' <<<"$no_auth") == "synthetic:false:false" ]] || + fail "Synthetic collector prints a valid record without credentials" "$no_auth" +[[ $(jq -r '.authHelpText' <<<"$no_auth") == *"SYNTHETIC_API_KEY"* ]] || + fail "Synthetic collector explains how to authenticate" "$no_auth" +pass "Synthetic collector handles missing credentials" + +result=$(python3 - "$collector" "$TEST_HOME" <<'PY' +import importlib.machinery +import importlib.util +import json +import os +import sys +import urllib.error +from pathlib import Path + +collector_path = sys.argv[1] +home = Path(sys.argv[2]) +os.environ["HOME"] = str(home) +os.environ["XDG_DATA_HOME"] = str(home / ".local" / "share") + +loader = importlib.machinery.SourceFileLoader("synthetic_collector", collector_path) +spec = importlib.util.spec_from_loader(loader.name, loader) +scanner = importlib.util.module_from_spec(spec) +loader.exec_module(scanner) + +pi_auth = home / ".pi" / "agent" / "auth.json" +opencode_auth = home / ".local" / "share" / "opencode" / "auth.json" +pi_auth.parent.mkdir(parents=True, exist_ok=True) +opencode_auth.parent.mkdir(parents=True, exist_ok=True) +pi_auth.write_text(json.dumps({"synthetic": {"key": "pi-key"}})) +opencode_auth.write_text(json.dumps({"synthetic": {"key": "opencode-key"}})) + +os.environ["SYNTHETIC_API_KEY"] = "env-key" +env_key = scanner.credentials() +del os.environ["SYNTHETIC_API_KEY"] +pi_key = scanner.credentials() +pi_auth.unlink() +opencode_key = scanner.credentials() + +modern_payload = { + "subscription": {"requests": 0, "limit": 1000, "renewsAt": "2026-08-17T00:00:00Z"}, + "rollingFiveHourLimit": { + "remaining": 998.3, + "max": 1000, + "limited": False, + "nextTickAt": "2026-08-16T02:00:00Z", + }, + "weeklyTokenLimit": { + "percentRemaining": 64.69978866666666, + "remainingCredits": "$31.05", + "maxCredits": "$48.00", + "nextRegenAt": "2026-08-16T04:00:00Z", + }, + "search": { + "hourly": {"requests": 4, "limit": 250, "renewsAt": "2026-08-16T03:00:00Z"}, + }, +} +modern, status, help_text = scanner.parse_limits(modern_payload) +legacy, _, _ = scanner.parse_limits({ + "subscription": {"requests": 27, "limit": 135, "renewsAt": "2026-08-17T00:00:00Z"}, +}) +invalid_modern, _, _ = scanner.parse_limits({ + "rollingFiveHourLimit": {}, + "subscription": {"requests": 10, "limit": 100, "renewsAt": "2026-08-17T00:00:00Z"}, +}) + +captured = {} +class Response: + def __enter__(self): + return self + + def __exit__(self, *args): + return False + + def read(self, *args): + return json.dumps(modern_payload).encode() + + +def fake_urlopen(request, timeout): + captured["url"] = request.full_url + captured["authorization"] = request.get_header("Authorization") + captured["timeout"] = timeout + return Response() + +real_urlopen = scanner.urllib.request.urlopen +scanner.urllib.request.urlopen = fake_urlopen +fetched = scanner.SyntheticClient("secret", "https://synthetic.test").quotas() + +scanner.urllib.request.urlopen = lambda *_args, **_kwargs: (_ for _ in ()).throw( + urllib.error.URLError("offline") +) +try: + scanner.SyntheticClient("secret").quotas() + retry_advised = None +except scanner.SyntheticError as error: + retry_advised = error.retry_advised +finally: + scanner.urllib.request.urlopen = real_urlopen + +print(json.dumps({ + "credentials": [env_key, pi_key, opencode_key], + "modern": modern, + "modernStatus": [status, help_text], + "legacy": legacy, + "invalidModern": invalid_modern, + "request": captured, + "fetchedMatches": fetched == modern_payload, + "retryAdvised": retry_advised, +})) +PY +) + +[[ $(jq -c '.credentials' <<<"$result") == '["env-key","pi-key","opencode-key"]' ]] || + fail "Synthetic credentials follow environment, Pi, then OpenCode priority" "$result" +pass "Synthetic credentials follow the documented priority" + +# Round percentages to 4 decimals so the assertion does not depend on exact +# IEEE-754 rendering across Python/jq versions. +[[ $(jq -c '[.modern[] | {label, percent: (.percent * 10000 | round / 10000), detail, resetLabel}]' <<<"$result") == '[{"label":"Requests used / 5h","percent":0.0017,"detail":"1.7 / 1,000 used","resetLabel":"Next tick in"},{"label":"Credits used / week","percent":0.353,"detail":"$31.05 / $48.00 left","resetLabel":"Next regen in"},{"label":"Search used / hour","percent":0.016,"detail":"4 / 250 used","resetLabel":null}]' ]] || + fail "Synthetic collector maps modern quota windows" "$result" +[[ $(jq -r '.modern | map(.label) | index("Subscription requests")' <<<"$result") == "null" ]] || + fail "Synthetic collector hides the stale legacy bucket on modern accounts" "$result" +pass "Synthetic collector maps modern quotas without the stale legacy bucket" + +[[ $(jq -c '.legacy[0] | {label,percent,detail}' <<<"$result") == '{"label":"Subscription requests","percent":0.2,"detail":"27 / 135 used"}' ]] || + fail "Synthetic collector supports the documented legacy subscription quota" "$result" +[[ $(jq -r '.invalidModern[0].label' <<<"$result") == "Subscription requests" ]] || + fail "Synthetic collector falls back when modern quota objects are unusable" "$result" +pass "Synthetic collector preserves legacy quota compatibility" + +[[ $(jq -r '.request.url' <<<"$result") == "https://synthetic.test/v2/quotas" ]] || + fail "Synthetic collector calls the documented quota endpoint" "$result" +[[ $(jq -r '.request.authorization' <<<"$result") == "Bearer secret" ]] || + fail "Synthetic collector authenticates quota requests" "$result" +[[ $(jq -r '.fetchedMatches' <<<"$result") == "true" ]] || + fail "Synthetic collector decodes quota responses" "$result" +[[ $(jq -r '.retryAdvised' <<<"$result") == "true" ]] || + fail "Synthetic collector advises one retry after transport failure" "$result" +pass "Synthetic collector handles the quota transport contract" From bd119b5962ee85a409ed94edf5d1b32bbee15e80 Mon Sep 17 00:00:00 2001 From: Jocelyn Knight Date: Wed, 14 Oct 2026 23:02:00 -0400 Subject: [PATCH 090/546] Add a LiteLLM usage collector The collector and test parts of #13052. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-litellm | 539 ++++++++++++++++++ .../agent-usage-litellm-scanner-test.sh | 314 ++++++++++ 2 files changed, 853 insertions(+) create mode 100755 bin/omarchy-agent-usage-litellm create mode 100755 test/shell.d/agent-usage-litellm-scanner-test.sh diff --git a/bin/omarchy-agent-usage-litellm b/bin/omarchy-agent-usage-litellm new file mode 100755 index 00000000000..e064b2ada72 --- /dev/null +++ b/bin/omarchy-agent-usage-litellm @@ -0,0 +1,539 @@ +#!/usr/bin/python3 +# omarchy:summary=Print the LiteLLM proxy usage record as JSON +# omarchy:args=[--force] [--limits-only] +# omarchy:hidden=true +"""Collect a LiteLLM proxy's spend logs into one display-ready JSON record. + +Token stats come from the proxy's daily activity endpoint for the last 30 +days. An admin key is asked for /global/activity, which is the whole +instance; a key that cannot read that route falls back to +/user/daily/activity, which is only the key's own user. A max_budget on the +key, or maxBudget in ~/.config/omarchy/agents/litellm.json, becomes the +prepaid balance the agents panel already draws. The panel only ever reads +the JSON this prints. +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +import urllib.error +import urllib.parse +import urllib.request +from datetime import date, datetime, timedelta, timezone +from pathlib import Path +from typing import Any + +AGENT_ID = "litellm" +AGENT_NAME = "LiteLLM" +AUTH_HELP = "Set LITELLM_BASE_URL and LITELLM_API_KEY, or write ~/.config/omarchy/agents/litellm.json." +WINDOW_DAYS = 30 + + +class LiteLLMError(Exception): + def __init__(self, message: str, status: int = 0) -> None: + super().__init__(message) + self.status = status + + +def number(value: Any) -> int: + try: + return max(0, round(float(value or 0))) + except (TypeError, ValueError): + return 0 + + +def money(value: Any) -> float: + try: + return max(0.0, float(value or 0)) + except (TypeError, ValueError): + return 0.0 + + +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, + # Proxy spend is account-global: every machine that reads the same proxy + # reports the same truth, so aggregation must not sum the records. + "scope": "account", + # api_requests are gateway attempts, not coding-agent prompts. The panel + # would print "0 sessions" next to them, so leave that line off. + "hasPromptStats": False, + "tierLabel": "Proxy", + "usageStatusText": "", + "authHelpText": "", + "limits": [], + } + record.update(empty_stats()) + record.update(overrides) + return record + + +def metric_block(node: Any) -> dict[str, Any]: + # Current proxies nest the numbers under "metrics". Older daily-activity + # payloads put prompt_tokens on the model object itself. + if not isinstance(node, dict): + return {} + nested = node.get("metrics") + if isinstance(nested, dict): + return nested + return node + + +def token_parts(metrics: dict[str, Any]) -> tuple[int, int, int, int, int]: + prompt = number(metrics.get("prompt_tokens")) + output = number(metrics.get("completion_tokens")) + cache_read = number(metrics.get("cache_read_input_tokens")) + cache_write = number(metrics.get("cache_creation_input_tokens")) + cache = cache_read + cache_write + # OpenAI-style rows include cache tokens inside prompt_tokens. Split them + # back out so the four panel buckets stay mutually exclusive. A row whose + # prompt is already smaller than its cache is left alone. + if cache and prompt >= cache: + uncached = prompt - cache + else: + uncached = prompt + total = uncached + cache_read + cache_write + output + return uncached, output, cache_read, cache_write, total + + +MONTHS = { + "jan": 1, "feb": 2, "mar": 3, "apr": 4, "may": 5, "jun": 6, + "jul": 7, "aug": 8, "sep": 9, "oct": 10, "nov": 11, "dec": 12, +} + + +def parse_activity_date(raw: Any, end: date) -> str: + text = str(raw or "").strip() + if len(text) >= 10 and text[4] == "-" and text[7] == "-": + return text[:10] + # Older /global/activity rows use "Aug 25" with no year. The query window + # is at most 30 days, so a month past the window's end belongs to last year. + parts = text.replace(",", " ").split() + if len(parts) < 2: + return "" + month = MONTHS.get(parts[0][:3].lower()) + if month is None: + return "" + try: + day = int(parts[1]) + candidate = date(end.year, month, day) + except ValueError: + return "" + if candidate > end: + try: + candidate = date(end.year - 1, month, day) + except ValueError: + return "" + return candidate.isoformat() + + +def summarize_activity(payload: dict[str, Any], today: date | None = None) -> dict[str, Any]: + today = today or datetime.now().astimezone().date() + rows = payload.get("results") + daily = payload.get("daily_data") + if isinstance(rows, list) and rows: + return summarize_results(payload, today) + if isinstance(daily, list): + return summarize_daily_data(payload, today) + return summarize_results(payload, today) + + +def summarize_daily_data(payload: dict[str, Any], today: date) -> dict[str, Any]: + # LiteLLM's older global activity route is a day series of request and + # token totals. Model rows for that shape come from /global/activity/model. + recent_dates = [(today - timedelta(days=offset)).isoformat() for offset in range(6, -1, -1)] + recent = {day: 0 for day in recent_dates} + active_dates: set[str] = set() + total_requests = 0 + today_requests = 0 + today_tokens = 0 + rows = payload.get("daily_data") + if not isinstance(rows, list): + rows = [] + for raw_row in rows: + if not isinstance(raw_row, dict): + continue + day = parse_activity_date(raw_row.get("date"), today) + tokens = number(raw_row.get("total_tokens")) + requests = number(raw_row.get("api_requests")) + total_requests += requests + if tokens <= 0 and requests <= 0: + continue + if day: + active_dates.add(day) + if day in recent: + recent[day] += tokens + if day == today.isoformat(): + today_requests += requests + today_tokens += tokens + return { + "todayPrompts": today_requests, + "todayTotalTokens": today_tokens, + "todayTokensByModel": {}, + "recentDays": [{"date": day, "messageCount": recent[day]} for day in recent_dates], + "totalPrompts": total_requests, + "activeDays": len(active_dates), + "activeDates": sorted(active_dates), + "modelUsage": {}, + "windowSpend": 0.0, + } + + +def summarize_results(payload: dict[str, Any], today: date) -> dict[str, Any]: + recent_dates = [(today - timedelta(days=offset)).isoformat() 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() + total_requests = 0 + today_requests = 0 + window_spend = 0.0 + + rows = payload.get("results") + if not isinstance(rows, list): + rows = [] + + for raw_row in rows: + if not isinstance(raw_row, dict): + continue + day = parse_activity_date(raw_row.get("date"), today) + metrics = metric_block(raw_row) + _, _, _, _, total = token_parts(metrics) + requests = number(metrics.get("api_requests") or metrics.get("successful_requests")) + window_spend += money(metrics.get("spend")) + total_requests += requests + if total <= 0 and requests <= 0: + continue + if day: + active_dates.add(day) + if day in recent: + recent[day] += total + if day == today.isoformat(): + today_requests += requests + + breakdown = raw_row.get("breakdown") if isinstance(raw_row.get("breakdown"), dict) else {} + models = breakdown.get("models") if isinstance(breakdown.get("models"), dict) else {} + if models: + for name, raw_model in models.items(): + model_metrics = metric_block(raw_model) + uncached, output, cache_read, cache_write, model_total = token_parts(model_metrics) + if model_total <= 0: + continue + bucket = model_usage.setdefault(str(name), empty_bucket()) + bucket["inputTokens"] += uncached + bucket["outputTokens"] += output + bucket["cacheReadInputTokens"] += cache_read + bucket["cacheCreationInputTokens"] += cache_write + if day == today.isoformat(): + today_by_model[str(name)] = today_by_model.get(str(name), 0) + model_total + elif day == today.isoformat() and total > 0: + today_by_model["unknown"] = today_by_model.get("unknown", 0) + total + + metadata = payload.get("metadata") if isinstance(payload.get("metadata"), dict) else {} + if money(metadata.get("total_spend")) > 0: + window_spend = money(metadata.get("total_spend")) + + return { + "todayPrompts": today_requests, + "todayTotalTokens": sum(today_by_model.values()), + "todayTokensByModel": today_by_model, + "recentDays": [{"date": day, "messageCount": recent[day]} for day in recent_dates], + "totalPrompts": total_requests, + "activeDays": len(active_dates), + "activeDates": sorted(active_dates), + "modelUsage": model_usage, + "windowSpend": window_spend, + } + + +def summarize_model_series(payload: dict[str, Any], today: date) -> dict[str, Any]: + # /global/activity/model on LiteLLM 1.82 reports one total per model. The + # panel's bar is the sum of the four buckets, so the combined total sits in + # inputTokens. Prompt, cache, and output stay zero because this route does + # not split them. + model_usage: dict[str, dict[str, int]] = {} + today_by_model: dict[str, int] = {} + rows = payload.get("results") if isinstance(payload.get("results"), list) else [] + today_key = today.isoformat() + for raw in rows: + if not isinstance(raw, dict): + continue + name = str(raw.get("model") or "").strip() + if not name: + continue + daily = raw.get("daily_data") if isinstance(raw.get("daily_data"), list) else [] + summed = 0 + today_tokens = 0 + for day_row in daily: + if not isinstance(day_row, dict): + continue + tokens = number(day_row.get("total_tokens")) + summed += tokens + if parse_activity_date(day_row.get("date"), today) == today_key: + today_tokens += tokens + total = number(raw.get("sum_total_tokens")) or summed + if total <= 0: + continue + bucket = empty_bucket() + bucket["inputTokens"] = total + model_usage[name] = bucket + if today_tokens > 0: + today_by_model[name] = today_tokens + return {"modelUsage": model_usage, "todayTokensByModel": today_by_model} + + +def config_path() -> Path: + root = os.environ.get("XDG_CONFIG_HOME") or str(Path.home() / ".config") + return Path(root) / "omarchy" / "agents" / "litellm.json" + + +def read_config() -> dict[str, Any]: + path = config_path() + if not path.is_file(): + return {} + try: + parsed = json.loads(path.read_text()) + except (OSError, json.JSONDecodeError): + return {} + return parsed if isinstance(parsed, dict) else {} + + +def normalize_base_url(raw: str) -> str: + url = raw.strip().rstrip("/") + if url.endswith("/v1"): + url = url[: -len("/v1")] + parsed = urllib.parse.urlparse(url) + if parsed.scheme not in ("http", "https") or not parsed.hostname: + raise LiteLLMError("LiteLLM base URL must be http or https") + if parsed.username or parsed.password: + raise LiteLLMError("LiteLLM base URL must not carry credentials") + return urllib.parse.urlunparse((parsed.scheme, parsed.netloc, parsed.path.rstrip("/"), "", "", "")) + + +def credentials(config: dict[str, Any]) -> tuple[str, str, float | None]: + base = str(os.environ.get("LITELLM_BASE_URL") or config.get("baseUrl") or "").strip() + key = str( + os.environ.get("LITELLM_API_KEY") + or os.environ.get("LITELLM_MASTER_KEY") + or config.get("apiKey") + or "" + ).strip() + budget_raw = config.get("maxBudget") + budget = money(budget_raw) if budget_raw not in (None, "") else None + if budget == 0: + budget = None + return base, key, budget + + +class SameOriginRedirectHandler(urllib.request.HTTPRedirectHandler): + def __init__(self, host: str) -> None: + super().__init__() + self.host = host + + def redirect_request(self, req, fp, code, msg, headers, newurl): + parsed = urllib.parse.urlparse(newurl) + if parsed.scheme not in ("http", "https") or parsed.hostname != self.host: + raise urllib.error.HTTPError(newurl, code, "cross-origin redirect refused", headers, fp) + return super().redirect_request(req, fp, code, msg, headers, newurl) + + +class LiteLLMClient: + def __init__(self, base_url: str, api_key: str) -> None: + self.base_url = normalize_base_url(base_url) + self.api_key = api_key + self.host = urllib.parse.urlparse(self.base_url).hostname or "" + + def get(self, path: str, query: dict[str, str] | None = None) -> dict[str, Any]: + url = self.base_url + path + if query: + url += "?" + urllib.parse.urlencode(query) + request = urllib.request.Request( + url, + headers={"Authorization": "Bearer " + self.api_key, "Accept": "application/json"}, + method="GET", + ) + opener = urllib.request.build_opener(SameOriginRedirectHandler(self.host)) + try: + with opener.open(request, timeout=20) as response: + final = getattr(response, "geturl", lambda: url)() + if urllib.parse.urlparse(final).hostname != self.host: + raise LiteLLMError("LiteLLM redirected off the configured host") + payload = json.loads(response.read().decode("utf-8", errors="replace")) + except urllib.error.HTTPError as error: + if error.code in (401, 403): + raise LiteLLMError("LiteLLM rejected the API key", error.code) from error + raise LiteLLMError(f"LiteLLM returned HTTP {error.code}", error.code) from error + except LiteLLMError: + raise + except Exception as error: + raise LiteLLMError("LiteLLM did not respond") from error + if isinstance(payload, list): + return {"results": payload} + if not isinstance(payload, dict): + raise LiteLLMError("LiteLLM returned an unexpected payload") + return payload + + def _paged(self, path: str, query: dict[str, str]) -> dict[str, Any]: + payload = self.get(path, query) + page = 1 + while page < 20: + meta = payload.get("metadata") if isinstance(payload.get("metadata"), dict) else {} + if not meta.get("has_more"): + break + page += 1 + nxt = self.get(path, {**query, "page": str(page)}) + rows = nxt.get("results") + if isinstance(rows, list): + payload.setdefault("results", []) + if isinstance(payload["results"], list): + payload["results"].extend(rows) + payload["metadata"] = nxt.get("metadata") if isinstance(nxt.get("metadata"), dict) else {} + return payload + + def _has_usage(self, payload: dict[str, Any]) -> bool: + if isinstance(payload.get("daily_data"), list): + return True + if isinstance(payload.get("results"), list) and payload["results"]: + return True + return False + + def activity(self, start_day: date, end_day: date) -> tuple[dict[str, Any], str]: + query = {"start_date": start_day.isoformat(), "end_date": end_day.isoformat()} + try: + payload = self._paged("/global/activity", query) + except LiteLLMError as error: + if error.status not in (403, 404): + raise + payload = None + if payload is not None and self._has_usage(payload): + return payload, "global" + return self._paged("/user/daily/activity", query), "user" + + def model_series(self, start_day: date, end_day: date) -> dict[str, Any] | None: + query = {"start_date": start_day.isoformat(), "end_date": end_day.isoformat()} + try: + return self.get("/global/activity/model", query) + except LiteLLMError: + return None + + def key_budget(self) -> tuple[float, float] | None: + try: + payload = self.get("/key/info") + except LiteLLMError: + return None + info = payload.get("info") if isinstance(payload.get("info"), dict) else payload + if not isinstance(info, dict) or info.get("max_budget") in (None, ""): + return None + funded = money(info.get("max_budget")) + if funded <= 0: + return None + return money(info.get("spend")), funded + + +def balance_from(spent: float, funded: float | None, estimated: bool) -> dict[str, Any] | None: + if funded is None or funded <= 0: + return None + return { + "remaining": round(max(0.0, funded - spent), 4), + "funded": round(funded, 4), + "spent": round(spent, 4), + "currency": "USD", + "estimated": estimated, + } + + +def scan(config: dict[str, Any] | None = None) -> dict[str, Any]: + config = read_config() if config is None else config + base_url, api_key, configured_budget = credentials(config) + if not base_url or not api_key: + return base_record(authHelpText=AUTH_HELP) + + try: + client = LiteLLMClient(base_url, api_key) + today = datetime.now().astimezone().date() + start = today - timedelta(days=WINDOW_DAYS - 1) + payload, scope_name = client.activity(start, today) + except LiteLLMError as error: + return base_record(usageStatusText="LiteLLM unavailable", authHelpText=str(error)) + + stats = summarize_activity(payload, today) + if scope_name == "global" and not stats["modelUsage"]: + series = client.model_series(start, today) + if series: + models = summarize_model_series(series, today) + if models["modelUsage"]: + stats["modelUsage"] = models["modelUsage"] + stats["todayTokensByModel"] = models["todayTokensByModel"] + window_spend = float(stats.pop("windowSpend")) + record = base_record( + ready=True, + hasLocalStats=stats["todayTotalTokens"] > 0 or stats["totalPrompts"] > 0 or bool(stats["modelUsage"]), + tierLabel="Proxy" if scope_name == "global" else "This key", + **stats, + ) + + key_budget = client.key_budget() + if key_budget is not None: + spent, funded = key_budget + ledger = balance_from(spent, funded, estimated=False) + else: + ledger = balance_from(window_spend, configured_budget, estimated=True) + if ledger is not None: + record["balance"] = ledger + return record + + +def main() -> int: + parser = argparse.ArgumentParser(description="Print the LiteLLM usage record as JSON") + # One activity request is the whole record, so there is no cache to force + # past and no faster limits-only path. The flags exist so every collector + # accepts the same invocation. + parser.add_argument("--force", action="store_true") + parser.add_argument("--limits-only", action="store_true") + parser.parse_args() + + try: + record = scan() + except LiteLLMError as error: + record = base_record(usageStatusText="LiteLLM unavailable", authHelpText=str(error)) + except Exception: + record = base_record(usageStatusText="LiteLLM unavailable", authHelpText="LiteLLM usage scan failed") + print("omarchy-agent-usage-litellm: scan failed", file=sys.stderr) + print(json.dumps(record, separators=(",", ":"))) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/test/shell.d/agent-usage-litellm-scanner-test.sh b/test/shell.d/agent-usage-litellm-scanner-test.sh new file mode 100755 index 00000000000..16eadf2ae4c --- /dev/null +++ b/test/shell.d/agent-usage-litellm-scanner-test.sh @@ -0,0 +1,314 @@ +#!/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/.config/omarchy/agents" + +# Without a URL and key the collector must still print a full record. The +# update runner writes whatever valid JSON appears on stdout, and the panel +# hides a record that has no usage. +no_key=$(HOME="$TEST_HOME" XDG_CONFIG_HOME="$TEST_HOME/.config" \ + env -u LITELLM_BASE_URL -u LITELLM_API_KEY -u LITELLM_MASTER_KEY \ + "$ROOT/bin/omarchy-agent-usage-litellm") + +[[ $(jq -r '.id + ":" + (.ready | tostring) + ":" + (.hasPromptStats | tostring) + ":" + .scope' <<<"$no_key") == "litellm:false:false:account" ]] || + fail "LiteLLM collector prints a valid record without credentials" "$no_key" +pass "LiteLLM collector prints a valid record without credentials" + +[[ $(jq -r '.authHelpText | contains("LITELLM_BASE_URL")' <<<"$no_key") == "true" ]] || + fail "LiteLLM collector tells the user how to configure the proxy" "$no_key" +pass "LiteLLM collector tells the user how to configure the proxy" + +result=$(python3 - "$ROOT/bin/omarchy-agent-usage-litellm" <<'PY' +import importlib.machinery +import importlib.util +import json +import os +import sys +import time +from datetime import date +from pathlib import Path + +os.environ["TZ"] = "UTC" +time.tzset() +for name in ("LITELLM_BASE_URL", "LITELLM_API_KEY", "LITELLM_MASTER_KEY"): + os.environ.pop(name, None) + +collector_path = str(Path(sys.argv[1])) +loader = importlib.machinery.SourceFileLoader("litellm_collector", collector_path) +spec = importlib.util.spec_from_loader(loader.name, loader) +scanner = importlib.util.module_from_spec(spec) +loader.exec_module(scanner) + +payload = { + "results": [ + { + "date": "2026-07-31", + "metrics": { + "spend": 0.4, + "prompt_tokens": 100, + "completion_tokens": 20, + "cache_read_input_tokens": 40, + "cache_creation_input_tokens": 10, + "api_requests": 3, + }, + "breakdown": { + "models": { + "ollama/qwen2.5": { + "metrics": { + "prompt_tokens": 100, + "completion_tokens": 20, + "cache_read_input_tokens": 40, + "cache_creation_input_tokens": 10, + } + } + } + }, + }, + { + "date": "2026-07-30T00:00:00Z", + "metrics": { + "prompt_tokens": 50, + "completion_tokens": 10, + "api_requests": 1, + }, + "breakdown": { + "models": { + "gpt-4": {"prompt_tokens": 50, "completion_tokens": 10} + } + }, + }, + ], + "metadata": {"total_spend": 1.25}, +} + +summary = scanner.summarize_activity(payload, date(2026, 7, 31)) +legacy = scanner.summarize_activity({ + "daily_data": [ + {"date": "Jul 31", "api_requests": 4, "total_tokens": 90}, + {"date": "Jul 30", "api_requests": 1, "total_tokens": 10}, + {"date": "Jan 02", "api_requests": 1, "total_tokens": 5}, + ], + "sum_total_tokens": 105, +}, date(2026, 7, 31)) +summary["legacyToday"] = legacy["todayTotalTokens"] +summary["legacyRequests"] = legacy["totalPrompts"] +summary["legacyModels"] = legacy["modelUsage"] +summary["legacyYear"] = scanner.parse_activity_date("Dec 28", date(2026, 1, 5)) == "2025-12-28" +series = scanner.summarize_model_series({ + "results": [ + { + "model": "ollama/qwen", + "sum_total_tokens": 90, + "daily_data": [{"date": "Jul 31", "total_tokens": 70, "api_requests": 2}], + }, + {"model": "", "sum_total_tokens": 10, "daily_data": []}, + {"model": "quiet", "sum_total_tokens": 0, "daily_data": []}, + ] +}, date(2026, 7, 31)) +summary["seriesTotal"] = series["modelUsage"]["ollama/qwen"]["inputTokens"] +summary["seriesSplit"] = series["modelUsage"]["ollama/qwen"]["outputTokens"] +summary["seriesToday"] = series["todayTokensByModel"]["ollama/qwen"] +summary["seriesSkipped"] = list(series["modelUsage"]) == ["ollama/qwen"] +summary["uncached"] = summary["modelUsage"]["ollama/qwen2.5"]["inputTokens"] +summary["cacheRead"] = summary["modelUsage"]["ollama/qwen2.5"]["cacheReadInputTokens"] +summary["flatModel"] = summary["modelUsage"]["gpt-4"]["outputTokens"] + +calls = [] +RealClient = scanner.LiteLLMClient + +class GlobalClient: + def __init__(self, base_url, api_key): + calls.append(("init", base_url, api_key)) + + def activity(self, start_day, end_day): + calls.append(("activity", start_day.isoformat(), end_day.isoformat())) + return payload, "global" + + def key_budget(self): + return (4.0, 20.0) + +class ForbiddenGlobal: + def __init__(self, base_url, api_key): + pass + + def activity(self, start_day, end_day): + raise scanner.LiteLLMError("LiteLLM rejected the API key", 403) + +class UserFallback(ForbiddenGlobal): + def activity(self, start_day, end_day): + return payload, "user" + + def key_budget(self): + return None + +scanner.LiteLLMClient = GlobalClient +record = scanner.scan({"baseUrl": "http://proxy.example:4000/v1", "apiKey": "sk-test"}) +summary["record"] = { + "id": record["id"], + "ready": record["ready"], + "scope": record["scope"], + "tierLabel": record["tierLabel"], + "hasPromptStats": record["hasPromptStats"], + "limits": record["limits"], + "balance": record["balance"], + "todayTotalTokens": record["todayTotalTokens"], +} +summary["normalizedBase"] = scanner.normalize_base_url("http://proxy.example:4000/v1") +start = date.fromisoformat(calls[1][1]) +end = date.fromisoformat(calls[1][2]) +summary["windowDays"] = (end - start).days == 29 and end == date.today() + +scanner.LiteLLMClient = UserFallback +user_record = scanner.scan({"baseUrl": "https://proxy.example", "apiKey": "sk-user", "maxBudget": 10}) +summary["userTier"] = user_record["tierLabel"] +summary["estimatedBalance"] = user_record["balance"] + +class LegacyModelClient: + def __init__(self, base_url, api_key): + pass + + def activity(self, start_day, end_day): + return {"daily_data": [{"date": "Jul 31", "api_requests": 4, "total_tokens": 90}]}, "global" + + def model_series(self, start_day, end_day): + return { + "results": [{ + "model": "ollama/qwen", + "sum_total_tokens": 90, + "daily_data": [{"date": "Jul 31", "total_tokens": 90}], + }] + } + + def key_budget(self): + return None + +scanner.LiteLLMClient = LegacyModelClient +legacy_record = scanner.scan({"baseUrl": "http://proxy.example:4000", "apiKey": "sk-old"}) +summary["legacyScanModel"] = legacy_record["modelUsage"]["ollama/qwen"]["inputTokens"] +summary["legacyScanReady"] = legacy_record["ready"] and legacy_record["tierLabel"] == "Proxy" + +refused = False +try: + scanner.normalize_base_url("http://user:secret@proxy.example:4000") +except scanner.LiteLLMError: + refused = True +summary["refusesCredentialsInUrl"] = refused + +# A 401 is the key itself, not a missing admin route, so activity() must not +# fall through to the user endpoint. +class RejectingGet: + def __init__(self): + self.paths = [] + + def get(self, path, query=None): + self.paths.append(path) + raise scanner.LiteLLMError("LiteLLM rejected the API key", 401) + +client = RealClient.__new__(RealClient) +rejecting = RejectingGet() +client.get = rejecting.get +try: + client.activity(date(2026, 7, 1), date(2026, 7, 31)) + summary["authStops"] = False +except scanner.LiteLLMError: + summary["authStops"] = rejecting.paths == ["/global/activity"] + +class FallbackGet: + def __init__(self): + self.paths = [] + + def get(self, path, query=None): + self.paths.append(path) + if path == "/global/activity": + raise scanner.LiteLLMError("LiteLLM rejected the API key", 403) + return payload + +client.get = FallbackGet().get +body, scope_name = client.activity(date(2026, 7, 1), date(2026, 7, 31)) +summary["fallsBack"] = scope_name == "user" and body["metadata"]["total_spend"] == 1.25 +print(json.dumps(summary, separators=(",", ":"))) +PY +) + +[[ $(jq -r '.todayTotalTokens' <<<"$result") == "120" ]] || + fail "LiteLLM collector totals today's uncached, cache, and output tokens once" "$result" +pass "LiteLLM collector totals today's uncached, cache, and output tokens once" + +[[ $(jq -r '.uncached' <<<"$result") == "50" ]] || + fail "LiteLLM collector splits cache tokens out of prompt tokens" "$result" +pass "LiteLLM collector splits cache tokens out of prompt tokens" + +[[ $(jq -r '.cacheRead' <<<"$result") == "40" ]] || + fail "LiteLLM collector keeps cache reads on their own counter" "$result" +pass "LiteLLM collector keeps cache reads on their own counter" + +[[ $(jq -r '.flatModel' <<<"$result") == "10" ]] || + fail "LiteLLM collector accepts model rows that are not nested under metrics" "$result" +pass "LiteLLM collector accepts model rows that are not nested under metrics" + +[[ $(jq -r '.recentDays[-1].messageCount' <<<"$result") == "120" ]] || + fail "LiteLLM collector builds the seven-day token series" "$result" +pass "LiteLLM collector builds the seven-day token series" + +[[ $(jq -r '.totalPrompts' <<<"$result") == "4" ]] || + fail "LiteLLM collector counts logged API requests" "$result" +pass "LiteLLM collector counts logged API requests" + +[[ $(jq -r '.windowSpend' <<<"$result") == "1.25" ]] || + fail "LiteLLM collector prefers the metadata spend total" "$result" +pass "LiteLLM collector prefers the metadata spend total" + +[[ $(jq -c '.record | {id, ready, scope, tierLabel, hasPromptStats, limits}' <<<"$result") == '{"id":"litellm","ready":true,"scope":"account","tierLabel":"Proxy","hasPromptStats":false,"limits":[]}' ]] || + fail "LiteLLM collector prints the display-ready record contract" "$result" +pass "LiteLLM collector prints the display-ready record contract" + +[[ $(jq -c '.record.balance' <<<"$result") == '{"remaining":16.0,"funded":20.0,"spent":4.0,"currency":"USD","estimated":false}' ]] || + fail "LiteLLM collector uses the key budget when the proxy reports one" "$result" +pass "LiteLLM collector uses the key budget when the proxy reports one" + +[[ $(jq -r '.normalizedBase' <<<"$result") == "http://proxy.example:4000" ]] || + fail "LiteLLM collector strips a trailing /v1 from the proxy URL" "$result" +pass "LiteLLM collector strips a trailing /v1 from the proxy URL" + +[[ $(jq -r '.windowDays' <<<"$result") == "true" ]] || + fail "LiteLLM collector asks for the last 30 days" "$result" +pass "LiteLLM collector asks for the last 30 days" + +[[ $(jq -r '.userTier' <<<"$result") == "This key" ]] || + fail "LiteLLM collector labels a non-admin key as its own usage" "$result" +pass "LiteLLM collector labels a non-admin key as its own usage" + +[[ $(jq -c '.estimatedBalance' <<<"$result") == '{"remaining":8.75,"funded":10.0,"spent":1.25,"currency":"USD","estimated":true}' ]] || + fail "LiteLLM collector estimates a balance from the configured budget" "$result" +pass "LiteLLM collector estimates a balance from the configured budget" + +[[ $(jq -r '.refusesCredentialsInUrl' <<<"$result") == "true" ]] || + fail "LiteLLM collector refuses a base URL that embeds a password" "$result" +pass "LiteLLM collector refuses a base URL that embeds a password" + +[[ $(jq -r '.authStops' <<<"$result") == "true" ]] || + fail "LiteLLM collector does not fall back when the key itself is rejected" "$result" +pass "LiteLLM collector does not fall back when the key itself is rejected" + +[[ $(jq -r '.fallsBack' <<<"$result") == "true" ]] || + fail "LiteLLM collector falls back to the user activity route" "$result" +pass "LiteLLM collector falls back to the user activity route" + +[[ $(jq -r '.legacyToday' <<<"$result") == "90" && $(jq -r '.legacyRequests' <<<"$result") == "6" && $(jq -r '.legacyYear' <<<"$result") == "true" && $(jq -c '.legacyModels' <<<"$result") == "{}" ]] || + fail "LiteLLM collector reads the older global activity day series" "$result" +pass "LiteLLM collector reads the older global activity day series" + +[[ $(jq -r '.seriesTotal' <<<"$result") == "90" && $(jq -r '.seriesSplit' <<<"$result") == "0" && $(jq -r '.seriesToday' <<<"$result") == "70" && $(jq -r '.seriesSkipped' <<<"$result") == "true" ]] || + fail "LiteLLM collector reads per-model totals from the model activity route" "$result" +pass "LiteLLM collector reads per-model totals from the model activity route" + +[[ $(jq -r '.legacyScanModel' <<<"$result") == "90" && $(jq -r '.legacyScanReady' <<<"$result") == "true" ]] || + fail "LiteLLM collector attaches model rows when the day series has none" "$result" +pass "LiteLLM collector attaches model rows when the day series has none" From 0862f7f7628cb60453b98e4e908a8a05c289a865 Mon Sep 17 00:00:00 2001 From: Nate Ashby Date: Wed, 14 Oct 2026 23:02:00 -0400 Subject: [PATCH 091/546] Add a Local collector for self-hosted model usage The collector and test parts of #12034; its Codex provider fix is already upstream. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-local | 533 ++++++++++++++++++ .../shell.d/agent-usage-local-scanner-test.sh | 199 +++++++ 2 files changed, 732 insertions(+) create mode 100755 bin/omarchy-agent-usage-local create mode 100755 test/shell.d/agent-usage-local-scanner-test.sh diff --git a/bin/omarchy-agent-usage-local b/bin/omarchy-agent-usage-local new file mode 100755 index 00000000000..44a96685626 --- /dev/null +++ b/bin/omarchy-agent-usage-local @@ -0,0 +1,533 @@ +#!/usr/bin/python3 +# omarchy:summary=Print the self-hosted model usage record as JSON +# omarchy:args=[--force] [--limits-only] +# omarchy:hidden=true +"""Collect self-hosted model usage into one display-ready JSON record. + +Every other collector answers "what did this subscription cost me". This one +answers the question none of them can: how much work ran on hardware you own. +It counts the turns the subscription collectors deliberately skip — Codex CLI +pointed at Ollama, opencode on a local gateway, pi on llama.cpp — whether the +endpoint is this machine or a box down the hall. + +An endpoint is self-hosted when its URL resolves to a loopback, private, or +link-local address, or to a bare/LAN hostname (`server.home`, `gpu-box`). +Provider ids with no URL to resolve fall back to a list of well-known local +runtimes. Neither is exhaustive, so `~/.config/omarchy/agents/local.json` can +name more: + + {"providers": ["homelab", "my-gateway"]} + +There are no limits and no balance to report: the meter on self-hosted work is +the electricity bill. +""" + +import argparse +import fcntl +import hashlib +import ipaddress +import json +import os +import sqlite3 +import sys +import tempfile +import time +import tomllib +from datetime import datetime, timedelta +from pathlib import Path +from urllib.parse import urlparse + +AGENT_ID = "local" +AGENT_NAME = "Local" +TIER_LABEL = "Self-hosted" + +# Runtimes that serve models from your own hardware. Used when a provider id +# carries no URL to classify — opencode and pi record the id but not the +# endpoint, and Codex's built-in `--oss` provider has no config.toml entry. +KNOWN_LOCAL_PROVIDERS = { + "gpt4all", "jan", "koboldcpp", "llama-cpp", "llamacpp", "llamafile", + "lm-studio", "lmstudio", "local", "localai", "mlx", "ollama", + "openai-local", "oss", "sglang", "tabby", "text-generation-webui", "vllm", +} + +# Session files are scanned back this far, matching the Codex collector's +# window. Pi transcripts carry no rollover, so they are read in full. +NATIVE_WINDOW_DAYS = 30 + +# Same reuse windows as the subscription collectors: a no-flag run reuses a +# scan only long enough to dedup the concurrent collectors the update command +# backgrounds, while --limits-only promises nothing about stats freshness and +# may reuse one for far longer. This collector has no limits to fetch, so +# --limits-only is purely a cheap way to redraw. +SCAN_REUSE_SECONDS = 20 +LIMITS_ONLY_REUSE_SECONDS = 900 + +today = datetime.now().strftime("%Y-%m-%d") + + +def local_day(value): + if value is None: + return datetime.now().strftime("%Y-%m-%d") + if isinstance(value, (int, float)): + # pi message timestamps are milliseconds; Codex timestamps are seconds. + if value > 10_000_000_000: + value = value / 1000 + return datetime.fromtimestamp(value).strftime("%Y-%m-%d") + text = str(value) + try: + if text.endswith("Z"): + dt = datetime.fromisoformat(text[:-1] + "+00:00") + else: + dt = datetime.fromisoformat(text) + if dt.tzinfo is not None: + dt = dt.astimezone() + return dt.strftime("%Y-%m-%d") + except Exception: + return datetime.now().strftime("%Y-%m-%d") + + +def number(value): + try: + return max(0, int(value or 0)) + except Exception: + return 0 + + +def model_name(value): + text = str(value or "").strip() + return text or "local" + + +def is_self_hosted(url): + """True when a base URL points at hardware the user controls.""" + try: + host = urlparse(str(url or "")).hostname + except Exception: + return False + if not host: + return False + host = host.lower().rstrip(".") + if host == "localhost": + return True + try: + ip = ipaddress.ip_address(host) + return ip.is_loopback or ip.is_private or ip.is_link_local + except ValueError: + pass + # A bare hostname is a LAN name by construction; the rest are the suffixes + # home routers and mDNS hand out. + if "." not in host: + return True + return host.endswith((".home", ".lan", ".local", ".internal", ".localdomain", ".home.arpa")) + + +def configured_providers(): + path = Path(os.environ.get("XDG_CONFIG_HOME") or (Path.home() / ".config")) / "omarchy" / "agents" / "local.json" + try: + data = json.loads(path.read_text()) + except Exception: + return set() + names = data.get("providers") if isinstance(data, dict) else None + if not isinstance(names, list): + return set() + return {str(name).strip().lower() for name in names if str(name).strip()} + + +def codex_endpoints(): + """provider id -> base_url, from the Codex CLI's own config.""" + codex_home = Path(os.environ.get("CODEX_HOME") or (Path.home() / ".codex")) + try: + with (codex_home / "config.toml").open("rb") as handle: + config = tomllib.load(handle) + except Exception: + return {} + providers = config.get("model_providers") + if not isinstance(providers, dict): + return {} + endpoints = {} + for name, entry in providers.items(): + if isinstance(entry, dict): + endpoints[str(name).lower()] = str(entry.get("base_url") or "") + return endpoints + + +class Scan: + def __init__(self): + self.extra_providers = configured_providers() + self.endpoints = codex_endpoints() + self.today = today + self.recent_days = [ + (datetime.now() - timedelta(days=offset)).strftime("%Y-%m-%d") for offset in range(6, -1, -1) + ] + self.recent = {day: 0 for day in self.recent_days} + self.today_prompts = 0 + self.today_sessions = set() + self.today_tokens = 0 + self.today_by_model = {} + self.total_prompts = 0 + self.total_sessions = set() + self.active_days = set() + self.model_usage = {} + self.providers_seen = set() + self.endpoints_seen = set() + self.seen_pi_messages = set() + self.complete = True + + def provider_is_local(self, provider): + pid = str(provider or "").strip().lower() + if not pid: + return False + if pid in self.extra_providers or pid in KNOWN_LOCAL_PROVIDERS: + return True + url = self.endpoints.get(pid) + return bool(url) and is_self_hosted(url) + + def local_provider_ids(self): + """Every id this scan would accept, for pushing filters down into SQL.""" + ids = set(KNOWN_LOCAL_PROVIDERS) | set(self.extra_providers) + for pid, url in self.endpoints.items(): + if url and is_self_hosted(url): + ids.add(pid) + return ids + + def note_provider(self, provider): + pid = str(provider or "").strip().lower() + self.providers_seen.add(pid) + url = self.endpoints.get(pid) + if url: + host = urlparse(url).hostname + if host: + self.endpoints_seen.add(host.lower()) + + def add(self, day, session_key, model, input_tokens, output_tokens, cache_read, cache_write): + total = input_tokens + output_tokens + cache_read + cache_write + if total <= 0: + return + self.total_prompts += 1 + self.total_sessions.add(session_key) + self.active_days.add(day) + if day in self.recent: + self.recent[day] += total + if day == self.today: + self.today_prompts += 1 + self.today_sessions.add(session_key) + self.today_tokens += total + self.today_by_model[model] = self.today_by_model.get(model, 0) + total + bucket = self.model_usage.setdefault( + model, + {"inputTokens": 0, "outputTokens": 0, "cacheReadInputTokens": 0, "cacheCreationInputTokens": 0}, + ) + bucket["inputTokens"] += input_tokens + bucket["outputTokens"] += output_tokens + bucket["cacheReadInputTokens"] += cache_read + bucket["cacheCreationInputTokens"] += cache_write + + # ------------------------------------------------------------ native codex + def scan_codex(self): + codex_home = Path(os.environ.get("CODEX_HOME") or (Path.home() / ".codex")) + cutoff = time.time() - NATIVE_WINDOW_DAYS * 24 * 60 * 60 + for root in (codex_home / "sessions", codex_home / "archived_sessions"): + if not root.exists(): + continue + for path in root.rglob("*.jsonl"): + try: + if path.stat().st_mtime < cutoff: + continue + except OSError: + continue + self._scan_codex_file(path) + + def _scan_codex_file(self, path): + current_model = "local" + try: + with path.open(errors="replace") as handle: + for raw in handle: + try: + entry = json.loads(raw) + except Exception: + continue + if entry.get("type") == "session_meta": + provider = (entry.get("payload") or {}).get("model_provider") + # The subscription collector counts these; this one must not. + if not self.provider_is_local(provider): + return + self.note_provider(provider) + continue + if entry.get("type") == "turn_context": + payload = entry.get("payload") or {} + current_model = model_name(payload.get("model") or payload.get("model_slug") or current_model) + continue + payload = entry.get("payload") or entry + if entry.get("type") == "response_item" and isinstance(payload, dict): + payload = payload.get("payload") or payload + if not isinstance(payload, dict) or payload.get("type") != "token_count": + continue + usage = (payload.get("info") or {}).get("last_token_usage") or {} + cache_read = number(usage.get("cached_input_tokens")) + cache_write = number(usage.get("cache_write_input_tokens")) + input_tokens = max(0, number(usage.get("input_tokens")) - cache_read - cache_write) + output_tokens = number(usage.get("output_tokens")) + day = local_day(entry.get("timestamp") or path.stat().st_mtime) + self.add(day, str(path), current_model, input_tokens, output_tokens, cache_read, cache_write) + except Exception: + return + + # ----------------------------------------------------------------- pi / omp + def scan_pi(self): + for root in (Path.home() / ".pi" / "agent" / "sessions", Path.home() / ".omp" / "agent" / "sessions"): + if not root.exists(): + continue + for path in root.rglob("*.jsonl"): + try: + with path.open(errors="replace") as handle: + for raw in handle: + try: + entry = json.loads(raw) + except Exception: + continue + message = entry.get("message") or {} + if not isinstance(message, dict) or message.get("role") != "assistant": + continue + provider = message.get("provider") or message.get("api") + if not self.provider_is_local(provider): + continue + key = str(entry.get("id") or entry.get("uuid") or f"{path}:{raw[:64]}") + if key in self.seen_pi_messages: + continue + self.seen_pi_messages.add(key) + self.note_provider(provider) + usage = message.get("usage") or {} + day = local_day(entry.get("timestamp") or message.get("timestamp")) + self.add( + day, str(path), model_name(message.get("model")), + number(usage.get("input")), number(usage.get("output")), + number(usage.get("cacheRead")), number(usage.get("cacheWrite")), + ) + except Exception: + continue + + # ----------------------------------------------------------------- opencode + def scan_opencode(self): + db = Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) / "opencode" / "opencode.db" + if not db.exists(): + return + try: + conn = sqlite3.connect(f"file:{db}?mode=ro&immutable=1", uri=True) + except Exception: + return + # OpenCode DBs hold every historical message with its JSON, so the + # provider filter belongs in SQL rather than in a Python loop over every + # row. The candidate ids are known before the query runs: the well-known + # runtimes, whatever the user configured, and any Codex provider whose + # endpoint resolves to hardware they host. json_valid() guards the + # extract so one malformed row cannot abort the scan. + candidates = sorted(self.local_provider_ids()) + if not candidates: + return + placeholders = ",".join("?" * len(candidates)) + try: + for session_id, raw in conn.execute( + "SELECT session_id, data FROM message" + " WHERE data LIKE '%\"role\"%:%\"assistant\"%'" + " 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" + f" IN ({placeholders})", + candidates, + ): + try: + entry = json.loads(raw) + except Exception: + continue + # Exact match again in Python: the SQL narrowed the rows, it did not + # vet them. + if not isinstance(entry, dict) or not self.provider_is_local(entry.get("providerID")): + continue + self.note_provider(entry.get("providerID")) + tokens = entry.get("tokens") or {} + cache = tokens.get("cache") or {} + created = (entry.get("time") or {}).get("created") + # opencode keeps thinking tokens out of output; both are generated. + self.add( + local_day(created), "opencode:" + str(session_id), model_name(entry.get("modelID")), + number(tokens.get("input")), number(tokens.get("output")) + number(tokens.get("reasoning")), + number(cache.get("read")), number(cache.get("write")), + ) + except Exception: + # A scan cut short by a database error must not be cached as the whole + # story, or the missing usage stays suppressed until the cache expires. + self.complete = False + return + finally: + try: + conn.close() + except Exception: + pass + + def stats(self): + """Snapshot the aggregated local usage into the record's stats dict.""" + return { + "todayPrompts": self.today_prompts, + "todaySessions": len(self.today_sessions), + "todayTotalTokens": self.today_tokens, + "todayTokensByModel": self.today_by_model, + "recentDays": [{"date": day, "messageCount": self.recent[day]} for day in self.recent_days], + "totalPrompts": self.total_prompts, + "totalSessions": len(self.total_sessions), + "activeDays": len(self.active_days), + "activeDates": sorted(self.active_days), + "modelUsage": self.model_usage, + } + + +def cache_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(): + codex_home = Path(os.environ.get("CODEX_HOME") or (Path.home() / ".codex")) + db = Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) / "opencode" / "opencode.db" + config = Path(os.environ.get("XDG_CONFIG_HOME") or (Path.home() / ".config")) / "omarchy" / "agents" / "local.json" + # The digest covers every path the scan reads: the codex session roots and + # its config.toml, the opencode DB, the provider config, and (via + # Path.home()) the pi/omp session roots. + digest = hashlib.sha1( + "\n".join(str(part) for part in (Path.home(), codex_home, db, config)).encode("utf-8") + ).hexdigest()[:16] + root = cache_root() + return root / f"local-scan-{digest}.json", root / f"local-scan-{digest}.lock" + + +def read_fresh_json(path, max_age_seconds): + 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, payload): + # 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") + # mkstemp opens at 0600; nothing in the cache is sensitive, so open it + # up to the usual 0644. + tmp.chmod(0o644) + tmp.replace(path) + except BaseException: + tmp.unlink(missing_ok=True) + raise + + +# The cache payload is a versioned envelope around the stats dict, so a +# corrupted or foreign-shaped file is a cache miss (rescan + rewrite) instead +# of a crash or a garbage record. +def read_cached_stats(cache_file, max_age_seconds): + 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, + # not merely old, whatever its mtime says. + 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, stats): + try: + write_json(cache_file, {"schemaVersion": 1, "scanDate": today, "stats": stats}) + except Exception as exc: + print(f"omarchy-agent-usage-local: could not write usage cache ({exc})", file=sys.stderr) + + +def run_local_scan(): + scan = Scan() + scan.scan_codex() + scan.scan_pi() + scan.scan_opencode() + return scan.stats(), scan.complete + + +def cached_stats(max_age): + """Stats, with the cache as a pure optimization. + + The cache must never take the collector down: any cache-layer failure + (unwritable cache root, lock errors, disk full) degrades to a direct scan + and a warning on stderr. The JSON record is the contract; the cache is not. + """ + try: + return _cached_stats(max_age) + except Exception as exc: + print(f"omarchy-agent-usage-local: cache unavailable ({exc}); scanning directly", file=sys.stderr) + stats, _ = run_local_scan() + return stats + + +def _cached_stats(max_age): + cache_file, lock_file = scan_cache_paths() + + cached = read_cached_stats(cache_file, max_age) + if cached is not None: + return cached + + with lock_file.open("w") as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + cached = read_cached_stats(cache_file, max_age) + if cached is not None: + return cached + stats, complete = run_local_scan() + # An interrupted scan still serves this run, but caching it would + # suppress the missing usage for every reader until the cache expires. + if complete: + write_cached_stats(cache_file, stats) + return stats + + +def main(): + parser = argparse.ArgumentParser() + # Accepted for parity with the other collectors: the update command passes + # them to everything it runs. There are no limits to fetch here, so + # --limits-only just means "a redraw may reuse an older scan". + parser.add_argument("--force", action="store_true") + parser.add_argument("--limits-only", action="store_true") + args = parser.parse_args() + + max_age = 0 if args.force else (LIMITS_ONLY_REUSE_SECONDS if args.limits_only else SCAN_REUSE_SECONDS) + + record = { + "schemaVersion": 1, + "id": AGENT_ID, + "name": AGENT_NAME, + "updatedAt": datetime.now().astimezone().isoformat(), + "ready": True, + "hasLocalStats": True, + "tierLabel": TIER_LABEL, + "usageStatusText": "", + "authHelpText": "", + # Nothing meters self-hosted work, so the panel draws no limit rows. + "limits": [], + } + record.update(cached_stats(max_age)) + print(json.dumps(record, separators=(",", ":"))) + + +if __name__ == "__main__": + main() diff --git a/test/shell.d/agent-usage-local-scanner-test.sh b/test/shell.d/agent-usage-local-scanner-test.sh new file mode 100755 index 00000000000..3c6940f47e5 --- /dev/null +++ b/test/shell.d/agent-usage-local-scanner-test.sh @@ -0,0 +1,199 @@ +#!/bin/bash + +source "$(dirname "$0")/base-test.sh" + +require_command jq +require_command python3 + +LOCAL_HOME=$(mktemp -d) +trap 'rm -rf "$LOCAL_HOME"' EXIT +mkdir -p "$LOCAL_HOME/.codex/sessions/$(date +%Y/%m/%d)" + +cat >"$LOCAL_HOME/.codex/config.toml" <<'EOF' +[model_providers.homelab] +base_url = "http://server.home:11434/v1" + +[model_providers.openrouter] +base_url = "https://openrouter.ai/api/v1" +EOF + +timestamp="$(date +%Y-%m-%d)T12:00:00Z" +rollout() { + local name=$1 provider=$2 model=$3 + cat >"$LOCAL_HOME/.codex/sessions/$(date +%Y/%m/%d)/rollout-$name.jsonl" <"$LOCAL_HOME/.config/omarchy/agents/local.json" +rollout mystery mystery qwen-custom +# --force: the config changed inside the scan-reuse window, so a no-flag run +# would legitimately serve the previous scan. +local_result=$(env "${env_args[@]}" "$ROOT/bin/omarchy-agent-usage-local" --force) +[[ $(jq -r '.modelUsage["qwen-custom"].inputTokens' <<<"$local_result") == "1000" ]] || + fail "Local collector honors configured provider ids" "$local_result" +pass "Local collector honors configured provider ids" + +# The local and Codex collectors partition the same sessions directory: every +# rollout is counted by exactly one of them, so the fix cannot lose a turn and +# the new tab cannot double-count one. +cat >"$LOCAL_HOME/codex" <<'EOF' +#!/bin/bash +while read -r request; do + id=$(jq -r '.id // empty' <<<"$request") + case "$(jq -r '.method // empty' <<<"$request")" in + initialize) jq -cn --argjson id "$id" '{id: $id, result: {}}' ;; + account/read) jq -cn --argjson id "$id" '{id: $id, result: {account: {}}}' ;; + account/rateLimits/read) jq -cn --argjson id "$id" '{id: $id, result: {rateLimits: {}}}' ;; + esac +done +EOF +chmod +x "$LOCAL_HOME/codex" + +codex_result=$(env "${env_args[@]}" PATH="$LOCAL_HOME:$PATH" "$ROOT/bin/omarchy-agent-usage-codex" --force) + +[[ $(jq -r '.modelUsage | keys | join(",")' <<<"$codex_result") == "gpt-sub" ]] || + fail "Codex collector keeps only the subscription's own turns" "$codex_result" + +overlap=$(jq -n --argjson a "$(jq -c '.modelUsage | keys' <<<"$local_result")" \ + --argjson b "$(jq -c '.modelUsage | keys' <<<"$codex_result")" \ + '$a - ($a - $b) | length') +[[ $overlap == "0" ]] || + fail "Local and Codex collectors must not count the same model twice" "$local_result$codex_result" +pass "Local and Codex collectors partition sessions without gaps or overlap" + +# opencode records the provider id but not the endpoint, so the well-known +# runtimes and configured ids carry the filter - pushed into SQL, since these +# databases keep every historical message with its JSON. +OC_HOME=$(mktemp -d) +trap 'rm -rf "$LOCAL_HOME" "$OC_HOME"' EXIT +mkdir -p "$OC_HOME/.config/omarchy/agents" +echo '{"providers":["homelab"]}' >"$OC_HOME/.config/omarchy/agents/local.json" + +python3 - "$OC_HOME/.local/share/opencode/opencode.db" <<'PY' +import json, sqlite3, sys, 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(mid, provider, model, role="assistant", inp=0, out=0, reasoning=0): + return (mid, "ses_1", now_ms, now_ms, json.dumps({ + "role": role, "providerID": provider, "modelID": model, + "tokens": {"input": inp, "output": out, "reasoning": reasoning, "cache": {"read": 0, "write": 0}}, + "time": {"created": now_ms}, + })) + +conn.executemany("INSERT INTO message VALUES (?, ?, ?, ?, ?)", [ + message("m1", "ollama", "qwen-oc", inp=100, out=10, reasoning=5), + message("m2", "homelab", "qwen-cfg", inp=50, out=5), + message("m3", "anthropic", "claude-paid", inp=999, out=999), + message("m4", "openai", "gpt-sub", inp=999, out=999), + message("m5", "ollama", "qwen-oc", role="user", inp=999), +]) +conn.execute("INSERT INTO message VALUES ('m6', 'ses_1', ?, ?, 'not json at all')", (now_ms, now_ms)) +conn.commit(); conn.close() +PY + +oc_env=(HOME="$OC_HOME" CODEX_HOME="$OC_HOME/.codex" XDG_CACHE_HOME="$OC_HOME/.cache" + XDG_DATA_HOME="$OC_HOME/.local/share" XDG_CONFIG_HOME="$OC_HOME/.config") +oc=$(env "${oc_env[@]}" "$ROOT/bin/omarchy-agent-usage-local") + +# 115 = 100 + 10 + 5 reasoning; 55 = 50 + 5. Paid providers and the user +# message contribute nothing, and the malformed row does not abort the scan. +[[ $(jq -r '.todayTotalTokens' <<<"$oc") == "170" ]] || + fail "Local collector counts opencode turns on self-hosted providers" "$oc" +[[ $(jq -r '.modelUsage | keys | sort | join(",")' <<<"$oc") == "qwen-cfg,qwen-oc" ]] || + fail "Local collector filters opencode rows to self-hosted providers" "$oc" +[[ $(jq -r '.modelUsage["qwen-oc"].outputTokens' <<<"$oc") == "15" ]] || + fail "Local collector counts opencode reasoning tokens as generated" "$oc" +pass "Local collector counts self-hosted opencode turns past paid and malformed rows" + +# The scan is cached like every other collector's: one envelope per data-path +# set, readable, versioned, and bypassed by --force. +cache_file=$(ls "$OC_HOME/.cache/omarchy/agent-usage/"local-scan-*.json 2>/dev/null | head -n 1) +[[ -n $cache_file && $(stat -c %a "$cache_file") == "644" ]] || + fail "Local collector leaves a readable cache file behind" "$oc" +[[ $(jq -r '.schemaVersion' "$cache_file") == "1" && $(jq -r '.stats.todayTotalTokens' "$cache_file") == "170" ]] || + fail "Local collector writes a versioned cache envelope" "$oc" +pass "Local collector caches its scan" + +python3 - "$OC_HOME/.local/share/opencode/opencode.db" <<'PY' +import json, sqlite3, sys, time +conn = sqlite3.connect(sys.argv[1]); now_ms = int(time.time() * 1000) +conn.execute("INSERT INTO message VALUES (?, ?, ?, ?, ?)", ("m7", "ses_1", now_ms, now_ms, json.dumps({ + "role": "assistant", "providerID": "ollama", "modelID": "qwen-oc", + "tokens": {"input": 30, "output": 0, "reasoning": 0, "cache": {"read": 0, "write": 0}}, + "time": {"created": now_ms}}))) +conn.commit(); conn.close() +PY + +[[ $(jq -r '.todayTotalTokens' <<<"$(env "${oc_env[@]}" "$ROOT/bin/omarchy-agent-usage-local")") == "170" ]] || + fail "Local collector reuses a seconds-old scan" "$oc" +[[ $(jq -r '.todayTotalTokens' <<<"$(env "${oc_env[@]}" "$ROOT/bin/omarchy-agent-usage-local" --force)") == "200" ]] || + fail "Local collector --force rescans past the cache" "$oc" +pass "Local collector --force rescans past the cache" + +# A cache stamped with another local date holds another day's today* stats. +jq -c '.scanDate = "1999-01-01"' "$cache_file" >"$cache_file.tmp" && mv "$cache_file.tmp" "$cache_file" +[[ $(jq -r '.scanDate' <<<"$(cat "$cache_file")") == "1999-01-01" ]] && + [[ $(jq -r '.todayTotalTokens' <<<"$(env "${oc_env[@]}" "$ROOT/bin/omarchy-agent-usage-local" --limits-only)") == "200" ]] || + fail "Local collector treats a cache from another day as a miss" "$oc" +pass "Local collector treats a cache from another day as a miss" + +# A corrupt-but-parseable cache is a miss, not a garbage record. +printf '[]' >"$cache_file" +[[ $(jq -r '.todayTotalTokens' <<<"$(env "${oc_env[@]}" "$ROOT/bin/omarchy-agent-usage-local" --limits-only)") == "200" ]] || + fail "Local collector recovers from a corrupt cache file" "$oc" +pass "Local collector recovers from a corrupt cache file" + +# An unwritable cache must not kill the collector: the record is the contract. +touch "$OC_HOME/not-a-dir" +broken=$(env "${oc_env[@]}" XDG_CACHE_HOME="$OC_HOME/not-a-dir" "$ROOT/bin/omarchy-agent-usage-local" 2>/dev/null) +[[ $(jq -r '.todayTotalTokens' <<<"$broken") == "200" ]] || + fail "Local collector still prints a record when the cache is unwritable" "$broken" +pass "Local collector survives an unwritable cache" + +# A scan cut short by a database error must not be cached as the whole story. +BROKEN_HOME=$(mktemp -d) +trap 'rm -rf "$LOCAL_HOME" "$OC_HOME" "$BROKEN_HOME"' EXIT +python3 - "$BROKEN_HOME/.local/share/opencode/opencode.db" <<'PY' +import sqlite3, sys +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 unrelated (id text PRIMARY KEY)") +conn.commit(); conn.close() +PY +env HOME="$BROKEN_HOME" CODEX_HOME="$BROKEN_HOME/.codex" XDG_CACHE_HOME="$BROKEN_HOME/.cache" \ + XDG_DATA_HOME="$BROKEN_HOME/.local/share" XDG_CONFIG_HOME="$BROKEN_HOME/.config" \ + "$ROOT/bin/omarchy-agent-usage-local" >/dev/null +[[ -z $(ls "$BROKEN_HOME/.cache/omarchy/agent-usage/"local-scan-*.json 2>/dev/null) ]] || + fail "Local collector must not cache an interrupted scan" "broken-db" +pass "Local collector does not cache an interrupted opencode scan" From 108c4719f3f3e63eed2f9b3653c27a9615098564 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:07:12 -0400 Subject: [PATCH 092/546] Bring the Devin collector in line with the collector contract Honor XDG_DATA_HOME and XDG_CACHE_HOME, count pi/omp sessions on the devin provider, keep a quota window that reset at zero instead of dropping it, and answer an unused machine with ready: false rather than an install hint. Read the plan's hidden windows, credit buckets and tier names from the cached GetUserStatus, and fall back to the older per-session databases when the CLI's own is absent. Untrusted counts and timestamps skip themselves, and a scan that couldn't read a database is not cached. Co-authored-by: betim-hodza Co-authored-by: Joao Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-devin | 1025 +++++++++++------ .../shell.d/agent-usage-devin-scanner-test.sh | 566 ++++----- 2 files changed, 925 insertions(+), 666 deletions(-) diff --git a/bin/omarchy-agent-usage-devin b/bin/omarchy-agent-usage-devin index 63d29ad0440..a7d524af3fb 100755 --- a/bin/omarchy-agent-usage-devin +++ b/bin/omarchy-agent-usage-devin @@ -1,226 +1,192 @@ #!/usr/bin/python3 -# omarchy:summary=Print the Devin CLI usage record as JSON +# omarchy:summary=Print the Devin usage record as JSON # omarchy:args=[--force] [--limits-only] # omarchy:hidden=true -"""Collect Devin CLI usage into one display-ready JSON record. +"""Collect Devin usage into one display-ready JSON record. Token stats come from the Devin CLI session database at -~/.local/share/devin/cli/sessions.db. Each assistant message carries -metadata.metrics with input/output/cache token counts and the generation -model, so the collector aggregates those by day and by model the same way -the Claude and Codex collectors do. - -Rate limits (daily and weekly quota percentages with reset times) come from -the user-status cache that the Devin TUI writes at -~/.cache/devin/cli/user_status.*.bin on every startup. The TUI fetches -these from the server via a gRPC call (GetUserStatus); the collector reads -the cached result so it doesn't need to make its own network call. +$XDG_DATA_HOME/devin/cli/sessions.db, where each assistant message carries +metadata.metrics with its input/output/cache token counts and the model that +generated it, and from pi/omp sessions that ran on the devin provider. + +The daily and weekly quota comes from the user-status cache the Devin CLI +writes at $XDG_CACHE_HOME/devin/cli/user_status.*.bin every time it starts: +the GetUserStatus answer it fetched, base64 protobuf. The collector reads that +cache, the same data the CLI's /usage command renders, and makes no network +call of its own. The agents panel only ever reads the JSON this prints. """ import argparse import base64 -import datetime as dt import fcntl import hashlib import json +import math import os +import shutil import sqlite3 +import subprocess import sys import tempfile import time +from datetime import datetime, timedelta, timezone from pathlib import Path -from typing import Any AGENT_ID = "devin" AGENT_NAME = "Devin" +PROVIDER = "devin" +AUTH_HELP = "Run `devin auth login` so Devin can record its quota." +REFRESH_HELP = "Devin hasn't recorded its quota since it was signed in. Start Devin to refresh it." # 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, however low -# refreshIntervalSec is set. --limits-only promises only fresh limits, so it -# may reuse a scan for up to 15 minutes. Devin's limits come from the TUI's -# user-status cache, not a live API call, so they are always read fresh. +# own). --limits-only promises only fresh limits, so it may reuse a scan for +# up to 15 minutes. The quota is a file read, so it is read on every run. SCAN_REUSE_SECONDS = 20 LIMITS_ONLY_REUSE_SECONDS = 900 -# Devin's quota comes from a cache the TUI writes, never a live fetch, and -# cloud sessions keep spending it meanwhile. Numbers this recent still -# describe the account; older ones are flagged stale with their fetch time. +# The quota comes from a cache the CLI writes when it starts, never from a +# check of ours, and cloud sessions keep spending it meanwhile. Numbers this +# recent still describe the account; older ones are reported stale, with +# their age. LIMITS_CURRENT_ENOUGH_SECONDS = 900 +# The plan tiers GetUserStatus numbers, named as Devin names them. Only asked +# when the plan carries no name of its own. +TEAMS_TIER = { + 1: "Teams", 2: "Pro", 3: "Enterprise Saas", 4: "Hybrid", 5: "Enterprise Self Hosted", + 6: "Waitlist Pro", 7: "Teams Ultimate", 9: "Trial", 10: "Enterprise Self Serve", + 12: "Devin Enterprise", 14: "Devin Teams", 15: "Devin Teams V2", 16: "Devin Pro", + 17: "Devin Max", 18: "Max", 19: "Devin Free", 20: "Devin Trial", +} -def data_dir() -> Path: - return Path(os.environ.get("DEVIN_DATA_DIR") or (Path.home() / ".local" / "share" / "devin" / "cli")) +# Past this a token count or timestamp from disk is garbage, not usage. +MAX_TOKENS = 10_000_000_000 +EARLIEST_EPOCH = 946_684_800 # 2000-01-01 +LATEST_EPOCH = 4_102_444_800 # 2100-01-01 -def db_path() -> Path: - return data_dir() / "sessions.db" +# --------------------------------------------------------------- utilities -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 obj(value): + return value if isinstance(value, dict) else {} -def cache_paths() -> tuple[Path, Path]: - digest = hashlib.sha1(str(data_dir()).encode("utf-8")).hexdigest()[:16] - root = cache_root() - return root / f"devin-scan-{digest}.json", root / f"devin-scan-{digest}.lock" +# Token counts and timestamps on disk are untrusted: a boolean, NaN, infinite, +# negative, huge or non-numeric value reads as nothing. +def number(value): + if isinstance(value, bool) or not isinstance(value, (int, float)): + return 0 + if not math.isfinite(value) or value < 0 or value > MAX_TOKENS: + return 0 + return int(value) -def date_string(value: dt.date) -> str: - return value.strftime("%Y-%m-%d") +def epoch_seconds(value): + if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value): + return None + # Session timestamps are seconds or milliseconds. + seconds = value / 1000 if value > 100_000_000_000 else value + if not EARLIEST_EPOCH <= seconds < LATEST_EPOCH: + return None + return seconds -def today_string() -> str: - return date_string(dt.datetime.now().date()) +def parse_time(value): + """A timestamp from disk as an aware datetime, or None when it won't read.""" + seconds = epoch_seconds(value) + if seconds is not None: + return datetime.fromtimestamp(seconds, timezone.utc) + if not isinstance(value, str) or not value.strip(): + return None + try: + parsed = datetime.fromisoformat(value.strip().replace("Z", "+00:00")) + except ValueError: + return None + if parsed.tzinfo is None: + parsed = parsed.astimezone() + if epoch_seconds(parsed.timestamp()) is None: + return None + return parsed -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_day(value): + parsed = parse_time(value) + return parsed.astimezone().strftime("%Y-%m-%d") if parsed else None -def local_date_string() -> str: - return date_string(dt.datetime.now().date()) +def iso_from_unix(value): + seconds = epoch_seconds(value) + return datetime.fromtimestamp(seconds, timezone.utc).isoformat() if seconds else "" -def number(value: Any) -> int: - try: - return int(value or 0) - except (TypeError, ValueError): - return 0 +def model_name(raw): + return str(raw or "").strip() or AGENT_ID -def model_name(raw: Any) -> str: - value = str(raw or "devin") - return value if value else "devin" +def runtime_env(): + home = str(Path.home()) + path_parts = [ + os.environ.get("PATH", ""), + f"{home}/.local/bin", + f"{home}/.npm-global/bin", + f"{home}/.local/share/mise/shims", + ] + env = os.environ.copy() + env["PATH"] = os.pathsep.join(part for part in path_parts if part) + return env -def empty_bucket() -> dict[str, int]: - return { - "inputTokens": 0, - "outputTokens": 0, - "cacheReadInputTokens": 0, - "cacheCreationInputTokens": 0, - } +ENV = runtime_env() -# ---------------------------------------------------------------- local scan - - -def run_scan(conn: sqlite3.Connection) -> dict[str, Any]: - """Run the full SQLite scan and return the stats dict for caching.""" - today = dt.datetime.now().date() - recent_dates = recent_date_strings() - recent: dict[str, int] = {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() - - # json_extract raises on malformed JSON, aborting the whole query — and the - # direct-scan fallback — on a single bad row. CASE only evaluates the - # matching branch, so invalid chat_message rows are skipped instead; a bare - # AND would not order json_valid before json_extract. - for row in conn.execute( - """ - SELECT - date(mn.created_at, 'unixepoch', 'localtime') AS day, - json_extract(mn.chat_message, '$.metadata.metrics.input_tokens') AS input_tokens, - json_extract(mn.chat_message, '$.metadata.metrics.output_tokens') AS output_tokens, - json_extract(mn.chat_message, '$.metadata.metrics.cache_read_tokens') AS cache_read_tokens, - json_extract(mn.chat_message, '$.metadata.metrics.cache_creation_tokens') AS cache_creation_tokens, - json_extract(mn.chat_message, '$.metadata.generation_model') AS model - FROM message_nodes mn - WHERE CASE WHEN json_valid(mn.chat_message) - THEN json_extract(mn.chat_message, '$.metadata.metrics') IS NOT NULL - ELSE 0 END - """ - ): - inp = number(row["input_tokens"]) - out = number(row["output_tokens"]) - cache_read = number(row["cache_read_tokens"]) - cache_write = number(row["cache_creation_tokens"]) - total = inp + out + cache_read + cache_write - if total <= 0: - continue +def find_command(name): + return shutil.which(name, path=ENV.get("PATH")) - model = model_name(row["model"]) - day = row["day"] or local_date_string() - bucket = model_usage.setdefault(model, empty_bucket()) - bucket["inputTokens"] += inp - bucket["outputTokens"] += out - bucket["cacheReadInputTokens"] += cache_read - bucket["cacheCreationInputTokens"] += cache_write +def data_home(): + return Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) - if day: - active_dates.add(day) - if day in recent: - recent[day] += total - if day == date_string(today): - today_by_model[model] = today_by_model.get(model, 0) + total - - today_prompts = 0 - today_sessions = 0 - total_prompts = 0 - - for row in conn.execute( - """ - SELECT - date(timestamp, 'unixepoch', 'localtime') AS day, - COUNT(*) AS prompts, - COUNT(DISTINCT session_id) AS sessions - FROM prompt_history - GROUP BY day - """ - ): - day = row["day"] or local_date_string() - prompts = number(row["prompts"]) - total_prompts += prompts - if day == date_string(today): - today_prompts = prompts - today_sessions = number(row["sessions"]) - - # Per-day distinct counts can't be summed: a session spanning midnight - # belongs to two days. All-time sessions need their own distinct count. - total_sessions = number( - conn.execute("SELECT COUNT(DISTINCT session_id) FROM prompt_history").fetchone()[0] - ) - return { - "todayPrompts": today_prompts, - "todaySessions": today_sessions, - "todayTotalTokens": sum(today_by_model.values()), - "todayTokensByModel": today_by_model, - "recentDays": [{"date": day, "messageCount": recent[day]} for day in recent_dates], - "totalPrompts": total_prompts, - "totalSessions": total_sessions, - "activeDays": len(active_dates), - "activeDates": sorted(active_dates), - "modelUsage": model_usage, - } +def xdg_cache_home(): + return Path(os.environ.get("XDG_CACHE_HOME") or (Path.home() / ".cache")) + + +def cache_root(): + root = xdg_cache_home() / "omarchy" / "agent-usage" + root.mkdir(parents=True, exist_ok=True) + return root + + +def read_json(path): + try: + return json.loads(path.read_text(encoding="utf-8")) + except Exception: + return None -def read_fresh_json(path: Path, max_age_seconds: int) -> Any: - if max_age_seconds <= 0 or not path.exists(): +def read_fresh_json(path, max_age_seconds): + if max_age_seconds <= 0: 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: + except OSError: return None - return None + return read_json(path) if 0 <= age < max_age_seconds else None -def write_json(path: Path, payload: Any) -> None: +def write_json(path, payload): + # A temp name unique to this writer: 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=(",", ":"), sort_keys=True) + "\n") + handle.write(json.dumps(payload, separators=(",", ":")) + "\n") tmp.chmod(0o644) tmp.replace(path) except BaseException: @@ -228,231 +194,569 @@ def write_json(path: Path, payload: Any) -> None: raise -def read_cached_stats(cache_file: Path, max_age_seconds: int) -> 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, - # not merely old, whatever its mtime says. - if cached.get("scanDate") != today_string(): - 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 +# ------------------------------------------------------------------- stats -def write_cached_stats(cache_file: Path, stats: dict[str, Any]) -> None: - try: - write_json(cache_file, {"schemaVersion": 1, "scanDate": today_string(), "stats": stats}) - except Exception as exc: - print(f"omarchy-agent-usage-devin: could not write usage cache ({exc})", file=sys.stderr) +def empty_bucket(): + return {"inputTokens": 0, "outputTokens": 0, "cacheReadInputTokens": 0, "cacheCreationInputTokens": 0} -def cached_scan(conn: sqlite3.Connection, force: bool, limits_only: bool) -> dict[str, Any]: - """Local stats, with the cache as a pure optimization. +class Stats: + """Token and activity totals shared by every local source. Sessions and days + are sets, so a day active in two sources counts once toward activeDays.""" - The cache must never take the collector down: any cache-layer failure - (unwritable cache root, lock errors, disk full) degrades to a direct scan - and a warning on stderr. The JSON record is the contract; the cache is not. - """ - max_age = 0 if force else (LIMITS_ONLY_REUSE_SECONDS if limits_only else SCAN_REUSE_SECONDS) + def __init__(self): + 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() + self.today_tokens_by_model = {} + self.total_prompts = 0 + self.total_sessions = set() + self.active_days = set() + self.model_usage = {} + + # A session of None is settled later, with add_session. + def add_session(self, session, today): + if session is not None: + self.total_sessions.add(session) + if today: + self.today_sessions.add(session) + + def add_prompt(self, session, day): + self.total_prompts += 1 + if day == self.today: + self.today_prompts += 1 + self.add_session(session, day == self.today) + + def add(self, session, day, model, input_tokens, output_tokens, cache_read, cache_write, prompt=True): + total = input_tokens + output_tokens + cache_read + cache_write + if total <= 0: + return + if prompt: + self.add_prompt(session, day) + self.add_session(session, day == self.today) + self.active_days.add(day) + bucket = self.model_usage.setdefault(model, empty_bucket()) + bucket["inputTokens"] += input_tokens + bucket["outputTokens"] += output_tokens + bucket["cacheReadInputTokens"] += cache_read + bucket["cacheCreationInputTokens"] += cache_write + if day in self.recent: + # recentDays.messageCount holds tokens, despite the legacy name. + self.recent[day] += total + if day == self.today: + self.today_total_tokens += total + self.today_tokens_by_model[model] = self.today_tokens_by_model.get(model, 0) + total + + def record(self): + 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, + } + + +# ------------------------------------------------------------- Devin CLI + + +def data_dir(): + return Path(os.environ.get("DEVIN_DATA_DIR") or (data_home() / "devin" / "cli")) + + +def credentials_path(): + return data_dir().parent / "credentials.toml" + + +def session_databases(): + """The CLI's own database, or failing that the older layout some builds + wrote straight under $XDG_DATA_HOME/devin. Never both: a machine that moved + to the newer layout may still have the old files, holding the same work.""" + primary = data_dir() / "sessions.db" + if primary.is_file(): + return [primary] + root = data_home() / "devin" + return [path for path in (root / "cli_sessions.db", root / "sessions.db") if path.is_file()] + + +def table_columns(conn, table): + return {row[1] for row in conn.execute(f"PRAGMA table_info({table})")} + + +def scan_message_nodes(conn, stats, prefix): + # Every assistant message records what it spent under metadata.metrics. + # The LIKE gate skips rows that cannot hold metrics before any parse. + for session_id, created_at, raw in conn.execute( + "SELECT session_id, created_at, chat_message FROM message_nodes WHERE chat_message LIKE '%\"metrics\"%'" + ): + try: + metadata = obj(obj(json.loads(raw)).get("metadata")) + except (TypeError, ValueError): + continue + metrics = obj(metadata.get("metrics")) + day = local_day(created_at) + if not metrics or day is None: + continue + # Prompts are counted from prompt_history, where each one is a row. + stats.add( + prefix + str(session_id), day, model_name(metadata.get("generation_model")), + number(metrics.get("input_tokens")), number(metrics.get("output_tokens")), + number(metrics.get("cache_read_tokens")), number(metrics.get("cache_creation_tokens")), + prompt=False, + ) + + +def scan_prompt_history(conn, stats, prefix): + for session_id, timestamp in conn.execute("SELECT session_id, timestamp FROM prompt_history"): + day = local_day(timestamp) + if day is not None: + stats.add_prompt(prefix + str(session_id), day) + + +def scan_session_totals(conn, stats, prefix, tables): + # The older layout keeps one lifetime total per session and no split by + # kind, so the tokens land on the day the session was last active and in + # inputTokens, the one bucket that cannot overstate a split it doesn't have. + for session_id, model, created_at, last_activity_at, total_tokens in conn.execute( + "SELECT id, model, created_at, last_activity_at, total_tokens FROM sessions" + ): + day = local_day(last_activity_at) or local_day(created_at) + if day is not None: + stats.add(prefix + str(session_id), day, model_name(model), number(total_tokens), 0, 0, 0, prompt=False) + if "messages" in tables and {"role", "timestamp", "session_id"} <= table_columns(conn, "messages"): + for session_id, role, timestamp in conn.execute("SELECT session_id, role, timestamp FROM messages WHERE role = 'assistant'"): + day = local_day(timestamp) + if day is not None: + stats.add_prompt(prefix + str(session_id), day) + + +def scan_devin_database(path, stats): + """Add one session database. Returns whether it was read to the end.""" try: - return _cached_scan(conn, max_age) + # Read-only: Devin may be writing right now. + conn = sqlite3.connect(path.resolve().as_uri() + "?mode=ro", uri=True, timeout=2) except sqlite3.Error: - raise # a database failure is not a cache-layer problem - except Exception as exc: - print(f"omarchy-agent-usage-devin: cache unavailable ({exc}); scanning directly", file=sys.stderr) - return run_scan(conn) + return False + prefix = "devin:" + path.name + ":" + try: + conn.execute("PRAGMA query_only = ON") + tables = {name for (name,) in conn.execute("SELECT name FROM sqlite_master WHERE type = 'table'")} + known = False + if "message_nodes" in tables and {"session_id", "created_at", "chat_message"} <= table_columns(conn, "message_nodes"): + known = True + scan_message_nodes(conn, stats, prefix) + if "prompt_history" in tables and {"session_id", "timestamp"} <= table_columns(conn, "prompt_history"): + known = True + scan_prompt_history(conn, stats, prefix) + if not known and "sessions" in tables and {"id", "model", "created_at", "last_activity_at", "total_tokens"} <= table_columns(conn, "sessions"): + known = True + scan_session_totals(conn, stats, prefix, tables) + if not known: + # A layout this collector doesn't know: reading it again won't help. + print(f"omarchy-agent-usage-devin: {path} has no session tables this collector reads", file=sys.stderr) + except sqlite3.Error as error: + # Transient lock, schema migration, corruption: the numbers stop here, + # incomplete. + print(f"omarchy-agent-usage-devin: {path}: {error}", file=sys.stderr) + return False + finally: + conn.close() + return True -def _cached_scan(conn: sqlite3.Connection, max_age: int) -> dict[str, Any]: - cache_file, lock_file = cache_paths() +# --------------------------------------------------------------- pi and omp - cached = read_cached_stats(cache_file, max_age) - if cached is not None: - return cached - with lock_file.open("w") as lock: - fcntl.flock(lock, fcntl.LOCK_EX) - cached = read_cached_stats(cache_file, max_age) - if cached is not None: - return cached - stats = run_scan(conn) - write_cached_stats(cache_file, stats) - return stats +# `omp --profile=` (and pi's equivalent) moves the agent tree under +# /profiles//, and PI_CODING_AGENT_DIR moves pi's default one. +# Sessions are keyed by path, so a profile adds sessions rather than counting +# the default root again. +def pi_pattern(): + return r'"provider"\s*:\s*"devin"' -# ------------------------------------------------------------------- limits -# -# Devin does not expose rate limits through a CLI command or REST API. The -# TUI fetches GetUserStatus from the server on startup and caches the -# protobuf response at ~/.cache/devin/cli/user_status.*.bin. The collector -# reads that cache — the same data the /usage slash command renders — so it -# doesn't need credentials or a network call of its own. +def accept_pi(message): + return str(message.get("provider") or "") == PROVIDER -def user_status_dir() -> Path: - return Path(os.environ.get("DEVIN_CACHE_DIR") or (Path.home() / ".cache" / "devin" / "cli")) +def pi_session_roots(): + pi_agent = os.environ.get("PI_CODING_AGENT_DIR") + agent_dirs = [Path(pi_agent).expanduser() if pi_agent else Path.home() / ".pi" / "agent", Path.home() / ".omp" / "agent"] + for base in (Path.home() / ".pi", Path.home() / ".omp"): + try: + agent_dirs.extend(sorted(child / "agent" for child in (base / "profiles").iterdir() if child.is_dir())) + except OSError: + pass + roots = [] + for agent_dir in agent_dirs: + root = agent_dir / "sessions" + if root.is_dir() and root.resolve() not in [r.resolve() for r in roots]: + roots.append(root) + return roots + + +def pi_parent_session(path): + # A fork names the session it was cut from in its header line. + try: + with open(path, encoding="utf-8", errors="replace") as handle: + header = json.loads(handle.readline()) + except (OSError, ValueError): + return None + parent = header.get("parentSession") if isinstance(header, dict) else None + return os.path.abspath(os.path.expanduser(parent)) if isinstance(parent, str) and parent else None -def credentials_mtime() -> float | None: - # The CLI rewrites credentials.toml on sign-in, so a user-status file fetched - # before that write describes a previous account. The fetch can land a beat - # before the file flush, hence the small skew allowance. +def scan_pi_sessions(stats): + """Add the pi/omp assistant messages this collector counts. Returns whether + every root was read.""" + roots = pi_session_roots() + if not roots: + return True + rg = find_command("rg") or "rg" try: - return (data_dir().parent / "credentials.toml").stat().st_mtime - 30 + proc = subprocess.Popen( + # --no-ignore: these are data files, not a source tree. A $HOME that is + # a git checkout would otherwise hide every session behind its ignores. + [rg, "--json", "--no-ignore", "-e", pi_pattern(), *map(str, roots)], + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + text=True, + errors="replace", + env=ENV, + ) except OSError: + return False + + counted = set() + owned = {} + today_owned = {} + assert proc.stdout is not None + for raw in proc.stdout: + try: + event = json.loads(raw) + if event.get("type") != "match": + continue + path = os.path.abspath(event["data"]["path"]["text"]) + line_number = event["data"].get("line_number") + entry = json.loads(event["data"]["lines"]["text"]) + except Exception: + continue + if not isinstance(entry, dict) or entry.get("type") != "message": + continue + message = obj(entry.get("message")) + if message.get("role") != "assistant" or not accept_pi(message): + continue + usage = obj(message.get("usage")) + 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 not (input_tokens or output_tokens or cache_read or cache_write): + input_tokens = number(usage.get("totalTokens")) + timestamp = entry.get("timestamp") + day = local_day(timestamp) or local_day(message.get("timestamp")) + if day is None or not (input_tokens or output_tokens or cache_read or cache_write): + continue + # A fork repeats its parent's messages with their ids and timestamps, so + # the pair says two lines are one message. An id alone is only unique + # within its file, and a line without one always counts. + message_id = str(entry.get("id") or "") + if message_id and timestamp is not None: + key = ("entry", message_id, str(timestamp)) + elif message_id: + key = ("path", path, message_id) + else: + key = ("line", path, line_number) + owned.setdefault(path, set()).add(key) + if day == stats.today: + today_owned.setdefault(path, set()).add(key) + if key in counted: + continue + counted.add(key) + stats.add(None, day, model_name(message.get("model")), input_tokens, output_tokens, cache_read, cache_write) + + try: + code = proc.wait(timeout=5) + except subprocess.TimeoutExpired: + proc.kill() + return False + + # A message belongs to the oldest session holding it, so a fork is a + # session of its own only for what it added, whichever file rg read first. + parents = {} + for path, keys in owned.items(): + inherited = set() + visited = {path} + parent = parents.setdefault(path, pi_parent_session(path)) + while parent and parent not in visited: + visited.add(parent) + inherited |= owned.get(parent, set()) + parent = parents.setdefault(parent, pi_parent_session(parent)) + if keys - inherited: + stats.add_session(path, bool(today_owned.get(path, set()) - inherited)) + # rg exits 1 when nothing matched and 2 when a file couldn't be read. + return code in (0, 1) + + +# -------------------------------------------------------------- scan cache + + +def scan_cache_paths(): + digest = hashlib.sha1("\n".join([ + str(Path.home()), str(data_dir()), str(data_home()), os.environ.get("PI_CODING_AGENT_DIR", ""), + ]).encode("utf-8")).hexdigest()[:16] + root = cache_root() + return root / f"devin-scan-{digest}.json", root / f"devin-scan-{digest}.lock" + + +def run_local_scans(): + stats = Stats() + complete = scan_pi_sessions(stats) + for path in session_databases(): + complete = scan_devin_database(path, stats) and complete + return stats.record(), complete + + +def read_cached_stats(cache_file, max_age): + cached = read_fresh_json(cache_file, max_age) + if not isinstance(cached, dict) or cached.get("schemaVersion") != 2: + return None + # today* fields only mean "today" on the day they were scanned. + if cached.get("scanDate") != datetime.now().strftime("%Y-%m-%d"): return None + stats = cached.get("stats") + return stats if isinstance(stats, dict) and set(Stats().record()) <= set(stats) else None + + +def cached_local_stats(max_age): + """Local stats, with the cache as a pure optimization: any cache-layer + failure degrades to a direct scan and a warning on stderr.""" + try: + cache_file, lock_file = scan_cache_paths() + cached = read_cached_stats(cache_file, max_age) + if cached is not None: + return cached + with lock_file.open("w") as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + cached = read_cached_stats(cache_file, max_age) + if cached is not None: + return cached + stats, complete = run_local_scans() + # A scan that couldn't read everything serves this run, but caching it + # would hide the missing usage from every reader until it expires. + if complete: + write_json(cache_file, {"schemaVersion": 2, "scanDate": datetime.now().strftime("%Y-%m-%d"), "stats": stats}) + return stats + except Exception as exc: + print(f"omarchy-agent-usage-devin: cache unavailable ({exc}); scanning directly", file=sys.stderr) + return run_local_scans()[0] + +# ------------------------------------------------------------------ quota +# +# GetUserStatus is protobuf, and the record needs a handful of its fields, so +# they are walked by hand rather than through generated code. In the cached +# UserStatus, field 10 is the plan tier and field 13 the PlanStatus; inside +# that, field 1 is the PlanInfo, 3 the plan's end, 5-7 the credits used, +# 14/15 the daily/weekly percent remaining and 17/18 their reset times. -def _decode_varint(data: bytes, pos: int) -> tuple[int, int]: + +def decode_varint(data, pos): result = 0 shift = 0 while True: byte = data[pos] result |= (byte & 0x7F) << shift pos += 1 - if not (byte & 0x80): - break + if not byte & 0x80: + return result, pos shift += 7 - return result, pos + if shift > 63: + raise ValueError("varint too long") -def _find_field(data: bytes, target_fn: int) -> bytes | int | None: - """Return the value of field `target_fn` in a protobuf message, or None.""" +def proto_fields(data): + """field number -> first value: an int for a varint, bytes for a message.""" + found = {} pos = 0 - while pos < len(data): - try: - tag, pos = _decode_varint(data, pos) - fn = tag >> 3 - wt = tag & 0x07 - if wt == 0: # varint - val, pos = _decode_varint(data, pos) - if fn == target_fn: - return val - elif wt == 1: # 64-bit - pos += 8 - elif wt == 2: # length-delimited - length, pos = _decode_varint(data, pos) - val = data[pos:pos + length] - pos += length - if fn == target_fn: - return val - elif wt == 5: # 32-bit - pos += 4 + try: + while pos < len(data): + tag, pos = decode_varint(data, pos) + number_, wire = tag >> 3, tag & 0x07 + if wire == 0: + value, pos = decode_varint(data, pos) + elif wire == 1: + value, pos = None, pos + 8 + elif wire == 2: + length, pos = decode_varint(data, pos) + value, pos = bytes(data[pos:pos + length]), pos + length + elif wire == 5: + value, pos = None, pos + 4 else: break - except IndexError: - break - return None + if value is not None: + found.setdefault(number_, value) + except (IndexError, ValueError): + pass + return found + + +def message(fields, number_): + value = fields.get(number_) + return proto_fields(value) if isinstance(value, bytes) else {} -def _iso_from_unix(unix_ts: int | None) -> str: - if not unix_ts or unix_ts <= 0: - return "" - return dt.datetime.fromtimestamp(unix_ts, tz=dt.timezone.utc).isoformat() +def varint(fields, number_): + value = fields.get(number_) + if not isinstance(value, int): + return None + return value - (1 << 64) if value >= (1 << 63) else value + +def text(fields, number_): + value = fields.get(number_) + return value.decode("utf-8", errors="replace").strip() if isinstance(value, bytes) else "" -# A cached percentage outlives the TUI session that measured it, but only -# until its window rolls over: once a window has reset, the figure describes -# a period that is over, and a stale 69% would misreport a fresh allowance. -# A window with no reset time, or one that won't parse, is kept — an -# unreadable timestamp is no reason to throw away a real number. -def _limit_window_open(entry: dict[str, Any], now: dt.datetime) -> bool: + +def quota_window(label, remaining, reset): + # The percent is what's left; the panel wants what's spent. Proto3 leaves a + # zero off the wire, so a window with nothing left carries no percent. + remaining = max(0, min(100, remaining or 0)) + return {"label": label, "percent": (100 - remaining) / 100.0, "resetsAt": iso_from_unix(reset)} + + +def parse_user_status(payload): + status = proto_fields(payload) + plan_status = message(status, 13) + plan = message(plan_status, 1) + devin_info = message(plan, 33) + tier = text(devin_info, 2) or text(plan, 2) or TEAMS_TIER.get(varint(plan, 1)) or TEAMS_TIER.get(varint(status, 10)) or "" + + limits = [] + if plan_status: + # A credit plan leaves the quota percents at their defaults, so a window + # only shows when the server dates its reset or the plan bills by quota, + # and never when the plan hides it. + quota_plan = varint(plan, 35) == 2 + windows = (("Daily", 36, 14, 17), ("Weekly", 37, 15, 18)) + for label, hidden, remaining, reset in windows: + if varint(plan, hidden): + continue + if (varint(plan_status, reset) or 0) > 0 or quota_plan: + limits.append(quota_window(label, varint(plan_status, remaining), varint(plan_status, reset))) + + # Monthly credit buckets exist only on credit plans, and last until the + # plan's period ends; a quota plan reports no grant and is skipped. + plan_end = iso_from_unix(varint(message(plan_status, 3), 1)) + for label, limit_field, used_field in (("Prompt credits", 12, 6), ("Flow credits", 13, 5), ("Flex credits", 14, 7)): + limit = varint(plan, limit_field) or 0 + used = max(0, varint(plan_status, used_field) or 0) + if limit > 0: + limits.append({"label": label, "percent": min(1.0, used / limit), "resetsAt": plan_end}) + return {"tierLabel": tier, "limits": limits} + + +def user_status_dir(): + return Path(os.environ.get("DEVIN_CACHE_DIR") or (xdg_cache_home() / "devin" / "cli")) + + +def sign_in_time(): + # The CLI rewrites credentials.toml on sign-in, so a status fetched before + # that write describes the account signed in before. The fetch can land a + # beat before the file is flushed, hence the small allowance. + try: + return credentials_path().stat().st_mtime - 30 + except OSError: + return None + + +def limit_window_open(entry, now): raw = str(entry.get("resetsAt") or "") if raw == "": return True try: - resets_at = dt.datetime.fromisoformat(raw.replace("Z", "+00:00")) - except Exception: + 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=dt.timezone.utc) + resets_at = resets_at.replace(tzinfo=timezone.utc) return resets_at > now -def parse_quota() -> dict[str, Any]: - """Read the user-status caches and return limits + tier label. +# A window whose reset has passed is untouched, not unknown: it comes back at +# zero with no reset time until the CLI fetches the next one. +def usable_limits(entries): + now = datetime.now(timezone.utc) + return [entry if limit_window_open(entry, now) else dict(entry, percent=0.0, resetsAt="") for entry in entries] + + +def read_quota(): + """The newest user-status cache fetched under the current sign-in, or None. - The PlanStatus nested message carries daily/weekly quota percentages and - reset timestamps. Multiple cache files may exist (one per sign-in session), - and the CLI exposes no active-identity marker. The freshest fetch wins, but - files fetched before the last credentials write belong to a signed-out - account: after a switch with no fetch since, nothing qualifies and no quota - is reported rather than the previous account's. Without a credentials file - to date the sign-in, stale data beats none. + The CLI keeps one cache per sign-in and no marker of which is active, so + the newest fetch wins, and fetches from before the last sign-in are dropped: + after a switch with no fetch since, nothing qualifies rather than the + previous account's numbers. Without a credentials file to date the sign-in, + the newest fetch stands. """ - candidates: list[tuple[float, dict[str, Any]]] = [] - for path in sorted(user_status_dir().glob("user_status.*.bin")): + cutoff = sign_in_time() + best = None + try: + paths = sorted(user_status_dir().glob("user_status.*.bin")) + except OSError: + paths = [] + for path in paths: + cache = read_json(path) + if not isinstance(cache, dict) or not isinstance(cache.get("payload"), str): + continue + fetched_at = epoch_seconds(cache.get("fetched_at_secs")) + if fetched_at is None or (cutoff is not None and fetched_at < cutoff): + continue try: - cache = json.loads(path.read_text(encoding="utf-8")) - fetched_at = number(cache.get("fetched_at_secs")) - payload_b64 = cache.get("payload") - if not payload_b64: - continue - payload = base64.b64decode(payload_b64) - plan_status = _find_field(payload, 13) - if not isinstance(plan_status, (bytes, bytearray)): - continue - plan_info = _find_field(plan_status, 1) - plan_name = None - if isinstance(plan_info, (bytes, bytearray)): - raw_name = _find_field(plan_info, 2) - if isinstance(raw_name, (bytes, bytearray)): - plan_name = raw_name.decode("utf-8", errors="replace") - daily_remaining_pct = _find_field(plan_status, 14) - weekly_remaining_pct = _find_field(plan_status, 15) - daily_reset = _find_field(plan_status, 17) - weekly_reset = _find_field(plan_status, 18) - limits: list[dict[str, Any]] = [] - # Both fields are *remaining* percent from GetUserStatus; the panel - # expects *used*, so invert: 99% remaining -> 1% used. - if isinstance(daily_remaining_pct, int) and daily_remaining_pct >= 0: - limits.append({ - "label": "Daily", - "percent": (100 - daily_remaining_pct) / 100.0, - "resetsAt": _iso_from_unix(daily_reset if isinstance(daily_reset, int) else None), - }) - if isinstance(weekly_remaining_pct, int) and weekly_remaining_pct >= 0: - limits.append({ - "label": "Weekly", - "percent": (100 - weekly_remaining_pct) / 100.0, - "resetsAt": _iso_from_unix(weekly_reset if isinstance(weekly_reset, int) else None), - }) - now = dt.datetime.now(dt.timezone.utc) - limits = [e for e in limits if _limit_window_open(e, now)] - candidates.append((fetched_at, { - "tierLabel": plan_name or "", - "limits": limits, - })) - except Exception: + quota = parse_user_status(base64.b64decode(cache["payload"], validate=True)) + except ValueError: continue - cutoff = credentials_mtime() - if cutoff is not None: - candidates = [c for c in candidates if c[0] >= cutoff] - if not candidates: - return {} - fetched_at, quota = max(candidates, key=lambda c: c[0]) - quota["limitsFetchedAt"] = fetched_at * 1000 - quota["limitsStale"] = time.time() - fetched_at > LIMITS_CURRENT_ENOUGH_SECONDS - return quota + if not quota["limits"] and not quota["tierLabel"]: + continue + if best is None or fetched_at > best[0]: + best = (fetched_at, quota) + if best is None: + return None + fetched_at, quota = best + age = time.time() - fetched_at + return { + "tierLabel": quota["tierLabel"], + "limits": usable_limits(quota["limits"]), + "limitsFetchedAt": round(fetched_at * 1000), + "limitsStale": not 0 <= age < LIMITS_CURRENT_ENOUGH_SECONDS, + } + +# ------------------------------------------------------------------ record -# -------------------------------------------------------------------- record +def build_record(stats): + quota = read_quota() + signed_in = credentials_path().is_file() + in_use = number(stats.get("totalPrompts")) > 0 or bool(stats.get("modelUsage")) -def base_record(**overrides: Any) -> dict[str, Any]: - record: dict[str, Any] = { + record = { "schemaVersion": 1, "id": AGENT_ID, "name": AGENT_NAME, - "updatedAt": dt.datetime.now(dt.timezone.utc).isoformat(), - "ready": False, - "hasLocalStats": False, + "updatedAt": datetime.now(timezone.utc).isoformat(), + "ready": in_use or signed_in or bool(quota and quota["limits"]), + "hasLocalStats": True, "tierLabel": "", "usageStatusText": "", "authHelpText": "", @@ -460,73 +764,26 @@ def base_record(**overrides: Any) -> dict[str, Any]: "limitsStale": False, "limitsFetchedAt": 0, } - record.update({ - "todayPrompts": 0, - "todaySessions": 0, - "todayTotalTokens": 0, - "todayTokensByModel": {}, - "recentDays": [{"date": day, "messageCount": 0} for day in recent_date_strings()], - "totalPrompts": 0, - "totalSessions": 0, - "activeDays": 0, - "activeDates": [], - "modelUsage": {}, - }) - record.update(overrides) + if quota: + record.update(quota) + elif signed_in: + record["usageStatusText"] = "Devin limits unavailable" + record["authHelpText"] = REFRESH_HELP + elif in_use: + record["usageStatusText"] = "Waiting for auth" + record["authHelpText"] = AUTH_HELP + record.update(stats) return record -def scan(force: bool = False, limits_only: bool = False) -> dict[str, Any]: - quota = parse_quota() - - path = db_path() - if not path.is_file(): - if quota.get("limits"): - # A live status cache means the CLI ran under this account; the database - # is absent or lives elsewhere, so this is "no local stats", not "not - # installed". - return base_record(ready=True, **quota) - return base_record(usageStatusText="Devin CLI not installed", authHelpText="Install it with: curl -fsSL https://cli.devin.ai/install.sh | bash", **quota) - - conn = None - try: - conn = sqlite3.connect(f"file:{path}?mode=ro", uri=True) - conn.row_factory = sqlite3.Row - conn.execute("PRAGMA query_only = ON") - stats = cached_scan(conn, force, limits_only) - ready = number(stats.get("totalPrompts")) > 0 or len(quota.get("limits", [])) > 0 - return base_record(ready=ready, hasLocalStats=True, **stats, **quota) - except Exception as error: - # A locked or damaged history database must not hide quota already read - # from the user-status cache; it only costs the local stats. - detail = str(error) if isinstance(error, sqlite3.Error) else f"{type(error).__name__}: {error}" - print(f"omarchy-agent-usage-devin: {detail}", file=sys.stderr) - return base_record( - ready=len(quota.get("limits", [])) > 0, - usageStatusText="Devin usage unavailable", - authHelpText=detail, - **quota, - ) - finally: - if conn is not None: - conn.close() - - -def main() -> int: - parser = argparse.ArgumentParser(description="Print the Devin CLI usage record as JSON") - parser.add_argument("--force", action="store_true") - parser.add_argument("--limits-only", action="store_true") +def main(): + parser = argparse.ArgumentParser(description="Print the Devin usage record as JSON") + parser.add_argument("--force", action="store_true", help="rescan local sessions, ignoring the scan cache") + parser.add_argument("--limits-only", action="store_true", help="reuse a recent local scan") args = parser.parse_args() - try: - record = scan(force=args.force, limits_only=args.limits_only) - except sqlite3.Error as error: - record = base_record(usageStatusText="Devin usage unavailable", authHelpText=str(error)) - print(f"omarchy-agent-usage-devin: {error}", file=sys.stderr) - except Exception as error: - record = base_record(usageStatusText="Devin usage unavailable", authHelpText=f"{type(error).__name__}: {error}") - print(f"omarchy-agent-usage-devin: {type(error).__name__}: {error}", file=sys.stderr) - print(json.dumps(record, separators=(",", ":"), sort_keys=True)) + max_age = 0 if args.force else (LIMITS_ONLY_REUSE_SECONDS if args.limits_only else SCAN_REUSE_SECONDS) + print(json.dumps(build_record(cached_local_stats(max_age)), separators=(",", ":"), sort_keys=True)) return 0 diff --git a/test/shell.d/agent-usage-devin-scanner-test.sh b/test/shell.d/agent-usage-devin-scanner-test.sh index 925277f1e5b..6b1ae5a2052 100644 --- a/test/shell.d/agent-usage-devin-scanner-test.sh +++ b/test/shell.d/agent-usage-devin-scanner-test.sh @@ -4,323 +4,325 @@ source "$(dirname "$0")/base-test.sh" require_command jq require_command python3 +require_command rg TEST_HOME=$(mktemp -d) trap 'rm -rf "$TEST_HOME"' EXIT -# Every collector run must use the test cache root: without this, scans would -# write to (or read from) the real XDG cache when it is set in the environment. -export XDG_CACHE_HOME="$TEST_HOME/.cache" +# Every path the collector reads or writes lives in the test home. +export HOME="$TEST_HOME" +export XDG_CACHE_HOME="$TEST_HOME/cache" +export XDG_DATA_HOME="$TEST_HOME/data" +export XDG_CONFIG_HOME="$TEST_HOME/config" +unset DEVIN_DATA_DIR DEVIN_CACHE_DIR PI_CODING_AGENT_DIR -DEVIN_DATA_DIR="$TEST_HOME/.local/share/devin/cli" -mkdir -p "$DEVIN_DATA_DIR" +collector="$ROOT/bin/omarchy-agent-usage-devin" +DATA_DIR="$XDG_DATA_HOME/devin/cli" +STATUS_DIR="$XDG_CACHE_HOME/devin/cli" +SCAN_CACHE="$XDG_CACHE_HOME/omarchy/agent-usage" now=$(date +%s) today=$(date +%Y-%m-%d) -yesterday=$(date -d yesterday +%Y-%m-%d 2>/dev/null || date -v-1d +%Y-%m-%d) +yesterday=$(date -d yesterday +%Y-%m-%d) -# Build a mock sessions.db with the schema the collector expects. -python3 - "$DEVIN_DATA_DIR/sessions.db" "$now" "$today" "$yesterday" <<'PY' -import json, sqlite3, sys +run() { + "$collector" "$@" +} + +# ------------------------------------------------------------ unused machine + +result=$(run) +[[ $(jq -r '.ready' <<<"$result") == "false" ]] || + fail "Devin collector is not ready on an unused machine" "$result" +[[ $(jq -c '[.limits, .usageStatusText, .authHelpText, .tierLabel, .todayTotalTokens, (.recentDays | length), .modelUsage, has("balance"), has("retryAdvised")]' <<<"$result") == '[[],"","","",0,7,{},false,false]' ]] || + fail "Devin collector prints an empty record on an unused machine" "$result" +[[ $(jq -r '.id + ":" + .name' <<<"$result") == "devin:Devin" ]] || + fail "Devin collector identifies itself" "$result" +pass "Devin collector prints an empty, not-ready record on an unused machine" + +# ------------------------------------------------------- the CLI's database -db_path, now, today, yesterday = sys.argv[1], int(sys.argv[2]), sys.argv[3], sys.argv[4] +mkdir -p "$DATA_DIR" +python3 - "$DATA_DIR/sessions.db" "$now" <<'PY' +import json, sqlite3, sys +db_path, now = sys.argv[1], int(sys.argv[2]) conn = sqlite3.connect(db_path) conn.execute("CREATE TABLE message_nodes (row_id INTEGER PRIMARY KEY, session_id TEXT, node_id INTEGER, parent_node_id INTEGER, chat_message TEXT, created_at INTEGER, metadata TEXT)") conn.execute("CREATE TABLE prompt_history (id INTEGER PRIMARY KEY, content TEXT, timestamp INTEGER, session_id TEXT, is_shell INTEGER)") -def assistant_msg(model, inp, out, cache_read=0, cache_write=0): - return json.dumps({ - "metadata": { - "metrics": { - "input_tokens": inp, - "output_tokens": out, - "cache_read_tokens": cache_read, - "cache_creation_tokens": cache_write, - }, - "generation_model": model, - } - }) - -# Today: two turns with glm-5-2, one with devin-pro -conn.execute("INSERT INTO message_nodes (session_id, node_id, chat_message, created_at) VALUES (?, 1, ?, ?)", - ("s1", assistant_msg("glm-5-2", 100, 20, 60, 0), now)) -conn.execute("INSERT INTO message_nodes (session_id, node_id, chat_message, created_at) VALUES (?, 2, ?, ?)", - ("s1", assistant_msg("glm-5-2", 80, 10, 50, 0), now)) -conn.execute("INSERT INTO message_nodes (session_id, node_id, chat_message, created_at) VALUES (?, 3, ?, ?)", - ("s1", assistant_msg("devin-pro", 200, 40, 0, 10), now)) - -# Yesterday: one turn -conn.execute("INSERT INTO message_nodes (session_id, node_id, chat_message, created_at) VALUES (?, 4, ?, ?)", - ("s2", assistant_msg("glm-5-2", 50, 5, 0, 0), now - 86400)) - -# A row with no metrics should be skipped -conn.execute("INSERT INTO message_nodes (session_id, node_id, chat_message, created_at) VALUES (?, 5, ?, ?)", - ("s1", json.dumps({"metadata": {}}), now)) - -# A malformed row must not abort the whole scan -conn.execute("INSERT INTO message_nodes (session_id, node_id, chat_message, created_at) VALUES (?, 6, ?, ?)", - ("s1", "{not json", now)) - -# Prompt history: 3 prompts today across 2 sessions, 2 yesterday — one prompt -# belongs to today's session s1, so per-day session counts sum to 4 while the -# true distinct total is 3 -conn.execute("INSERT INTO prompt_history (content, timestamp, session_id, is_shell) VALUES (?, ?, 's1', 0)", ("hello", now)) -conn.execute("INSERT INTO prompt_history (content, timestamp, session_id, is_shell) VALUES (?, ?, 's1', 0)", ("world", now)) -conn.execute("INSERT INTO prompt_history (content, timestamp, session_id, is_shell) VALUES (?, ?, 's3', 0)", ("test", now)) -conn.execute("INSERT INTO prompt_history (content, timestamp, session_id, is_shell) VALUES (?, ?, 's2', 0)", ("yesterday", now - 86400)) -conn.execute("INSERT INTO prompt_history (content, timestamp, session_id, is_shell) VALUES (?, ?, 's1', 0)", ("late night", now - 86400)) - +def node(session, created, model, inp, out, cache_read=0, cache_write=0): + metrics = {"input_tokens": inp, "output_tokens": out, "cache_read_tokens": cache_read, "cache_creation_tokens": cache_write} + conn.execute("INSERT INTO message_nodes (session_id, chat_message, created_at) VALUES (?, ?, ?)", + (session, json.dumps({"metadata": {"metrics": metrics, "generation_model": model}}), created)) + +node("s1", now, "glm-5-2", 100, 20, 60, 0) +node("s1", now, "glm-5-2", 80, 10, 50, 0) +node("s1", now, "devin-pro", 200, 40, 0, 10) +node("s2", now - 86400, "glm-5-2", 50, 5) +# Untrusted fields skip themselves, or the row, never the scan: a boolean, +# a negative count and a string count read as nothing; a row with a broken +# timestamp, no metrics, or no JSON at all adds nothing. +node("s1", now, "glm-5-2", True, -5, "7", 0) +node("s1", "yesterday-ish", "glm-5-2", 999, 999) +node("s1", 10**15, "glm-5-2", 999, 999) +conn.execute("INSERT INTO message_nodes (session_id, chat_message, created_at) VALUES ('s1', ?, ?)", (json.dumps({"metadata": {}}), now)) +conn.execute("INSERT INTO message_nodes (session_id, chat_message, created_at) VALUES ('s1', '{not json \"metrics\"', ?)", (now,)) + +for content, timestamp, session in (("a", now, "s1"), ("b", now, "s1"), ("c", now, "s3"), ("d", now - 86400, "s2"), ("e", now - 86400, "s1")): + conn.execute("INSERT INTO prompt_history (content, timestamp, session_id, is_shell) VALUES (?, ?, ?, 0)", (content, timestamp, session)) conn.commit() conn.close() PY -result=$(HOME="$TEST_HOME" DEVIN_DATA_DIR="$DEVIN_DATA_DIR" \ - "$ROOT/bin/omarchy-agent-usage-devin") - -# Today: (100+20+60) + (80+10+50) + (200+40+0+10) = 570 +result=$(run --force) +# (100+20+60) + (80+10+50) + (200+40+10) = 570; the boolean and negative +# counts read as nothing, so the odd row adds nothing. [[ $(jq -r '.todayTotalTokens' <<<"$result") == "570" ]] || fail "Devin collector counts today's tokens across models" "$result" -pass "Devin collector counts today's tokens across models" - -# 3 prompts today across 2 sessions -[[ $(jq -r '.todayPrompts' <<<"$result") == "3" ]] || - fail "Devin collector counts today's prompts" "$result" -[[ $(jq -r '.todaySessions' <<<"$result") == "2" ]] || - fail "Devin collector counts today's sessions" "$result" -pass "Devin collector counts today's prompts and sessions" - -# Model usage (all-time): glm-5-2 = (100+80+50) input, (20+10+5) output, (60+50+0) cache_read, 0 cache_write -# devin-pro = 200 input, 40 output, 0 cache_read, 10 cache_write +[[ $(jq -c '[.todayPrompts, .todaySessions, .totalPrompts, .totalSessions, .activeDays]' <<<"$result") == "[3,2,5,3,2]" ]] || + fail "Devin collector counts prompts from prompt_history and a midnight-spanning session once" "$result" [[ $(jq -c '.modelUsage["glm-5-2"]' <<<"$result") == '{"cacheCreationInputTokens":0,"cacheReadInputTokens":110,"inputTokens":230,"outputTokens":35}' ]] || fail "Devin collector aggregates glm-5-2 model usage" "$result" [[ $(jq -c '.modelUsage["devin-pro"]' <<<"$result") == '{"cacheCreationInputTokens":10,"cacheReadInputTokens":0,"inputTokens":200,"outputTokens":40}' ]] || fail "Devin collector aggregates devin-pro model usage" "$result" -pass "Devin collector aggregates model usage correctly" - -# Total prompts = 5 (3 today + 2 yesterday); total sessions = 3 distinct -# (s1, s2, s3) — s1 spans midnight and must not count twice -[[ $(jq -r '.totalPrompts' <<<"$result") == "5" ]] || - fail "Devin collector counts total prompts" "$result" -[[ $(jq -r '.totalSessions' <<<"$result") == "3" ]] || - fail "Devin collector counts a midnight-spanning session once" "$result" -pass "Devin collector counts total prompts and dedupes cross-day sessions" - -# Active days = 2 (today + yesterday) -[[ $(jq -r '.activeDays' <<<"$result") == "2" ]] || - fail "Devin collector counts active days" "$result" -pass "Devin collector counts active days" - -# Recent days: today should have 570, yesterday should have 55 (50+5) -[[ $(jq -r '[.recentDays[] | select(.date == "'$today'") | .messageCount] | .[0]' <<<"$result") == "570" ]] || - fail "Devin collector fills today's recent day" "$result" -[[ $(jq -r '[.recentDays[] | select(.date == "'$yesterday'") | .messageCount] | .[0]' <<<"$result") == "55" ]] || +[[ $(jq -r --arg d "$yesterday" '.recentDays[] | select(.date == $d) | .messageCount' <<<"$result") == "55" ]] || fail "Devin collector fills yesterday's recent day" "$result" -pass "Devin collector fills the seven-day token series" +pass "Devin collector reads the CLI's session database, skipping untrusted fields" -# No limits -[[ $(jq -c '.limits' <<<"$result") == "[]" ]] || - fail "Devin collector reports no limits" "$result" -pass "Devin collector reports no limits" +# In use without any sign-in: the quota has nowhere to come from. +[[ $(jq -c '[.ready, .usageStatusText, .limits]' <<<"$result") == '[true,"Waiting for auth",[]]' ]] || + fail "Devin collector waits for auth when the CLI was never signed in" "$result" +pass "Devin collector waits for auth when in use without a sign-in" -# Identifies itself -[[ $(jq -r '.id' <<<"$result") == "devin" ]] || - fail "Devin collector identifies itself" "$result" -[[ $(jq -r '.name' <<<"$result") == "Devin" ]] || - fail "Devin collector names itself" "$result" -pass "Devin collector identifies itself" - -# Missing database -rm "$DEVIN_DATA_DIR/sessions.db" -result=$(HOME="$TEST_HOME" DEVIN_DATA_DIR="$DEVIN_DATA_DIR" \ - "$ROOT/bin/omarchy-agent-usage-devin") -[[ $(jq -r '.usageStatusText' <<<"$result") == "Devin CLI not installed" ]] || - fail "Devin collector reports missing database" "$result" -[[ $(jq -r '.authHelpText' <<<"$result") != "" ]] || - fail "Devin collector shows install hint for missing database" "$result" -pass "Devin collector reports missing database gracefully" - -# Caching: a second run without --force should reuse the cached scan -# (rebuild the DB since we deleted it above) -mkdir -p "$DEVIN_DATA_DIR" -python3 - "$DEVIN_DATA_DIR/sessions.db" "$now" <<'PY' -import json, sqlite3, sys -db_path, now = sys.argv[1], int(sys.argv[2]) -conn = sqlite3.connect(db_path) -conn.execute("CREATE TABLE message_nodes (row_id INTEGER PRIMARY KEY, session_id TEXT, node_id INTEGER, chat_message TEXT, created_at INTEGER, metadata TEXT)") -conn.execute("CREATE TABLE prompt_history (id INTEGER PRIMARY KEY, content TEXT, timestamp INTEGER, session_id TEXT, is_shell INTEGER)") -conn.execute("INSERT INTO message_nodes (session_id, node_id, chat_message, created_at) VALUES ('s1', 1, ?, ?)", - (json.dumps({"metadata": {"metrics": {"input_tokens": 10, "output_tokens": 5, "cache_read_tokens": 0, "cache_creation_tokens": 0}, "generation_model": "test"}}), now)) -conn.execute("INSERT INTO prompt_history (content, timestamp, session_id, is_shell) VALUES ('hi', ?, 's1', 0)", (now,)) -conn.commit() -conn.close() -PY - -# Clear the scan cache the earlier scenarios wrote so this run starts cold -rm -f "$TEST_HOME/.cache/omarchy/agent-usage/devin-scan-"* - -# First run populates the cache -result1=$(HOME="$TEST_HOME" DEVIN_DATA_DIR="$DEVIN_DATA_DIR" \ - "$ROOT/bin/omarchy-agent-usage-devin") -[[ $(jq -r '.todayTotalTokens' <<<"$result1") == "15" ]] || - fail "Devin collector first scan populates cache" "$result1" +# ---------------------------------------------------------- scan reuse -# Add more data after the cache is written -python3 - "$DEVIN_DATA_DIR/sessions.db" "$now" <<'PY' +python3 - "$DATA_DIR/sessions.db" "$now" <<'PY' import json, sqlite3, sys -db_path, now = sys.argv[1], int(sys.argv[2]) -conn = sqlite3.connect(db_path) -conn.execute("INSERT INTO message_nodes (session_id, node_id, chat_message, created_at) VALUES ('s1', 2, ?, ?)", - (json.dumps({"metadata": {"metrics": {"input_tokens": 100, "output_tokens": 50, "cache_read_tokens": 0, "cache_creation_tokens": 0}, "generation_model": "test"}}), now)) +conn = sqlite3.connect(sys.argv[1]) +conn.execute("INSERT INTO message_nodes (session_id, chat_message, created_at) VALUES ('s1', ?, ?)", + (json.dumps({"metadata": {"metrics": {"input_tokens": 1000, "output_tokens": 0}, "generation_model": "glm-5-2"}}), int(sys.argv[2]))) conn.commit() conn.close() PY -# Second run without --force should still serve the cached value (15, not 165) -result2=$(HOME="$TEST_HOME" DEVIN_DATA_DIR="$DEVIN_DATA_DIR" \ - "$ROOT/bin/omarchy-agent-usage-devin") -[[ $(jq -r '.todayTotalTokens' <<<"$result2") == "15" ]] || - fail "Devin collector serves cached scan without --force" "$result2" -pass "Devin collector serves cached scan without --force" - -# --force bypasses the cache and sees the new data (165) -result3=$(HOME="$TEST_HOME" DEVIN_DATA_DIR="$DEVIN_DATA_DIR" \ - "$ROOT/bin/omarchy-agent-usage-devin" --force) -[[ $(jq -r '.todayTotalTokens' <<<"$result3") == "165" ]] || - fail "Devin collector --force rescans past the cache" "$result3" -pass "Devin collector --force rescans past the cache" - -# Quota / limits from user-status cache -# Build a minimal protobuf PlanStatus: field 1 (PlanInfo) with field 2 (plan name "Pro"), -# field 14 (daily_remaining_pct=69), field 15 (weekly_remaining_pct=53), -# field 17 (daily_reset_unix), field 18 (weekly_reset_unix) -python3 - "$TEST_HOME/.cache/devin/cli" "$now" <<'PY' -import base64, json, os, struct, sys - -cache_dir, now = sys.argv[1], int(sys.argv[2]) -os.makedirs(cache_dir, exist_ok=True) - -def varint(val): - out = b"" - while True: - byte = val & 0x7F - val >>= 7 - if val: - out += bytes([byte | 0x80]) - else: - out += bytes([byte]) - break - return out - -def field(fn, data): - if isinstance(data, int): - return varint((fn << 3) | 0) + varint(data) - return varint((fn << 3) | 2) + varint(len(data)) + data - -# PlanInfo: field 2 = "Pro" -plan_info = field(2, b"Pro") -# PlanStatus: field 1 = PlanInfo, field 14 = 69 (remaining), field 15 = 53 (remaining), -# field 17 = now+3600, field 18 = now+86400*4 -plan_status = field(1, plan_info) + field(14, 69) + field(15, 53) + field(17, now + 3600) + field(18, now + 86400 * 4) -# Top-level: field 13 = PlanStatus -payload = field(13, plan_status) - -cache = { - "version": 1, - "identity_digest": "test123", - "fetched_at_secs": now, - "payload": base64.b64encode(payload).decode("ascii"), -} -with open(os.path.join(cache_dir, "user_status.test123.bin"), "w") as f: - json.dump(cache, f) +[[ $(jq -r '.todayTotalTokens' <<<"$(run)") == "570" ]] || + fail "Devin collector reuses a seconds-old scan" "$(run)" +[[ $(jq -r '.todayTotalTokens' <<<"$(run --limits-only)") == "570" ]] || + fail "Devin collector reuses a scan for --limits-only" "$(run --limits-only)" +[[ $(jq -r '.todayTotalTokens' <<<"$(run --force)") == "1570" ]] || + fail "Devin collector --force rescans past the cache" "$(run --force)" +pass "Devin collector reuses a recent scan and --force rescans" + +# A cache from another day holds another day's today* numbers. +cache_file=$(ls "$SCAN_CACHE"/devin-scan-*.json | head -n 1) +jq -c '.scanDate = "1999-01-01" | .stats.todayTotalTokens = 1' "$cache_file" >"$cache_file.tmp" && mv "$cache_file.tmp" "$cache_file" +[[ $(jq -r '.todayTotalTokens' <<<"$(run --limits-only)") == "1570" ]] || + fail "Devin collector treats a scan from another day as a miss" "$(run --limits-only)" +pass "Devin collector treats a scan from another day as a miss" + +# A database it cannot read costs the stats, never the record, and is not +# cached as the whole story. +cp "$DATA_DIR/sessions.db" "$TEST_HOME/sessions.db.good" +echo "not a database" >"$DATA_DIR/sessions.db" +rm -f "$SCAN_CACHE"/devin-scan-*.json +result=$(run --force 2>/dev/null) +[[ $(jq -r '.todayTotalTokens' <<<"$result") == "0" && $(jq -r '.id' <<<"$result") == "devin" ]] || + fail "Devin collector survives an unreadable session database" "$result" +[[ -z $(ls "$SCAN_CACHE"/devin-scan-*.json 2>/dev/null) ]] || + fail "Devin collector does not cache a scan that could not read the database" "$result" +pass "Devin collector survives an unreadable database without caching it" +mv "$TEST_HOME/sessions.db.good" "$DATA_DIR/sessions.db" + +# ------------------------------------------------------------- quota cache + +# write_status [key=value...]: a user-status cache the CLI +# would write, built from the GetUserStatus fields the collector reads. +write_status() { + python3 - "$STATUS_DIR" "$@" <<'PY' +import base64, json, os, sys + +directory, name, fetched_at = sys.argv[1], sys.argv[2], int(sys.argv[3]) +options = dict(arg.split("=", 1) for arg in sys.argv[4:]) +os.makedirs(directory, exist_ok=True) + +def varint(value): + value &= (1 << 64) - 1 + out = b"" + while True: + byte = value & 0x7F + value >>= 7 + out += bytes([byte | 0x80]) if value else bytes([byte]) + if not value: + return out + +def field(number, data): + if isinstance(data, int): + return varint(number << 3) + varint(data) + return varint((number << 3) | 2) + varint(len(data)) + data + +plan = b"" +if options.get("plan"): + plan += field(2, options["plan"].encode()) +for key, number in (("tier", 1), ("billing", 35), ("hideDaily", 36), ("hideWeekly", 37), ("promptLimit", 12)): + if key in options: + plan += field(number, int(options[key])) +status = field(1, plan) +for key, number in (("daily", 14), ("weekly", 15), ("dailyReset", 17), ("weeklyReset", 18), ("promptUsed", 6)): + if key in options: + status += field(number, int(options[key])) +if "planEnd" in options: + status += field(3, field(1, int(options["planEnd"]))) +payload = field(13, status) +if "userTier" in options: + payload = field(10, int(options["userTier"])) + payload + +with open(os.path.join(directory, f"user_status.{name}.bin"), "w") as handle: + json.dump({"version": 1, "identity_digest": name, "fetched_at_secs": fetched_at, "payload": base64.b64encode(payload).decode()}, handle) PY +} -result=$(HOME="$TEST_HOME" DEVIN_DATA_DIR="$DEVIN_DATA_DIR" \ - "$ROOT/bin/omarchy-agent-usage-devin" --force) -[[ $(jq -r '.tierLabel' <<<"$result") == "Pro" ]] || - fail "Devin collector reads plan name from user-status cache" "$result" -[[ $(jq -r '.limits[0].label' <<<"$result") == "Daily" ]] || - fail "Devin collector reports daily limit label" "$result" -[[ $(jq -r '.limits[0].percent' <<<"$result") == "0.31" ]] || - fail "Devin collector reports daily limit percent (31% used = 69% remaining)" "$result" -[[ $(jq -r '.limits[1].label' <<<"$result") == "Weekly" ]] || - fail "Devin collector reports weekly limit label" "$result" -[[ $(jq -r '.limits[1].percent' <<<"$result") == "0.47" ]] || - fail "Devin collector reports weekly limit percent (47% used = 53% remaining)" "$result" -[[ $(jq -r '.limitsStale' <<<"$result") == "false" ]] || - fail "Devin collector marks a just-fetched status cache as current" "$result" -[[ $(jq -r '.limitsFetchedAt' <<<"$result") == "$((now * 1000))" ]] || - fail "Devin collector reports when the status cache was fetched" "$result" -pass "Devin collector parses quota from user-status cache" - -# Quota without a session database: the status cache alone means the CLI has -# run, so the record is ready with limits rather than "not installed" -rm "$DEVIN_DATA_DIR/sessions.db" -result=$(HOME="$TEST_HOME" DEVIN_DATA_DIR="$DEVIN_DATA_DIR" \ - "$ROOT/bin/omarchy-agent-usage-devin" --force) -[[ $(jq -r '.ready' <<<"$result") == "true" ]] || - fail "Devin collector stays ready on quota without a session database" "$result" -[[ $(jq -r '.limits | length' <<<"$result") == "2" ]] || - fail "Devin collector keeps limits without a session database" "$result" -[[ $(jq -r '.usageStatusText' <<<"$result") == "" ]] || - fail "Devin collector drops the install hint when quota proves the CLI ran" "$result" -pass "Devin collector reports quota without a session database" - -# Status caches fetched under a previous sign-in are dropped: once -# credentials.toml is newer than every fetch, no quota is reported rather than -# the signed-out account's -creds="$TEST_HOME/.local/share/devin/credentials.toml" -echo 'api_server_url = "https://api.devin.ai"' >"$creds" -touch -d "@$((now + 7200))" "$creds" -result=$(HOME="$TEST_HOME" DEVIN_DATA_DIR="$DEVIN_DATA_DIR" \ - "$ROOT/bin/omarchy-agent-usage-devin" --force) -[[ $(jq -r '.tierLabel' <<<"$result") == "" ]] || +write_status a "$now" plan=Pro daily=69 weekly=53 dailyReset=$((now + 3600)) weeklyReset=$((now + 4 * 86400)) +result=$(run) +[[ $(jq -c '[.tierLabel, [.limits[] | [.label, .percent]], .limitsStale, .limitsFetchedAt, .usageStatusText, .authHelpText]' <<<"$result") == "[\"Pro\",[[\"Daily\",0.31],[\"Weekly\",0.47]],false,$((now * 1000)),\"\",\"\"]" ]] || + fail "Devin collector reads fresh quota from the user-status cache" "$result" +[[ $(jq -r '.limits[0].resetsAt' <<<"$result") == "$(date -u -d "@$((now + 3600))" +%Y-%m-%dT%H:%M:%S+00:00)" ]] || + fail "Devin collector keeps the service's own reset time" "$result" +[[ $(jq -r 'has("retryAdvised")' <<<"$result") == "false" ]] || + fail "Devin collector never advises a retry, having no check to retry" "$result" +pass "Devin collector reads fresh quota from the CLI's cache" + +# The quota is read on every run, past the scan reuse, so a fetch the CLI +# just made shows without --force. +write_status a "$now" plan=Pro daily=50 weekly=53 dailyReset=$((now + 3600)) weeklyReset=$((now + 4 * 86400)) +[[ $(jq -r '.limits[0].percent' <<<"$(run)") == "0.5" ]] || + fail "Devin collector reads the quota cache on every run" "$(run)" +pass "Devin collector reads the quota cache on every run" + +# Spent windows: proto3 leaves a zero off the wire, so a dated window with +# no percent has nothing left. +write_status a "$now" plan=Pro dailyReset=$((now + 3600)) +[[ $(jq -c '[.limits[] | [.label, .percent]]' <<<"$(run)") == '[["Daily",1.0]]' ]] || + fail "Devin collector reads a missing percent on a dated window as spent" "$(run)" +pass "Devin collector reads a missing percent on a dated window as spent" + +# A cache from more than fifteen minutes ago still describes the account, but +# is reported stale with its age; a window that reset since is back at zero. +write_status a $((now - 7200)) plan=Pro daily=10 weekly=40 dailyReset=$((now - 60)) weeklyReset=$((now + 86400)) +result=$(run) +[[ $(jq -c '[.limitsStale, .limitsFetchedAt, [.limits[] | [.label, .percent, .resetsAt != ""]], .usageStatusText]' <<<"$result") == "[true,$(((now - 7200) * 1000)),[[\"Daily\",0.0,false],[\"Weekly\",0.6,true]],\"\"]" ]] || + fail "Devin collector keeps an old quota stale and zeroes an elapsed window" "$result" +pass "Devin collector reports an old quota stale, with elapsed windows at zero" + +# Fetches from before the last sign-in describe another account. +creds="$XDG_DATA_HOME/devin/credentials.toml" +echo 'api_server_url = "https://example.invalid"' >"$creds" +touch -d "@$((now + 600))" "$creds" +write_status a "$now" plan=Pro daily=69 weekly=53 dailyReset=$((now + 3600)) weeklyReset=$((now + 86400)) +result=$(run) +[[ $(jq -c '[.ready, .tierLabel, .limits, .usageStatusText]' <<<"$result") == '[true,"",[],"Devin limits unavailable"]' ]] || fail "Devin collector drops quota fetched before the last sign-in" "$result" -[[ $(jq -c '.limits' <<<"$result") == "[]" ]] || - fail "Devin collector drops limits fetched before the last sign-in" "$result" -pass "Devin collector drops status caches from before the last sign-in" - -# A credentials file older than the caches leaves the freshest fetch winning -touch -d "@$((now - 7200))" "$creds" -result=$(HOME="$TEST_HOME" DEVIN_DATA_DIR="$DEVIN_DATA_DIR" \ - "$ROOT/bin/omarchy-agent-usage-devin" --force) -[[ $(jq -r '.tierLabel' <<<"$result") == "Pro" ]] || - fail "Devin collector serves quota fetched under the current credentials" "$result" -pass "Devin collector serves quota fetched under the current credentials" - -# A status cache older than the freshness window keeps its limits but is -# flagged stale, with its fetch time reported so the panel can warn -python3 - "$TEST_HOME/.cache/devin/cli/user_status.test123.bin" "$now" <<'PY' -import json, sys -path = sys.argv[1] -with open(path) as f: - cache = json.load(f) -cache["fetched_at_secs"] = int(sys.argv[2]) - 3600 -with open(path, "w") as f: - json.dump(cache, f) +[[ $(jq -r '.authHelpText' <<<"$result") == *"Start Devin"* ]] || + fail "Devin collector says how to refresh the quota after a sign-in" "$result" +pass "Devin collector drops quota fetched before the last sign-in" + +# Of the caches fetched since, the newest wins. +touch -d "@$((now - 600))" "$creds" +write_status b $((now - 300)) plan=Max daily=0 weekly=0 dailyReset=$((now + 3600)) +[[ $(jq -r '.tierLabel' <<<"$(run)") == "Pro" ]] || + fail "Devin collector serves the newest fetch under the current sign-in" "$(run)" +rm "$STATUS_DIR/user_status.b.bin" +pass "Devin collector serves the newest fetch under the current sign-in" + +# A plan can hide a window; credit buckets run to the plan's end; a plan +# without a name of its own is named by its tier. +write_status a "$now" tier=16 hideDaily=1 daily=10 weekly=80 dailyReset=$((now + 3600)) weeklyReset=$((now + 86400)) promptLimit=500 promptUsed=125 planEnd=$((now + 20 * 86400)) +result=$(run) +[[ $(jq -c '[.tierLabel, [.limits[] | [.label, .percent]]]' <<<"$result") == '["Devin Pro",[["Weekly",0.2],["Prompt credits",0.25]]]' ]] || + fail "Devin collector honors hidden windows, credit buckets and tier names" "$result" +pass "Devin collector honors hidden windows, credit buckets and tier names" + +# A credit plan leaves the quota percents at their defaults with no reset: +# those windows are not meters. +write_status a "$now" plan=Pro billing=1 +[[ $(jq -c '.limits' <<<"$(run)") == "[]" ]] || + fail "Devin collector shows no quota windows on a credit plan" "$(run)" +pass "Devin collector shows no quota windows on a credit plan" + +# A broken cache file is no quota at all. +echo '{"fetched_at_secs": true, "payload": "%%%"}' >"$STATUS_DIR/user_status.a.bin" +[[ $(jq -r '.usageStatusText' <<<"$(run)") == "Devin limits unavailable" ]] || + fail "Devin collector ignores a broken quota cache" "$(run)" +pass "Devin collector ignores a broken quota cache" +rm -rf "$STATUS_DIR" "$creds" + +# ------------------------------------------------------------- pi and omp + +rm -rf "$DATA_DIR" +pi_sessions="$TEST_HOME/.pi/agent/sessions/--work--" +omp_profile="$TEST_HOME/.omp/profiles/work/agent/sessions/--work--" +mkdir -p "$pi_sessions" "$omp_profile" +stamp="${today}T12:00:00.000Z" +pi_line() { + local id=$1 provider=$2 input=$3 output=$4 + printf '{"type":"message","id":"%s","timestamp":"%s","message":{"role":"assistant","provider":"%s","model":"swe-1.6","usage":{"input":%s,"output":%s,"cacheRead":10,"cacheWrite":0}}}\n' "$id" "$stamp" "$provider" "$input" "$output" +} +{ + pi_line m1 devin 100 20 + pi_line m2 devin 200 40 + pi_line m3 anthropic 999 999 +} >"$pi_sessions/one.jsonl" +# A fork repeats its parent's message with the same id and timestamp. +{ + pi_line m2 devin 200 40 + pi_line m4 devin 1 1 +} >"$omp_profile/fork.jsonl" +# A $HOME that is a git checkout ignoring everything must not hide sessions. +printf '*\n' >"$TEST_HOME/.gitignore" +git -C "$TEST_HOME" init -q + +result=$(run --force) +# m1 130 + m2 250 + m4 12; the fork's copy of m2 and the anthropic message +# are not Devin's to count. +[[ $(jq -c '[.todayTotalTokens, .todayPrompts, .todaySessions, .modelUsage["swe-1.6"].inputTokens]' <<<"$result") == "[392,3,2,301]" ]] || + fail "Devin collector counts pi/omp devin messages once, profiles included" "$result" +pass "Devin collector counts pi and omp sessions on the devin provider once" + +# PI_CODING_AGENT_DIR moves pi's sessions. +moved="$TEST_HOME/elsewhere" +mkdir -p "$moved/sessions/--work--" +pi_line m9 devin 5 5 >"$moved/sessions/--work--/moved.jsonl" +[[ $(jq -r '.todayTotalTokens' <<<"$(PI_CODING_AGENT_DIR="$moved" run --force)") == "282" ]] || + fail "Devin collector follows PI_CODING_AGENT_DIR" "$(PI_CODING_AGENT_DIR="$moved" run --force)" +pass "Devin collector follows PI_CODING_AGENT_DIR" +rm -rf "$TEST_HOME/.pi" "$TEST_HOME/.omp" + +# --------------------------------------------------------- the older layout + +# Builds without the cli/ directory kept one lifetime total per session. +python3 - "$XDG_DATA_HOME/devin/cli_sessions.db" "$now" <<'PY' +import sqlite3, sys +conn = sqlite3.connect(sys.argv[1]); now = int(sys.argv[2]) +conn.execute("CREATE TABLE sessions (id TEXT, working_directory TEXT, backend_type TEXT, model TEXT, permission_mode TEXT, created_at INTEGER, last_activity_at INTEGER, total_tokens INTEGER, current_epoch INTEGER)") +conn.execute("CREATE TABLE messages (id TEXT, session_id TEXT, sequence_number INTEGER, role TEXT, content TEXT, tool_call_id TEXT, tool_calls_json TEXT, timestamp INTEGER, compaction_epoch INTEGER)") +conn.execute("INSERT INTO sessions (id, model, created_at, last_activity_at, total_tokens) VALUES ('o1', 'swe-1.5', ?, ?, 4000)", (now - 86400, now)) +conn.execute("INSERT INTO messages (id, session_id, role, timestamp) VALUES ('1', 'o1', 'user', ?), ('2', 'o1', 'assistant', ?), ('3', 'o1', 'assistant', ?)", (now, now, now)) +conn.commit(); conn.close() PY - -result=$(HOME="$TEST_HOME" DEVIN_DATA_DIR="$DEVIN_DATA_DIR" \ - "$ROOT/bin/omarchy-agent-usage-devin" --force) -[[ $(jq -r '.limitsStale' <<<"$result") == "true" ]] || - fail "Devin collector flags an old status cache as stale" "$result" -[[ $(jq -r '.limitsFetchedAt' <<<"$result") == "$(( (now - 3600) * 1000 ))" ]] || - fail "Devin collector reports the stale cache's fetch time" "$result" -[[ $(jq -r '.limits | length' <<<"$result") == "2" ]] || - fail "Devin collector keeps limits from a stale cache" "$result" -pass "Devin collector flags old quota as stale with its fetch time" - -# A damaged session database must not hide working quota: the scan fails, but -# the limits already read from the user-status cache still render, with the -# database error reported alongside -rm -f "$TEST_HOME/.cache/omarchy/agent-usage/devin-scan-"* -echo "not a database" > "$DEVIN_DATA_DIR/sessions.db" -result=$(HOME="$TEST_HOME" DEVIN_DATA_DIR="$DEVIN_DATA_DIR" \ - "$ROOT/bin/omarchy-agent-usage-devin" --force) -[[ $(jq -r '.usageStatusText' <<<"$result") == "Devin usage unavailable" ]] || - fail "Devin collector reports an unreadable session database" "$result" -[[ $(jq -r '.ready' <<<"$result") == "true" ]] || - fail "Devin collector stays ready on quota with an unreadable database" "$result" -[[ $(jq -r '.limits | length' <<<"$result") == "2" ]] || - fail "Devin collector keeps quota when the session database is unreadable" "$result" -pass "Devin collector keeps quota when the session database is unreadable" +result=$(run --force) +[[ $(jq -c '[.todayTotalTokens, .todayPrompts, .totalSessions, .modelUsage["swe-1.5"].inputTokens]' <<<"$result") == "[4000,2,1,4000]" ]] || + fail "Devin collector falls back to the older per-session layout" "$result" +pass "Devin collector falls back to the older per-session layout" + +# Once the CLI's own database exists, the older files are not read again. +mkdir -p "$DATA_DIR" +python3 - "$DATA_DIR/sessions.db" <<'PY' +import sqlite3, sys +conn = sqlite3.connect(sys.argv[1]) +conn.execute("CREATE TABLE message_nodes (row_id INTEGER PRIMARY KEY, session_id TEXT, chat_message TEXT, created_at INTEGER)") +conn.execute("CREATE TABLE prompt_history (id INTEGER PRIMARY KEY, timestamp INTEGER, session_id TEXT)") +conn.commit(); conn.close() +PY +[[ $(jq -r '.todayTotalTokens' <<<"$(run --force)") == "0" ]] || + fail "Devin collector reads only the CLI's own database when it exists" "$(run --force)" +pass "Devin collector reads only the CLI's own database when it exists" From f2c4fc35b5ddd20aec86403943db0df5e8966d29 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:13:40 -0400 Subject: [PATCH 093/546] Bring the Synthetic collector in line with the collector contract Count pi, omp and opencode usage on the synthetic provider, and read the key from omp too. Keep each key's last quotas in its own cache, reuse a check from the last fifteen seconds unless forced, and keep the last quotas stale past a failed check with elapsed windows at zero. Advise a retry only when the endpoint couldn't be reached, and report a refused key as an expired sign-in. Leave the plan out of the tier rather than the credits left, drop the detail and reset-label fields the panel doesn't read, and don't read the regeneration ticks as reset times. Co-authored-by: aserper Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-synthetic | 945 +++++++++++++++------ test/shell.d/agent-usage-synthetic-test.sh | 361 +++++--- 2 files changed, 916 insertions(+), 390 deletions(-) diff --git a/bin/omarchy-agent-usage-synthetic b/bin/omarchy-agent-usage-synthetic index b101706e37d..7c196a050a9 100755 --- a/bin/omarchy-agent-usage-synthetic +++ b/bin/omarchy-agent-usage-synthetic @@ -1,314 +1,757 @@ #!/usr/bin/python3 -# omarchy:summary=Print the Synthetic Labs usage record as JSON +# omarchy:summary=Print the Synthetic usage record as JSON # omarchy:args=[--force] [--limits-only] # omarchy:hidden=true -"""Collect Synthetic Labs quotas into one display-ready JSON record. +"""Collect Synthetic usage into one display-ready JSON record. -The account-global limits come from Synthetic's /v2/quotas endpoint. There -are no local transcript stats: the agents panel displays the weekly credit, -rolling request, and search allowances this collector prints. +Local stats come from pi/omp sessions and opencode messages that ran on the +synthetic provider. The account's allowances (the rolling five-hour request +limit, the weekly credits, search and free tool calls) come from Synthetic's +documented GET https://api.synthetic.new/v2/quotas, which doesn't count +against them. The agents panel only ever reads the JSON this prints. """ -from __future__ import annotations - import argparse +import fcntl +import hashlib import json import math import os +import re +import shutil +import sqlite3 +import subprocess import sys +import tempfile +import time import urllib.error import urllib.request -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from pathlib import Path -from typing import Any AGENT_ID = "synthetic" -AGENT_NAME = "Synthetic Labs" +AGENT_NAME = "Synthetic" +PROVIDER = "synthetic" API_BASE_URL = "https://api.synthetic.new" -AUTH_HELP = "Set SYNTHETIC_API_KEY or sign in to Synthetic in Pi or OpenCode." +AUTH_HELP = "Set SYNTHETIC_API_KEY, or sign in to Synthetic in pi, omp or opencode, to read its quotas." +REFUSED_HELP = "Synthetic rejected the API key. Set a current one in SYNTHETIC_API_KEY or sign in to Synthetic again." + +# A panel opened and closed repeatedly must not become one request per flick. +PROBE_MIN_INTERVAL_SECONDS = 15 +SCAN_REUSE_SECONDS = 20 +LIMITS_ONLY_REUSE_SECONDS = 900 + +# Past this a token count or timestamp from disk is garbage, not usage. +MAX_TOKENS = 10_000_000_000 +EARLIEST_EPOCH = 946_684_800 # 2000-01-01 +LATEST_EPOCH = 4_102_444_800 # 2100-01-01 + + +# --------------------------------------------------------------- utilities -class SyntheticError(Exception): - def __init__(self, message: str, retry_advised: bool = False): - super().__init__(message) - self.retry_advised = retry_advised +def obj(value): + return value if isinstance(value, dict) else {} + + +# Token counts and timestamps on disk are untrusted: a boolean, NaN, infinite, +# negative, huge or non-numeric value reads as nothing. +def number(value): + if isinstance(value, bool) or not isinstance(value, (int, float)): + return 0 + if not math.isfinite(value) or value < 0 or value > MAX_TOKENS: + return 0 + return int(value) + + +def finite(value): + if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value): + return None + return float(value) -def utc_now() -> str: - return datetime.now(timezone.utc).isoformat() +def epoch_seconds(value): + if finite(value) is None: + return None + # Session timestamps are seconds or milliseconds. + seconds = value / 1000 if value > 100_000_000_000 else value + return seconds if EARLIEST_EPOCH <= seconds < LATEST_EPOCH else None -def finite_number(value: Any) -> float | None: +def local_day(value): + seconds = epoch_seconds(value) + if seconds is not None: + return datetime.fromtimestamp(seconds).strftime("%Y-%m-%d") + if not isinstance(value, str) or not value.strip(): + return None try: - parsed = float(value) - return parsed if math.isfinite(parsed) else None - except (TypeError, ValueError): + parsed = datetime.fromisoformat(value.strip().replace("Z", "+00:00")) + except ValueError: + return None + if epoch_seconds(parsed.timestamp()) is None: return None + return parsed.astimezone().strftime("%Y-%m-%d") if parsed.tzinfo else parsed.strftime("%Y-%m-%d") + + +def model_name(raw): + # Synthetic serves Hugging Face models as "hf:org/model"; the model is the + # last part. + value = str(raw or "").strip().removeprefix("hf:").rstrip("/").split("/")[-1] + return value or AGENT_ID + + +def runtime_env(): + home = str(Path.home()) + path_parts = [ + os.environ.get("PATH", ""), + f"{home}/.local/bin", + f"{home}/.npm-global/bin", + f"{home}/.local/share/mise/shims", + ] + env = os.environ.copy() + env["PATH"] = os.pathsep.join(part for part in path_parts if part) + return env + + +ENV = runtime_env() + +def find_command(name): + return shutil.which(name, path=ENV.get("PATH")) -def number(value: Any, fallback: float = 0.0) -> float: - parsed = finite_number(value) - if parsed is None: - return fallback - return parsed +def data_home(): + return Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) -def ratio(used: Any, limit: Any) -> float | None: - used_value = finite_number(used) - total = finite_number(limit) - if used_value is None or total is None or total <= 0: + +def cache_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 read_json(path): + try: + return json.loads(path.read_text(encoding="utf-8")) + except Exception: return None - return max(0.0, min(1.0, used_value / total)) -def remaining_ratio(remaining: Any, maximum: Any) -> float | None: - remaining_value = finite_number(remaining) - total = finite_number(maximum) - if remaining_value is None or total is None or total <= 0: +def read_fresh_json(path, max_age_seconds): + if max_age_seconds <= 0: + 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 + except OSError: return None - return max(0.0, min(1.0, 1.0 - remaining_value / total)) + return read_json(path) if 0 <= age < max_age_seconds else None -def reset_at(value: Any) -> str: - return str(value or "").strip() +def write_json(path, payload): + # A temp name unique to this writer: 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 + + +# ------------------------------------------------------------------- stats + + +def empty_bucket(): + return {"inputTokens": 0, "outputTokens": 0, "cacheReadInputTokens": 0, "cacheCreationInputTokens": 0} + + +class Stats: + """Token and activity totals shared by every local source. Sessions and days + are sets, so a day active in pi and opencode counts once toward activeDays.""" + + def __init__(self): + 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() + self.today_tokens_by_model = {} + self.total_prompts = 0 + self.total_sessions = set() + self.active_days = set() + self.model_usage = {} + + def add(self, session, day, model, input_tokens, output_tokens, cache_read, cache_write): + total = input_tokens + output_tokens + cache_read + cache_write + if total <= 0: + return + self.total_prompts += 1 + self.active_days.add(day) + if session is not None: + self.total_sessions.add(session) + bucket = self.model_usage.setdefault(model, empty_bucket()) + bucket["inputTokens"] += input_tokens + bucket["outputTokens"] += output_tokens + bucket["cacheReadInputTokens"] += cache_read + bucket["cacheCreationInputTokens"] += cache_write + if day in self.recent: + # recentDays.messageCount holds tokens, despite the legacy name. + self.recent[day] += total + if day == self.today: + self.today_prompts += 1 + if session is not None: + self.today_sessions.add(session) + self.today_total_tokens += total + self.today_tokens_by_model[model] = self.today_tokens_by_model.get(model, 0) + total + + def add_session(self, session, today): + self.total_sessions.add(session) + if today: + self.today_sessions.add(session) + + def record(self): + 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, + } + + +# --------------------------------------------------------------- pi and omp + + +def pi_agent_dirs(): + pi_agent = os.environ.get("PI_CODING_AGENT_DIR") + dirs = [Path(pi_agent).expanduser() if pi_agent else Path.home() / ".pi" / "agent", Path.home() / ".omp" / "agent"] + # `omp --profile=` (and pi's equivalent) moves the agent tree under + # /profiles//. + for base in (Path.home() / ".pi", Path.home() / ".omp"): + try: + dirs.extend(sorted(child / "agent" for child in (base / "profiles").iterdir() if child.is_dir())) + except OSError: + pass + return dirs -def format_quantity(value: Any) -> str: - parsed = number(value) - if parsed.is_integer(): - return f"{int(parsed):,}" - return f"{parsed:,.2f}".rstrip("0").rstrip(".") +def pi_pattern(): + return r'"provider"\s*:\s*"synthetic"' -def empty_stats() -> dict[str, Any]: - return { - "todayPrompts": 0, - "todaySessions": 0, - "todayTotalTokens": 0, - "todayTokensByModel": {}, - "recentDays": [], - "totalPrompts": 0, - "totalSessions": 0, - "activeDays": 0, - "activeDates": [], - "modelUsage": {}, - } +def accept_pi(message): + return str(message.get("provider") or "") == PROVIDER -def base_record(**overrides: Any) -> dict[str, Any]: - record: dict[str, Any] = { - "schemaVersion": 1, - "id": AGENT_ID, - "name": AGENT_NAME, - "updatedAt": utc_now(), - "ready": False, - "hasLocalStats": False, - "hasPromptStats": False, - "scope": "account", - "tierLabel": "Subscription", - "usageStatusText": "", - "authHelpText": "", - "limits": [], - } - record.update(empty_stats()) - record.update(overrides) - return record +def pi_session_roots(): + roots = [] + for agent_dir in pi_agent_dirs(): + root = agent_dir / "sessions" + if root.is_dir() and root.resolve() not in [r.resolve() for r in roots]: + roots.append(root) + return roots -def stored_key(path: Path) -> str: +def pi_parent_session(path): + # A fork names the session it was cut from in its header line. try: - payload = json.loads(path.read_text(encoding="utf-8")) - except (OSError, json.JSONDecodeError): - return "" - entry = payload.get("synthetic") if isinstance(payload, dict) else None - if not isinstance(entry, dict): + with open(path, encoding="utf-8", errors="replace") as handle: + header = json.loads(handle.readline()) + except (OSError, ValueError): + return None + parent = header.get("parentSession") if isinstance(header, dict) else None + return os.path.abspath(os.path.expanduser(parent)) if isinstance(parent, str) and parent else None + + +def scan_pi_sessions(stats): + """Add the pi/omp assistant messages this collector counts. Returns whether + every root was read.""" + roots = pi_session_roots() + if not roots: + return True + rg = find_command("rg") or "rg" + try: + proc = subprocess.Popen( + # --no-ignore: these are data files, not a source tree. A $HOME that is + # a git checkout would otherwise hide every session behind its ignores. + [rg, "--json", "--no-ignore", "-e", pi_pattern(), *map(str, roots)], + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + text=True, + errors="replace", + env=ENV, + ) + except OSError: + return False + + counted = set() + owned = {} + today_owned = {} + assert proc.stdout is not None + for raw in proc.stdout: + try: + event = json.loads(raw) + if event.get("type") != "match": + continue + path = os.path.abspath(event["data"]["path"]["text"]) + line_number = event["data"].get("line_number") + entry = json.loads(event["data"]["lines"]["text"]) + except Exception: + continue + if not isinstance(entry, dict) or entry.get("type") != "message": + continue + message = obj(entry.get("message")) + if message.get("role") != "assistant" or not accept_pi(message): + continue + usage = obj(message.get("usage")) + 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 not (input_tokens or output_tokens or cache_read or cache_write): + input_tokens = number(usage.get("totalTokens")) + timestamp = entry.get("timestamp") + day = local_day(timestamp) or local_day(message.get("timestamp")) + if day is None or not (input_tokens or output_tokens or cache_read or cache_write): + continue + # A fork repeats its parent's messages with their ids and timestamps, so + # the pair says two lines are one message. An id alone is only unique + # within its file, and a line without one always counts. + message_id = str(entry.get("id") or "") + if message_id and timestamp is not None: + key = ("entry", message_id, str(timestamp)) + elif message_id: + key = ("path", path, message_id) + else: + key = ("line", path, line_number) + owned.setdefault(path, set()).add(key) + if day == stats.today: + today_owned.setdefault(path, set()).add(key) + if key in counted: + continue + counted.add(key) + stats.add(None, day, model_name(message.get("model")), input_tokens, output_tokens, cache_read, cache_write) + + try: + code = proc.wait(timeout=5) + except subprocess.TimeoutExpired: + proc.kill() + return False + + # A message belongs to the oldest session holding it, so a fork is a + # session of its own only for what it added, whichever file rg read first. + parents = {} + for path, keys in owned.items(): + inherited = set() + visited = {path} + parent = parents.setdefault(path, pi_parent_session(path)) + while parent and parent not in visited: + visited.add(parent) + inherited |= owned.get(parent, set()) + parent = parents.setdefault(parent, pi_parent_session(parent)) + if keys - inherited: + stats.add_session(path, bool(today_owned.get(path, set()) - inherited)) + # rg exits 1 when nothing matched and 2 when a file couldn't be read. + return code in (0, 1) + + +# ---------------------------------------------------------------- opencode + + +def opencode_db(): + # OPENCODE_DB points opencode at another database, as the CLI itself reads + # it; a relative name lives beside the default one. Only a build whose + # channel was literally "v2" wrote opencode-v2.db. + root = data_home() / "opencode" + configured = os.environ.get("OPENCODE_DB", "").strip() + if configured: + path = Path(configured).expanduser() + return path if path.is_absolute() else root / path + normal = root / "opencode.db" + legacy = root / "opencode-v2.db" + return legacy if not normal.is_file() and legacy.is_file() else normal + + +def table_columns(conn, table): + return {row[1] for row in conn.execute(f"PRAGMA table_info({table})")} + + +def opencode_row(stats, seen, message_id, session_id, entry, row_type, created): + # The original `message` store keeps the provider and model at the top + # level; the V2 `session_message` store nests them under `model` and keeps + # the role in its `type` column, with a compaction's tokens spent like any + # reply's. + if not isinstance(entry, dict) or (entry.get("role") or row_type) not in ("assistant", "compaction"): + return + model = entry.get("model") + provider = entry.get("providerID") or (model.get("providerID") if isinstance(model, dict) else "") + # Exact match: opencode provider ids are free-form. + if str(provider or "") != PROVIDER: + return + tokens = obj(entry.get("tokens")) + cache = obj(tokens.get("cache")) + 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): + return + # A migrated message can sit in both stores under the same id. + if message_id in seen: + return + seen.add(message_id) + day = local_day(obj(entry.get("time")).get("created")) or local_day(created) + if day is None: + return + model_id = entry.get("modelID") or (model.get("id") or model.get("modelID") if isinstance(model, dict) else model) + stats.add("opencode:" + str(session_id), day, model_name(model_id), input_tokens, output_tokens, cache_read, cache_write) + + +def scan_opencode(stats): + """Add opencode messages on the synthetic provider, from both of its + stores. Returns whether the database was read to the end.""" + db = opencode_db() + if not db.is_file(): + return True + try: + # Read-only, and not immutable: opencode writes through a WAL, and an + # immutable open would miss what hasn't been checkpointed yet. + conn = sqlite3.connect(db.resolve().as_uri() + "?mode=ro", uri=True, timeout=2) + except sqlite3.Error: + return False + gate = '%"providerID"%:%"' + PROVIDER + '"%' + seen = set() + try: + conn.execute("PRAGMA query_only = ON") + tables = {name for (name,) in conn.execute("SELECT name FROM sqlite_master WHERE type = 'table'")} + if "session_message" in tables and {"id", "session_id", "type", "data"} <= table_columns(conn, "session_message"): + columns = table_columns(conn, "session_message") + created = "m.time_created" if "time_created" in columns else "NULL" + source, fork_filter = "session_message m", "" + if "time_created" in columns and "session_v2" in tables and {"id", "fork_session_id", "time_created"} <= table_columns(conn, "session_v2"): + # A fork copies its source session's messages under new ids; only + # what came after the fork was spent in it. + source += " LEFT JOIN session_v2 s ON s.id = m.session_id" + fork_filter = " AND (s.fork_session_id IS NULL OR m.time_created >= s.time_created)" + for message_id, session_id, row_type, time_created, raw in conn.execute( + f"SELECT m.id, m.session_id, m.type, {created}, m.data FROM {source}" + " WHERE m.type IN ('assistant', 'compaction') AND m.data LIKE ?" + fork_filter, (gate,) + ): + try: + opencode_row(stats, seen, str(message_id), session_id, json.loads(raw), str(row_type or ""), time_created) + except (TypeError, ValueError): + continue + if "message" in tables and {"session_id", "data"} <= table_columns(conn, "message"): + message_id = "id" if "id" in table_columns(conn, "message") else "rowid" + for row_id, session_id, raw in conn.execute(f"SELECT {message_id}, session_id, data FROM message WHERE data LIKE ?", (gate,)): + try: + opencode_row(stats, seen, str(row_id), session_id, json.loads(raw), "", None) + except (TypeError, ValueError): + continue + except sqlite3.Error: + # Transient lock, schema migration, corruption: the numbers stop here, + # incomplete. + return False + finally: + conn.close() + return True + + +# -------------------------------------------------------------- scan cache + + +def scan_cache_paths(): + digest = hashlib.sha1("\n".join([str(Path.home()), str(opencode_db()), os.environ.get("PI_CODING_AGENT_DIR", "")]).encode("utf-8")).hexdigest()[:16] + root = cache_root() + return root / f"synthetic-scan-{digest}.json", root / f"synthetic-scan-{digest}.lock" + + +def run_local_scans(): + stats = Stats() + complete = scan_pi_sessions(stats) + complete = scan_opencode(stats) and complete + return stats.record(), complete + + +def read_cached_stats(cache_file, max_age): + cached = read_fresh_json(cache_file, max_age) + if not isinstance(cached, dict) or cached.get("schemaVersion") != 1: + return None + # today* fields only mean "today" on the day they were scanned. + if cached.get("scanDate") != datetime.now().strftime("%Y-%m-%d"): + return None + stats = cached.get("stats") + return stats if isinstance(stats, dict) and set(Stats().record()) <= set(stats) else None + + +def cached_local_stats(max_age): + """Local stats, with the cache as a pure optimization: any cache-layer + failure degrades to a direct scan and a warning on stderr.""" + try: + cache_file, lock_file = scan_cache_paths() + cached = read_cached_stats(cache_file, max_age) + if cached is not None: + return cached + with lock_file.open("w") as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + cached = read_cached_stats(cache_file, max_age) + if cached is not None: + return cached + stats, complete = run_local_scans() + # A scan that couldn't read everything serves this run, but caching it + # would hide the missing usage from every reader until it expires. + if complete: + write_json(cache_file, {"schemaVersion": 1, "scanDate": datetime.now().strftime("%Y-%m-%d"), "stats": stats}) + return stats + except Exception as exc: + print(f"omarchy-agent-usage-synthetic: cache unavailable ({exc}); scanning directly", file=sys.stderr) + return run_local_scans()[0] + + +# ------------------------------------------------------------- credentials + + +PI_ENV_REFERENCE = re.compile(r"\$(?:\{([A-Za-z_][A-Za-z0-9_]*)\}|([A-Za-z_][A-Za-z0-9_]*))") + + +def stored_key(path): + # pi, omp and opencode keep {"synthetic": {"key": ...}}. pi's key may also + # be a $ENV_VAR reference, resolved here, or a !command, which stays pi's. + entry = obj(obj(read_json(path)).get(PROVIDER)) + key = str(entry.get("key") or "").strip() + if not key or key.startswith("!"): return "" - return str(entry.get("key") or "").strip() + unresolved = [] + def substitute(match): + value = os.environ.get(match.group(1) or match.group(2), "") + if not value: + unresolved.append(match.group(0)) + return value -def credentials() -> str: + key = PI_ENV_REFERENCE.sub(substitute, key).strip() + return "" if unresolved else key + + +def api_key(): explicit = str(os.environ.get("SYNTHETIC_API_KEY", "")).strip() if explicit: return explicit - # Pi and OpenCode use the same provider id and {"key": ...} auth shape. - pi_key = stored_key(Path.home() / ".pi" / "agent" / "auth.json") - if pi_key: - return pi_key - data_home = Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) - return stored_key(data_home / "opencode" / "auth.json") + paths = [agent_dir / "auth.json" for agent_dir in pi_agent_dirs()[:2]] + paths.append(data_home() / "opencode" / "auth.json") + for path in paths: + key = stored_key(path) + if key: + return key + return "" -def quota_limit(label: str, payload: Any) -> dict[str, Any] | None: - if not isinstance(payload, dict): - return None - used = payload.get("requests") - maximum = payload.get("limit") - percent = ratio(used, maximum) - if percent is None: +# ------------------------------------------------------------------ quotas + + +def used_fraction(used, maximum): + used, maximum = finite(used), finite(maximum) + if used is None or maximum is None or maximum <= 0: return None - return { - "label": label, - "title": label, - "percent": percent, - "detail": f"{format_quantity(used)} / {format_quantity(maximum)} used", - "resetsAt": reset_at(payload.get("renewsAt")), - } + return max(0.0, min(1.0, used / maximum)) -def parse_limits(payload: dict[str, Any]) -> tuple[list[dict[str, Any]], str, str]: - limits: list[dict[str, Any]] = [] - - rolling = payload.get("rollingFiveHourLimit") - weekly = payload.get("weeklyTokenLimit") - has_modern_limit = False - limited = False - if isinstance(rolling, dict): - percent = remaining_ratio(rolling.get("remaining"), rolling.get("max")) - if percent is not None: - has_modern_limit = True - maximum = number(rolling.get("max")) - used = max(0.0, min(maximum, maximum - number(rolling.get("remaining")))) - limits.append({ - "label": "Requests used / 5h", - "title": "Requests used / 5h", - "percent": percent, - "detail": f"{format_quantity(used)} / {format_quantity(maximum)} used", - "resetsAt": reset_at(rolling.get("nextTickAt")), - "resetLabel": "Next tick in", - }) - limited = rolling.get("limited") is True - - weekly_remaining = None - weekly_percent = finite_number(weekly.get("percentRemaining")) if isinstance(weekly, dict) else None - if isinstance(weekly, dict) and weekly_percent is not None: - has_modern_limit = True - remaining_percent = max(0.0, min(100.0, weekly_percent)) - weekly_remaining = remaining_percent - remaining_credits = str(weekly.get("remainingCredits") or "").strip() - maximum_credits = str(weekly.get("maxCredits") or "").strip() - detail = ( - f"{remaining_credits} / {maximum_credits} left" - if remaining_credits and maximum_credits - else f"{format_quantity(remaining_percent)}% left" - ) - limits.append({ - "label": "Credits used / week", - "title": "Credits used / week", - "percent": 1.0 - remaining_percent / 100.0, - "detail": detail, - "resetsAt": reset_at(weekly.get("nextRegenAt")), - "resetLabel": "Next regen in", - }) - - # The endpoint documented only this bucket originally. Modern accounts - # still receive it at zero alongside the credit/request system, so use it - # only when neither modern allowance parsed successfully. - if not has_modern_limit: - subscription = quota_limit("Subscription requests", payload.get("subscription")) - if subscription: - limits.insert(0, subscription) - - search = payload.get("search") - if isinstance(search, dict): - hourly = quota_limit("Search used / hour", search.get("hourly")) - if hourly: - limits.append(hourly) - - free_tools = quota_limit("Free tool calls", payload.get("freeToolCalls")) - if free_tools: - limits.append(free_tools) - - if limited: - return limits, "Rolling limit reached", "Synthetic reports that the rolling five-hour limit is currently active." - if weekly_remaining is not None and weekly_remaining <= 0: - return limits, "Weekly credits exhausted", "Synthetic reports no weekly credits remaining." - return limits, "", "" - - -class SyntheticClient: - def __init__(self, api_key: str, base_url: str = API_BASE_URL): - self.api_key = api_key - self.base_url = base_url.rstrip("/") - - def quotas(self) -> dict[str, Any]: - request = urllib.request.Request( - self.base_url + "/v2/quotas", - headers={ - "Authorization": "Bearer " + self.api_key, - "Accept": "application/json", - "User-Agent": "omarchy-agents-synthetic/1.0", - }, - ) - try: - with urllib.request.urlopen(request, timeout=15) as response: - payload = json.load(response) - except urllib.error.HTTPError as error: - if error.code == 401: - raise SyntheticError("Synthetic rejected the API key") from error - if error.code == 403: - raise SyntheticError("The Synthetic API key cannot read quota data") from error - if error.code == 429: - raise SyntheticError("Synthetic is rate limiting quota checks") from error - raise SyntheticError(f"Synthetic quota API returned HTTP {error.code}") from error - except urllib.error.URLError as error: - raise SyntheticError("Could not reach the Synthetic quota API", retry_advised=True) from error - except TimeoutError as error: - raise SyntheticError("Synthetic quota request timed out", retry_advised=True) from error - except json.JSONDecodeError as error: - raise SyntheticError("Synthetic returned an invalid quota response") from error - if not isinstance(payload, dict): - raise SyntheticError("Synthetic returned an invalid quota response") - return payload - - -def scan(base_url: str) -> dict[str, Any]: - api_key = credentials() - if not api_key: - return base_record(usageStatusText="Waiting for auth", authHelpText=AUTH_HELP) - - payload = SyntheticClient(api_key, base_url).quotas() - limits, status, help_text = parse_limits(payload) +def reset_time(value): + raw = str(value or "").strip() + if not raw: + return "" + try: + datetime.fromisoformat(raw.replace("Z", "+00:00")) + except ValueError: + return "" + return raw + + +def request_window(label, payload, title=None): + payload = obj(payload) + percent = used_fraction(payload.get("requests"), payload.get("limit")) + if percent is None: + return None + window = {"label": label, "percent": percent, "resetsAt": reset_time(payload.get("renewsAt"))} + if title: + window["title"] = title + return window + + +def parse_quotas(payload): + limits = [] + # The rolling limit and the weekly credits regenerate a little at a time. + # nextTickAt and nextRegenAt say when the next share comes back, not when + # the window starts over, so they are no reset time: reading them as one + # would zero a window that has only regained a sliver. + rolling = obj(payload.get("rollingFiveHourLimit")) + remaining, maximum = finite(rolling.get("remaining")), finite(rolling.get("max")) + if remaining is not None and maximum is not None and maximum > 0: + percent = 1.0 if rolling.get("limited") is True else max(0.0, min(1.0, 1.0 - remaining / maximum)) + limits.append({"label": "Rolling (5h)", "percent": percent, "resetsAt": ""}) + weekly = finite(obj(payload.get("weeklyTokenLimit")).get("percentRemaining")) + if weekly is not None: + limits.append({"label": "Weekly credits", "percent": max(0.0, min(1.0, 1.0 - weekly / 100.0)), "resetsAt": ""}) + + # The documented bucket. Accounts on the credit system still receive it, at + # zero, so it stands in only when neither of those parsed. if not limits: - return base_record( - usageStatusText="Synthetic quota unavailable", - authHelpText="The quota endpoint returned no supported limits.", - ) - tier_label = "Subscription" - weekly = payload.get("weeklyTokenLimit") - if isinstance(weekly, dict) and str(weekly.get("remainingCredits") or "").strip(): - tier_label = str(weekly["remainingCredits"]).strip() + " weekly credits left" - return base_record( - ready=True, - tierLabel=tier_label, - usageStatusText=status, - authHelpText=help_text, - limits=limits, + window = request_window("Subscription requests", payload.get("subscription")) + if window: + limits.append(window) + for window in ( + request_window("Search (hourly)", obj(payload.get("search")).get("hourly"), "Search"), + request_window("Free tool calls", payload.get("freeToolCalls")), + ): + if window: + limits.append(window) + return limits + + +def probe_quotas(key, base_url): + request = urllib.request.Request( + base_url.rstrip("/") + "/v2/quotas", + headers={ + "Authorization": "Bearer " + key, + "Accept": "application/json", + "User-Agent": "omarchy-agent-usage-synthetic/1.0", + }, ) + 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, "refused": True, "helpText": REFUSED_HELP} + return {"ok": False, "helpText": f"Synthetic's quota endpoint returned status {error.code}. Local Synthetic stats are still shown."} + except (urllib.error.URLError, TimeoutError, OSError): + # A transport failure reached no server at all: no route, no DNS. + return {"ok": False, "transport": True, "helpText": "Couldn't reach Synthetic's quota endpoint. Retrying shortly. Local Synthetic stats are still shown."} + except ValueError: + return {"ok": False, "helpText": "Synthetic returned an unreadable quota answer. Local Synthetic stats are still shown."} + limits = parse_quotas(payload) if isinstance(payload, dict) else [] + if not limits: + return {"ok": False, "helpText": "Synthetic's quota endpoint returned no quotas. Local Synthetic stats are still shown."} + return {"ok": True, "limits": limits} -def main() -> int: - parser = argparse.ArgumentParser(description="Print Synthetic quota status as JSON") - # The shared updater passes these flags to every collector. Synthetic's - # quota request is inexpensive and does not count against subscription limits. - parser.add_argument("--force", action="store_true") - parser.add_argument("--limits-only", action="store_true") - parser.parse_args() - +def limit_window_open(entry, now): + raw = str(entry.get("resetsAt") or "") + if raw == "": + return True try: - record = scan(os.environ.get("SYNTHETIC_API_BASE_URL", API_BASE_URL)) - except SyntheticError as error: - record = base_record( - usageStatusText="Synthetic unavailable", - authHelpText=str(error), - retryAdvised=error.retry_advised, - ) - except Exception as error: - record = base_record( - usageStatusText="Synthetic unavailable", - authHelpText="Synthetic quota scan failed.", - ) - print(f"omarchy-agent-usage-synthetic: {type(error).__name__}", file=sys.stderr) - + 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 + + +# A kept window whose reset has passed is untouched, not unknown: it comes +# back at zero with no reset time until the next check. +def usable_cached_limits(cached): + entries = obj(cached).get("limits") + if not isinstance(entries, list): + return [] + now = datetime.now(timezone.utc) + return [ + entry if limit_window_open(entry, now) else dict(entry, percent=0.0, resetsAt="") + for entry in entries if isinstance(entry, dict) + ] + + +# One cache per key, so another key never shows this one's quotas. +def limits_cache_path(key): + return cache_root() / f"synthetic-limits-{hashlib.sha256(key.encode('utf-8')).hexdigest()[:16]}.json" + + +def collect_limits(key, base_url, force): + result = {"limits": [], "usageStatusText": "", "authHelpText": "", "stale": False, "fetchedAtMs": 0, "retryAdvised": False} + if not key: + return result + + cache = limits_cache_path(key) + cached = read_json(cache) + fallback = usable_cached_limits(cached) + fetched_at = obj(cached).get("fetchedAtMs") + fetched_at = round(fetched_at) if epoch_seconds(fetched_at) else 0 + + # --force is a person asking for fresh numbers, so it skips the reuse window. + if fallback and not force and 0 <= time.time() - fetched_at / 1000 < PROBE_MIN_INTERVAL_SECONDS: + result.update(limits=fallback, fetchedAtMs=fetched_at) + return result + + probe = probe_quotas(key, base_url) + if probe["ok"]: + fetched_at = round(time.time() * 1000) + try: + write_json(cache, {"fetchedAtMs": fetched_at, "limits": probe["limits"]}) + except OSError as exc: + print(f"omarchy-agent-usage-synthetic: could not write the quota cache ({exc})", file=sys.stderr) + result.update(limits=probe["limits"], fetchedAtMs=fetched_at) + return result + + # Last good numbers, marked stale, beat an empty section while the check + # fails. + result.update(limits=fallback, stale=bool(fallback), fetchedAtMs=fetched_at if fallback else 0, retryAdvised=bool(probe.get("transport"))) + if probe.get("refused"): + result.update(usageStatusText="Sign-in expired", authHelpText=probe["helpText"]) + elif not fallback: + result.update(usageStatusText="Synthetic limits unavailable", authHelpText=probe["helpText"]) + return result + + +# ------------------------------------------------------------------ record + + +def main(): + parser = argparse.ArgumentParser(description="Print the Synthetic usage record as JSON") + parser.add_argument("--force", action="store_true", help="rescan local sessions and re-check quotas, ignoring caches") + parser.add_argument("--limits-only", action="store_true", help="reuse a recent local scan; only the quotas must be fresh") + args = parser.parse_args() + + max_age = 0 if args.force else (LIMITS_ONLY_REUSE_SECONDS if args.limits_only else SCAN_REUSE_SECONDS) + stats = cached_local_stats(max_age) + key = api_key() + limits = collect_limits(key, os.environ.get("SYNTHETIC_API_BASE_URL") or API_BASE_URL, args.force) + in_use = number(stats.get("totalPrompts")) > 0 + + record = { + "schemaVersion": 1, + "id": AGENT_ID, + "name": AGENT_NAME, + "updatedAt": datetime.now(timezone.utc).isoformat(), + "ready": in_use or bool(key), + "hasLocalStats": True, + # The quotas answer names no plan. + "tierLabel": "", + "usageStatusText": limits["usageStatusText"], + "authHelpText": limits["authHelpText"], + "limits": limits["limits"], + "limitsStale": limits["stale"], + "limitsFetchedAt": limits["fetchedAtMs"], + } + if in_use and not key: + record.update(usageStatusText="Waiting for auth", authHelpText=AUTH_HELP) + if limits["retryAdvised"]: + record["retryAdvised"] = True + record.update(stats) print(json.dumps(record, separators=(",", ":"), sort_keys=True)) return 0 diff --git a/test/shell.d/agent-usage-synthetic-test.sh b/test/shell.d/agent-usage-synthetic-test.sh index eeecb267a10..a6d8b7bc68b 100755 --- a/test/shell.d/agent-usage-synthetic-test.sh +++ b/test/shell.d/agent-usage-synthetic-test.sh @@ -1,154 +1,237 @@ #!/bin/bash -set -euo pipefail - -source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" +source "$(dirname "$0")/base-test.sh" require_command jq require_command python3 +require_command rg TEST_HOME=$(mktemp -d) -trap 'rm -rf "$TEST_HOME"' EXIT +SERVER_PID="" +cleanup() { + [[ -n $SERVER_PID ]] && kill "$SERVER_PID" 2>/dev/null + rm -rf "$TEST_HOME" +} +trap cleanup EXIT + +export HOME="$TEST_HOME" +export XDG_CACHE_HOME="$TEST_HOME/cache" +export XDG_DATA_HOME="$TEST_HOME/data" +export XDG_CONFIG_HOME="$TEST_HOME/config" +unset SYNTHETIC_API_KEY PI_CODING_AGENT_DIR OPENCODE_DB collector="$ROOT/bin/omarchy-agent-usage-synthetic" +STATE="$TEST_HOME/server" +mkdir -p "$STATE" +LIMITS_CACHE="$XDG_CACHE_HOME/omarchy/agent-usage" + +# A stand-in for api.synthetic.new: answers /v2/quotas with $STATE/body and +# the status in $STATE/status, and notes each request and its key. +python3 - "$STATE" <<'PY' & +import http.server, pathlib, sys +state = pathlib.Path(sys.argv[1]) + +class Handler(http.server.BaseHTTPRequestHandler): + def do_GET(self): + count = state / "count" + count.write_text(str(int(count.read_text() if count.exists() else "0") + 1)) + (state / "auth").write_text(self.headers.get("Authorization", "")) + status = int((state / "status").read_text()) if (state / "status").exists() else 200 + body = (state / "body").read_bytes() if self.path == "/v2/quotas" else b"{}" + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.end_headers() + self.wfile.write(body) + + def log_message(self, *args): + pass + +server = http.server.HTTPServer(("127.0.0.1", 0), Handler) +(state / "port").write_text(str(server.server_address[1])) +server.serve_forever() +PY +SERVER_PID=$! +for _ in $(seq 50); do + [[ -s $STATE/port ]] && break + sleep 0.1 +done +[[ -s $STATE/port ]] || fail "Synthetic test server starts" +export SYNTHETIC_API_BASE_URL="http://127.0.0.1:$(cat "$STATE/port")" + +requests() { + cat "$STATE/count" 2>/dev/null || echo 0 +} -no_auth=$(HOME="$TEST_HOME" XDG_DATA_HOME="$TEST_HOME/.local/share" SYNTHETIC_API_KEY="" "$collector") -[[ $(jq -r '.id + ":" + (.ready | tostring) + ":" + (.hasLocalStats | tostring)' <<<"$no_auth") == "synthetic:false:false" ]] || - fail "Synthetic collector prints a valid record without credentials" "$no_auth" -[[ $(jq -r '.authHelpText' <<<"$no_auth") == *"SYNTHETIC_API_KEY"* ]] || - fail "Synthetic collector explains how to authenticate" "$no_auth" -pass "Synthetic collector handles missing credentials" - -result=$(python3 - "$collector" "$TEST_HOME" <<'PY' -import importlib.machinery -import importlib.util -import json -import os -import sys -import urllib.error -from pathlib import Path +serve() { + printf '%s' "$1" >"$STATE/status" + printf "%s" "${2-}" >"$STATE/body" +} -collector_path = sys.argv[1] -home = Path(sys.argv[2]) -os.environ["HOME"] = str(home) -os.environ["XDG_DATA_HOME"] = str(home / ".local" / "share") +modern='{ + "subscription": {"requests": 0, "limit": 1000, "renewsAt": "2099-08-17T00:00:00Z"}, + "rollingFiveHourLimit": {"remaining": 750, "max": 1000, "limited": false, "nextTickAt": "2099-08-16T02:00:00Z"}, + "weeklyTokenLimit": {"percentRemaining": 64, "remainingCredits": "$31.05", "maxCredits": "$48.00", "nextRegenAt": "2099-08-16T04:00:00Z"}, + "search": {"hourly": {"requests": 25, "limit": 250, "renewsAt": "2099-08-16T03:00:00Z"}}, + "freeToolCalls": {"requests": "many", "limit": 100} +}' +serve 200 "$modern" + +# ------------------------------------------------------------ unused machine + +result=$("$collector") +[[ $(jq -c '[.ready, .limits, .usageStatusText, .authHelpText, .tierLabel, .todayTotalTokens, (.recentDays | length), has("retryAdvised")]' <<<"$result") == '[false,[],"","","",0,7,false]' ]] || + fail "Synthetic collector prints an empty, not-ready record on an unused machine" "$result" +[[ $(requests) == "0" ]] || + fail "Synthetic collector asks nothing without a key" "$(requests)" +pass "Synthetic collector prints an empty, not-ready record on an unused machine" + +# ---------------------------------------------------------------- credentials + +key_used() { + rm -f "$LIMITS_CACHE"/synthetic-limits-*.json + "$@" "$collector" --force >/dev/null + cat "$STATE/auth" +} -loader = importlib.machinery.SourceFileLoader("synthetic_collector", collector_path) +mkdir -p "$TEST_HOME/.pi/agent" "$TEST_HOME/.omp/agent" "$XDG_DATA_HOME/opencode" "$TEST_HOME/moved" +echo '{"synthetic":{"type":"api_key","key":"opencode-key"}}' >"$XDG_DATA_HOME/opencode/auth.json" +[[ $(key_used env) == "Bearer opencode-key" ]] || fail "Synthetic collector reads opencode's key" "$(cat "$STATE/auth")" +echo '{"synthetic":{"type":"api_key","key":"omp-key"}}' >"$TEST_HOME/.omp/agent/auth.json" +[[ $(key_used env) == "Bearer omp-key" ]] || fail "Synthetic collector prefers omp's key to opencode's" "$(cat "$STATE/auth")" +echo '{"synthetic":{"type":"api_key","key":"$PI_SYNTHETIC"}}' >"$TEST_HOME/.pi/agent/auth.json" +[[ $(key_used env PI_SYNTHETIC=pi-key) == "Bearer pi-key" ]] || fail "Synthetic collector resolves pi's \$ENV key" "$(cat "$STATE/auth")" +echo '{"synthetic":{"key":"moved-key"}}' >"$TEST_HOME/moved/auth.json" +[[ $(key_used env PI_CODING_AGENT_DIR="$TEST_HOME/moved") == "Bearer moved-key" ]] || fail "Synthetic collector follows PI_CODING_AGENT_DIR" "$(cat "$STATE/auth")" +[[ $(key_used env SYNTHETIC_API_KEY=env-key) == "Bearer env-key" ]] || fail "Synthetic collector prefers SYNTHETIC_API_KEY" "$(cat "$STATE/auth")" +rm -f "$TEST_HOME/.pi/agent/auth.json" "$TEST_HOME/.omp/agent/auth.json" "$XDG_DATA_HOME/opencode/auth.json" +pass "Synthetic collector reads its key from the environment, pi, omp, then opencode" + +export SYNTHETIC_API_KEY=secret + +# --------------------------------------------------------------- fresh quotas + +rm -f "$LIMITS_CACHE"/synthetic-limits-*.json +result=$("$collector") +[[ $(jq -c '[.limits[] | [.label, (.percent * 1000 | round / 1000), .resetsAt, (.title // "")]]' <<<"$result") == '[["Rolling (5h)",0.25,"",""],["Weekly credits",0.36,"",""],["Search (hourly)",0.1,"2099-08-16T03:00:00Z","Search"]]' ]] || + fail "Synthetic collector maps the quota windows" "$result" +[[ $(jq -c '[.ready, .tierLabel, .limitsStale, .usageStatusText, .authHelpText, ([.limits[] | keys[]] | unique)]' <<<"$result") == '[true,"",false,"","",["label","percent","resetsAt","title"]]' ]] || + fail "Synthetic collector reports a healthy record with only contract fields" "$result" +pass "Synthetic collector reads fresh quotas, leaving regeneration ticks out of reset times" + +limits() { + SYNTHETIC_COLLECTOR="$collector" PAYLOAD="$1" python3 - <<'PY' +import importlib.machinery, importlib.util, json, os +loader = importlib.machinery.SourceFileLoader("collector", os.environ["SYNTHETIC_COLLECTOR"]) spec = importlib.util.spec_from_loader(loader.name, loader) -scanner = importlib.util.module_from_spec(spec) -loader.exec_module(scanner) - -pi_auth = home / ".pi" / "agent" / "auth.json" -opencode_auth = home / ".local" / "share" / "opencode" / "auth.json" -pi_auth.parent.mkdir(parents=True, exist_ok=True) -opencode_auth.parent.mkdir(parents=True, exist_ok=True) -pi_auth.write_text(json.dumps({"synthetic": {"key": "pi-key"}})) -opencode_auth.write_text(json.dumps({"synthetic": {"key": "opencode-key"}})) - -os.environ["SYNTHETIC_API_KEY"] = "env-key" -env_key = scanner.credentials() -del os.environ["SYNTHETIC_API_KEY"] -pi_key = scanner.credentials() -pi_auth.unlink() -opencode_key = scanner.credentials() - -modern_payload = { - "subscription": {"requests": 0, "limit": 1000, "renewsAt": "2026-08-17T00:00:00Z"}, - "rollingFiveHourLimit": { - "remaining": 998.3, - "max": 1000, - "limited": False, - "nextTickAt": "2026-08-16T02:00:00Z", - }, - "weeklyTokenLimit": { - "percentRemaining": 64.69978866666666, - "remainingCredits": "$31.05", - "maxCredits": "$48.00", - "nextRegenAt": "2026-08-16T04:00:00Z", - }, - "search": { - "hourly": {"requests": 4, "limit": 250, "renewsAt": "2026-08-16T03:00:00Z"}, - }, +collector = importlib.util.module_from_spec(spec) +loader.exec_module(collector) +print(json.dumps(collector.parse_quotas(json.loads(os.environ["PAYLOAD"])))) +PY } -modern, status, help_text = scanner.parse_limits(modern_payload) -legacy, _, _ = scanner.parse_limits({ - "subscription": {"requests": 27, "limit": 135, "renewsAt": "2026-08-17T00:00:00Z"}, -}) -invalid_modern, _, _ = scanner.parse_limits({ - "rollingFiveHourLimit": {}, - "subscription": {"requests": 10, "limit": 100, "renewsAt": "2026-08-17T00:00:00Z"}, -}) - -captured = {} -class Response: - def __enter__(self): - return self - - def __exit__(self, *args): - return False - - def read(self, *args): - return json.dumps(modern_payload).encode() - - -def fake_urlopen(request, timeout): - captured["url"] = request.full_url - captured["authorization"] = request.get_header("Authorization") - captured["timeout"] = timeout - return Response() - -real_urlopen = scanner.urllib.request.urlopen -scanner.urllib.request.urlopen = fake_urlopen -fetched = scanner.SyntheticClient("secret", "https://synthetic.test").quotas() - -scanner.urllib.request.urlopen = lambda *_args, **_kwargs: (_ for _ in ()).throw( - urllib.error.URLError("offline") -) -try: - scanner.SyntheticClient("secret").quotas() - retry_advised = None -except scanner.SyntheticError as error: - retry_advised = error.retry_advised -finally: - scanner.urllib.request.urlopen = real_urlopen - -print(json.dumps({ - "credentials": [env_key, pi_key, opencode_key], - "modern": modern, - "modernStatus": [status, help_text], - "legacy": legacy, - "invalidModern": invalid_modern, - "request": captured, - "fetchedMatches": fetched == modern_payload, - "retryAdvised": retry_advised, -})) +[[ $(jq -c '[.[] | [.label, .percent]]' <<<"$(limits '{"subscription":{"requests":27,"limit":135,"renewsAt":"2099-08-17T00:00:00Z"}}')") == '[["Subscription requests",0.2]]' ]] || + fail "Synthetic collector reads the documented subscription bucket" "$(limits '{"subscription":{"requests":27,"limit":135}}')" +[[ $(jq -c '[.[] | .percent]' <<<"$(limits '{"rollingFiveHourLimit":{"remaining":900,"max":1000,"limited":true},"weeklyTokenLimit":{"percentRemaining":true}}')") == '[1.0]' ]] || + fail "Synthetic collector fills a limited rolling window and skips an unreadable one" "$(limits '{"rollingFiveHourLimit":{"remaining":900,"max":1000,"limited":true}}')" +pass "Synthetic collector reads the documented bucket and untrusted quota fields" + +# ------------------------------------------------------------- probe reuse + +before=$(requests) +"$collector" >/dev/null +[[ $(requests) == "$before" ]] || + fail "Synthetic collector reuses a check from seconds ago" "$(requests)" +"$collector" --force >/dev/null +[[ $(requests) == "$((before + 1))" ]] || + fail "Synthetic collector --force checks again" "$(requests)" +pass "Synthetic collector reuses a recent check and --force checks again" + +# -------------------------------------------------------------- failures + +cache_file=$(ls "$LIMITS_CACHE"/synthetic-limits-*.json) +age_cache() { + local fetched=$((($(date +%s) - 3600) * 1000)) + jq -c --argjson t "$fetched" '.fetchedAtMs = $t' "$cache_file" >"$cache_file.tmp" && mv "$cache_file.tmp" "$cache_file" +} + +age_cache +serve 500 +result=$("$collector") +[[ $(jq -c '[.limitsStale, (.limits | length), .usageStatusText, has("retryAdvised"), .limitsFetchedAt > 0]' <<<"$result") == '[true,3,"",false,true]' ]] || + fail "Synthetic collector keeps the last quotas stale after a failed check" "$result" +pass "Synthetic collector keeps the last quotas stale after a failed check" + +# A kept window whose reset passed comes back at zero. +jq -c '.limits[2].resetsAt = "2000-01-01T00:00:00Z"' "$cache_file" >"$cache_file.tmp" && mv "$cache_file.tmp" "$cache_file" +result=$("$collector") +[[ $(jq -c '.limits[2] | [.label, .percent, .resetsAt]' <<<"$result") == '["Search (hourly)",0.0,""]' ]] || + fail "Synthetic collector zeroes a kept window that reset" "$result" +pass "Synthetic collector zeroes a kept window that reset" + +result=$(SYNTHETIC_API_BASE_URL="http://127.0.0.1:1" "$collector") +[[ $(jq -c '[.retryAdvised, .limitsStale, .usageStatusText]' <<<"$result") == '[true,true,""]' ]] || + fail "Synthetic collector advises a retry when the endpoint is unreachable" "$result" +pass "Synthetic collector advises a retry only when the endpoint is unreachable" + +serve 401 +result=$("$collector") +[[ $(jq -c '[.usageStatusText, .limitsStale, (.limits | length)]' <<<"$result") == '["Sign-in expired",true,3]' ]] || + fail "Synthetic collector says a refused key needs signing in again" "$result" +pass "Synthetic collector says a refused key needs signing in again" + +rm -f "$cache_file" +serve 500 +result=$("$collector") +[[ $(jq -c '[.usageStatusText, .limits, .limitsFetchedAt]' <<<"$result") == '["Synthetic limits unavailable",[],0]' ]] || + fail "Synthetic collector says the limits are unavailable with nothing kept" "$result" +[[ $(jq -r '.authHelpText' <<<"$result") == *"status 500"* ]] || + fail "Synthetic collector says what failed" "$result" +pass "Synthetic collector says the limits are unavailable with nothing kept" + +# ------------------------------------------------------------- local stats + +serve 200 "$modern" +today=$(date +%Y-%m-%d) +sessions="$TEST_HOME/.pi/agent/sessions/--work--" +mkdir -p "$sessions" "$TEST_HOME/.omp/profiles/side/agent/sessions" +line() { + printf '{"type":"message","id":"%s","timestamp":"%sT12:00:00.000Z","message":{"role":"assistant","provider":"%s","model":"hf:moonshotai/Kimi-K2","usage":{"input":%s,"output":10,"cacheRead":0,"cacheWrite":0}}}\n' "$1" "$today" "$2" "$3" +} +{ line a synthetic 100; line b synthetic 200; line c openrouter 999; line d synthetic NaN; } >"$sessions/one.jsonl" +# A fork names its parent in its header and repeats the parent's messages. +{ + printf '{"type":"session","parentSession":"%s"}\n' "$sessions/one.jsonl" + line b synthetic 200 +} >"$TEST_HOME/.omp/profiles/side/agent/sessions/fork.jsonl" +printf '*\n' >"$TEST_HOME/.gitignore" +git -C "$TEST_HOME" init -q + +{ python3 - "$XDG_DATA_HOME/opencode/opencode.db" || true; } 2>/dev/null <<'PY' +import json, sqlite3, sys, time +conn = sqlite3.connect(sys.argv[1]); now = int(time.time() * 1000) +conn.execute("PRAGMA journal_mode=WAL") +conn.execute("CREATE TABLE message (id TEXT PRIMARY KEY, session_id TEXT, time_created INTEGER, time_updated INTEGER, data TEXT)") +conn.execute("CREATE TABLE session_message (id TEXT PRIMARY KEY, session_id TEXT, type TEXT, time_created INTEGER, data TEXT)") +def tokens(i, o, r=0): return {"input": i, "output": o, "reasoning": r, "cache": {"read": 0, "write": 0}} +conn.execute("INSERT INTO message VALUES ('m1', 's1', ?, ?, ?)", (now, now, json.dumps({"role": "assistant", "providerID": "synthetic", "modelID": "hf:zai-org/GLM-4.6", "tokens": tokens(40, 5, 5), "time": {"created": now}}))) +conn.execute("INSERT INTO message VALUES ('m2', 's1', ?, ?, ?)", (now, now, json.dumps({"role": "assistant", "providerID": "anthropic", "modelID": "claude", "tokens": tokens(999, 999), "time": {"created": now}}))) +conn.execute("INSERT INTO session_message VALUES ('v1', 's2', 'compaction', ?, ?)", (now, json.dumps({"model": {"providerID": "synthetic", "id": "hf:zai-org/GLM-4.6"}, "tokens": tokens(30, 0), "time": {"created": now}}))) +conn.execute("INSERT INTO session_message VALUES ('m1', 's1', 'assistant', ?, ?)", (now, json.dumps({"model": {"providerID": "synthetic", "id": "hf:zai-org/GLM-4.6"}, "tokens": tokens(40, 5, 5), "time": {"created": now}}))) +conn.commit() +# Left open, so the rows stay in the write-ahead log and an immutable open +# would not see them. +import os +os.kill(os.getpid(), 9) PY -) - -[[ $(jq -c '.credentials' <<<"$result") == '["env-key","pi-key","opencode-key"]' ]] || - fail "Synthetic credentials follow environment, Pi, then OpenCode priority" "$result" -pass "Synthetic credentials follow the documented priority" - -# Round percentages to 4 decimals so the assertion does not depend on exact -# IEEE-754 rendering across Python/jq versions. -[[ $(jq -c '[.modern[] | {label, percent: (.percent * 10000 | round / 10000), detail, resetLabel}]' <<<"$result") == '[{"label":"Requests used / 5h","percent":0.0017,"detail":"1.7 / 1,000 used","resetLabel":"Next tick in"},{"label":"Credits used / week","percent":0.353,"detail":"$31.05 / $48.00 left","resetLabel":"Next regen in"},{"label":"Search used / hour","percent":0.016,"detail":"4 / 250 used","resetLabel":null}]' ]] || - fail "Synthetic collector maps modern quota windows" "$result" -[[ $(jq -r '.modern | map(.label) | index("Subscription requests")' <<<"$result") == "null" ]] || - fail "Synthetic collector hides the stale legacy bucket on modern accounts" "$result" -pass "Synthetic collector maps modern quotas without the stale legacy bucket" - -[[ $(jq -c '.legacy[0] | {label,percent,detail}' <<<"$result") == '{"label":"Subscription requests","percent":0.2,"detail":"27 / 135 used"}' ]] || - fail "Synthetic collector supports the documented legacy subscription quota" "$result" -[[ $(jq -r '.invalidModern[0].label' <<<"$result") == "Subscription requests" ]] || - fail "Synthetic collector falls back when modern quota objects are unusable" "$result" -pass "Synthetic collector preserves legacy quota compatibility" - -[[ $(jq -r '.request.url' <<<"$result") == "https://synthetic.test/v2/quotas" ]] || - fail "Synthetic collector calls the documented quota endpoint" "$result" -[[ $(jq -r '.request.authorization' <<<"$result") == "Bearer secret" ]] || - fail "Synthetic collector authenticates quota requests" "$result" -[[ $(jq -r '.fetchedMatches' <<<"$result") == "true" ]] || - fail "Synthetic collector decodes quota responses" "$result" -[[ $(jq -r '.retryAdvised' <<<"$result") == "true" ]] || - fail "Synthetic collector advises one retry after transport failure" "$result" -pass "Synthetic collector handles the quota transport contract" + +result=$("$collector" --force) +# pi: a 110 + b 210 (the fork's copy counts once, NaN reads as 10 output); +# opencode: m1 50 once across both stores + compaction 30. +[[ $(jq -c '[.todayTotalTokens, .todayPrompts, .modelUsage["Kimi-K2"].inputTokens, .modelUsage["GLM-4.6"].inputTokens, .todaySessions]' <<<"$result") == "[410,5,300,70,3]" ]] || + fail "Synthetic collector counts pi, omp and opencode usage on the synthetic provider once" "$result" +pass "Synthetic collector counts pi, omp and opencode usage on the synthetic provider once" + +unset SYNTHETIC_API_KEY +result=$("$collector" --force) +[[ $(jq -c '[.ready, .usageStatusText, .limits]' <<<"$result") == '[true,"Waiting for auth",[]]' ]] || + fail "Synthetic collector waits for auth when in use without a key" "$result" +pass "Synthetic collector waits for auth when in use without a key" From f973a6ec2cc9fbe7d5700845cf972b590d46722d Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:16:31 -0400 Subject: [PATCH 094/546] Count LiteLLM usage locally and keep only the key's budget The proxy's activity routes counted requests that the harness and provider collectors already count, and with an admin key every user of the proxy, so the token stats now come from pi, omp and opencode sessions on the litellm provider instead. What the proxy still answers is the key's budget from /key/info, as a balance cached per key, reused for fifteen seconds unless forced and kept stale past a failed check. Redirects to another host are still refused, and an unused machine prints ready: false. Co-authored-by: Jocelyn Knight Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-litellm | 1074 ++++++++++------- .../agent-usage-litellm-scanner-test.sh | 496 +++----- 2 files changed, 821 insertions(+), 749 deletions(-) diff --git a/bin/omarchy-agent-usage-litellm b/bin/omarchy-agent-usage-litellm index e064b2ada72..d76585834f8 100755 --- a/bin/omarchy-agent-usage-litellm +++ b/bin/omarchy-agent-usage-litellm @@ -1,537 +1,711 @@ #!/usr/bin/python3 -# omarchy:summary=Print the LiteLLM proxy usage record as JSON +# omarchy:summary=Print the LiteLLM usage record as JSON # omarchy:args=[--force] [--limits-only] # omarchy:hidden=true -"""Collect a LiteLLM proxy's spend logs into one display-ready JSON record. - -Token stats come from the proxy's daily activity endpoint for the last 30 -days. An admin key is asked for /global/activity, which is the whole -instance; a key that cannot read that route falls back to -/user/daily/activity, which is only the key's own user. A max_budget on the -key, or maxBudget in ~/.config/omarchy/agents/litellm.json, becomes the -prepaid balance the agents panel already draws. The panel only ever reads -the JSON this prints. +"""Collect LiteLLM usage into one display-ready JSON record. + +Local stats come from pi/omp sessions and opencode messages whose provider is +litellm. The proxy's own spend logs are left alone: the harness and provider +collectors already count those requests, and an admin key would count every +user of the proxy. The key's budget, when it has one, comes from the proxy's +GET /key/info (max_budget and spend) as a balance. The base URL and key come +from LITELLM_BASE_URL and LITELLM_API_KEY, or baseUrl and apiKey in +~/.config/omarchy/agents/litellm.json. The agents panel only ever reads the +JSON this prints. """ -from __future__ import annotations - import argparse +import fcntl +import hashlib import json +import math import os +import shutil +import sqlite3 +import subprocess import sys +import tempfile +import time import urllib.error import urllib.parse import urllib.request -from datetime import date, datetime, timedelta, timezone +from datetime import datetime, timedelta, timezone from pathlib import Path -from typing import Any AGENT_ID = "litellm" AGENT_NAME = "LiteLLM" -AUTH_HELP = "Set LITELLM_BASE_URL and LITELLM_API_KEY, or write ~/.config/omarchy/agents/litellm.json." -WINDOW_DAYS = 30 +PROVIDER = "litellm" +AUTH_HELP = "Set LITELLM_BASE_URL and LITELLM_API_KEY, or write ~/.config/omarchy/agents/litellm.json, to read the key's budget." +# A panel opened and closed repeatedly must not become one request per flick. +PROBE_MIN_INTERVAL_SECONDS = 15 +SCAN_REUSE_SECONDS = 20 +LIMITS_ONLY_REUSE_SECONDS = 900 -class LiteLLMError(Exception): - def __init__(self, message: str, status: int = 0) -> None: - super().__init__(message) - self.status = status +# Past this a token count or timestamp from disk is garbage, not usage. +MAX_TOKENS = 10_000_000_000 +EARLIEST_EPOCH = 946_684_800 # 2000-01-01 +LATEST_EPOCH = 4_102_444_800 # 2100-01-01 -def number(value: Any) -> int: - try: - return max(0, round(float(value or 0))) - except (TypeError, ValueError): +# --------------------------------------------------------------- utilities + + +def obj(value): + return value if isinstance(value, dict) else {} + + +# Token counts and timestamps on disk are untrusted: a boolean, NaN, infinite, +# negative, huge or non-numeric value reads as nothing. +def number(value): + if isinstance(value, bool) or not isinstance(value, (int, float)): return 0 + if not math.isfinite(value) or value < 0 or value > MAX_TOKENS: + return 0 + return int(value) + + +def finite(value): + if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value): + return None + return float(value) -def money(value: Any) -> float: +def epoch_seconds(value): + if finite(value) is None: + return None + # Session timestamps are seconds or milliseconds. + seconds = value / 1000 if value > 100_000_000_000 else value + return seconds if EARLIEST_EPOCH <= seconds < LATEST_EPOCH else None + + +def local_day(value): + seconds = epoch_seconds(value) + if seconds is not None: + return datetime.fromtimestamp(seconds).strftime("%Y-%m-%d") + if not isinstance(value, str) or not value.strip(): + return None try: - return max(0.0, float(value or 0)) - except (TypeError, ValueError): - return 0.0 + parsed = datetime.fromisoformat(value.strip().replace("Z", "+00:00")) + except ValueError: + return None + if epoch_seconds(parsed.timestamp()) is None: + return None + return parsed.astimezone().strftime("%Y-%m-%d") if parsed.tzinfo else parsed.strftime("%Y-%m-%d") -def empty_bucket() -> dict[str, int]: - return { - "inputTokens": 0, - "outputTokens": 0, - "cacheReadInputTokens": 0, - "cacheCreationInputTokens": 0, - } +def model_name(raw): + # A proxy model is often routed as "provider/model"; the model is the last + # part. + value = str(raw or "").strip().rstrip("/").split("/")[-1] + return value or AGENT_ID -def empty_stats() -> dict[str, Any]: - return { - "todayPrompts": 0, - "todaySessions": 0, - "todayTotalTokens": 0, - "todayTokensByModel": {}, - "recentDays": [], - "totalPrompts": 0, - "totalSessions": 0, - "activeDays": 0, - "activeDates": [], - "modelUsage": {}, - } +def runtime_env(): + home = str(Path.home()) + path_parts = [ + os.environ.get("PATH", ""), + f"{home}/.local/bin", + f"{home}/.npm-global/bin", + f"{home}/.local/share/mise/shims", + ] + env = os.environ.copy() + env["PATH"] = os.pathsep.join(part for part in path_parts if part) + return env -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, - # Proxy spend is account-global: every machine that reads the same proxy - # reports the same truth, so aggregation must not sum the records. - "scope": "account", - # api_requests are gateway attempts, not coding-agent prompts. The panel - # would print "0 sessions" next to them, so leave that line off. - "hasPromptStats": False, - "tierLabel": "Proxy", - "usageStatusText": "", - "authHelpText": "", - "limits": [], - } - record.update(empty_stats()) - record.update(overrides) - return record - - -def metric_block(node: Any) -> dict[str, Any]: - # Current proxies nest the numbers under "metrics". Older daily-activity - # payloads put prompt_tokens on the model object itself. - if not isinstance(node, dict): - return {} - nested = node.get("metrics") - if isinstance(nested, dict): - return nested - return node - - -def token_parts(metrics: dict[str, Any]) -> tuple[int, int, int, int, int]: - prompt = number(metrics.get("prompt_tokens")) - output = number(metrics.get("completion_tokens")) - cache_read = number(metrics.get("cache_read_input_tokens")) - cache_write = number(metrics.get("cache_creation_input_tokens")) - cache = cache_read + cache_write - # OpenAI-style rows include cache tokens inside prompt_tokens. Split them - # back out so the four panel buckets stay mutually exclusive. A row whose - # prompt is already smaller than its cache is left alone. - if cache and prompt >= cache: - uncached = prompt - cache - else: - uncached = prompt - total = uncached + cache_read + cache_write + output - return uncached, output, cache_read, cache_write, total - - -MONTHS = { - "jan": 1, "feb": 2, "mar": 3, "apr": 4, "may": 5, "jun": 6, - "jul": 7, "aug": 8, "sep": 9, "oct": 10, "nov": 11, "dec": 12, -} - - -def parse_activity_date(raw: Any, end: date) -> str: - text = str(raw or "").strip() - if len(text) >= 10 and text[4] == "-" and text[7] == "-": - return text[:10] - # Older /global/activity rows use "Aug 25" with no year. The query window - # is at most 30 days, so a month past the window's end belongs to last year. - parts = text.replace(",", " ").split() - if len(parts) < 2: - return "" - month = MONTHS.get(parts[0][:3].lower()) - if month is None: - return "" +ENV = runtime_env() + + +def find_command(name): + return shutil.which(name, path=ENV.get("PATH")) + + +def data_home(): + return Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) + + +def cache_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 read_json(path): try: - day = int(parts[1]) - candidate = date(end.year, month, day) - except ValueError: - return "" - if candidate > end: + return json.loads(path.read_text(encoding="utf-8")) + except Exception: + return None + + +def read_fresh_json(path, max_age_seconds): + if max_age_seconds <= 0: + 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 + except OSError: + return None + return read_json(path) if 0 <= age < max_age_seconds else None + + +def write_json(path, payload): + # A temp name unique to this writer: 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 + + +# ------------------------------------------------------------------- stats + + +def empty_bucket(): + return {"inputTokens": 0, "outputTokens": 0, "cacheReadInputTokens": 0, "cacheCreationInputTokens": 0} + + +class Stats: + """Token and activity totals shared by every local source. Sessions and days + are sets, so a day active in pi and opencode counts once toward activeDays.""" + + def __init__(self): + 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() + self.today_tokens_by_model = {} + self.total_prompts = 0 + self.total_sessions = set() + self.active_days = set() + self.model_usage = {} + + def add(self, session, day, model, input_tokens, output_tokens, cache_read, cache_write): + total = input_tokens + output_tokens + cache_read + cache_write + if total <= 0: + return + self.total_prompts += 1 + self.active_days.add(day) + if session is not None: + self.total_sessions.add(session) + bucket = self.model_usage.setdefault(model, empty_bucket()) + bucket["inputTokens"] += input_tokens + bucket["outputTokens"] += output_tokens + bucket["cacheReadInputTokens"] += cache_read + bucket["cacheCreationInputTokens"] += cache_write + if day in self.recent: + # recentDays.messageCount holds tokens, despite the legacy name. + self.recent[day] += total + if day == self.today: + self.today_prompts += 1 + if session is not None: + self.today_sessions.add(session) + self.today_total_tokens += total + self.today_tokens_by_model[model] = self.today_tokens_by_model.get(model, 0) + total + + def add_session(self, session, today): + self.total_sessions.add(session) + if today: + self.today_sessions.add(session) + + def record(self): + 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, + } + + +# --------------------------------------------------------------- pi and omp + + +def pi_agent_dirs(): + pi_agent = os.environ.get("PI_CODING_AGENT_DIR") + dirs = [Path(pi_agent).expanduser() if pi_agent else Path.home() / ".pi" / "agent", Path.home() / ".omp" / "agent"] + # `omp --profile=` (and pi's equivalent) moves the agent tree under + # /profiles//. + for base in (Path.home() / ".pi", Path.home() / ".omp"): try: - candidate = date(end.year - 1, month, day) - except ValueError: - return "" - return candidate.isoformat() - - -def summarize_activity(payload: dict[str, Any], today: date | None = None) -> dict[str, Any]: - today = today or datetime.now().astimezone().date() - rows = payload.get("results") - daily = payload.get("daily_data") - if isinstance(rows, list) and rows: - return summarize_results(payload, today) - if isinstance(daily, list): - return summarize_daily_data(payload, today) - return summarize_results(payload, today) - - -def summarize_daily_data(payload: dict[str, Any], today: date) -> dict[str, Any]: - # LiteLLM's older global activity route is a day series of request and - # token totals. Model rows for that shape come from /global/activity/model. - recent_dates = [(today - timedelta(days=offset)).isoformat() for offset in range(6, -1, -1)] - recent = {day: 0 for day in recent_dates} - active_dates: set[str] = set() - total_requests = 0 - today_requests = 0 - today_tokens = 0 - rows = payload.get("daily_data") - if not isinstance(rows, list): - rows = [] - for raw_row in rows: - if not isinstance(raw_row, dict): - continue - day = parse_activity_date(raw_row.get("date"), today) - tokens = number(raw_row.get("total_tokens")) - requests = number(raw_row.get("api_requests")) - total_requests += requests - if tokens <= 0 and requests <= 0: - continue - if day: - active_dates.add(day) - if day in recent: - recent[day] += tokens - if day == today.isoformat(): - today_requests += requests - today_tokens += tokens - return { - "todayPrompts": today_requests, - "todayTotalTokens": today_tokens, - "todayTokensByModel": {}, - "recentDays": [{"date": day, "messageCount": recent[day]} for day in recent_dates], - "totalPrompts": total_requests, - "activeDays": len(active_dates), - "activeDates": sorted(active_dates), - "modelUsage": {}, - "windowSpend": 0.0, - } + dirs.extend(sorted(child / "agent" for child in (base / "profiles").iterdir() if child.is_dir())) + except OSError: + pass + return dirs -def summarize_results(payload: dict[str, Any], today: date) -> dict[str, Any]: - recent_dates = [(today - timedelta(days=offset)).isoformat() 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() - total_requests = 0 - today_requests = 0 - window_spend = 0.0 +def pi_pattern(): + return r'"provider"\s*:\s*"litellm"' - rows = payload.get("results") - if not isinstance(rows, list): - rows = [] - for raw_row in rows: - if not isinstance(raw_row, dict): - continue - day = parse_activity_date(raw_row.get("date"), today) - metrics = metric_block(raw_row) - _, _, _, _, total = token_parts(metrics) - requests = number(metrics.get("api_requests") or metrics.get("successful_requests")) - window_spend += money(metrics.get("spend")) - total_requests += requests - if total <= 0 and requests <= 0: - continue - if day: - active_dates.add(day) - if day in recent: - recent[day] += total - if day == today.isoformat(): - today_requests += requests - - breakdown = raw_row.get("breakdown") if isinstance(raw_row.get("breakdown"), dict) else {} - models = breakdown.get("models") if isinstance(breakdown.get("models"), dict) else {} - if models: - for name, raw_model in models.items(): - model_metrics = metric_block(raw_model) - uncached, output, cache_read, cache_write, model_total = token_parts(model_metrics) - if model_total <= 0: - continue - bucket = model_usage.setdefault(str(name), empty_bucket()) - bucket["inputTokens"] += uncached - bucket["outputTokens"] += output - bucket["cacheReadInputTokens"] += cache_read - bucket["cacheCreationInputTokens"] += cache_write - if day == today.isoformat(): - today_by_model[str(name)] = today_by_model.get(str(name), 0) + model_total - elif day == today.isoformat() and total > 0: - today_by_model["unknown"] = today_by_model.get("unknown", 0) + total - - metadata = payload.get("metadata") if isinstance(payload.get("metadata"), dict) else {} - if money(metadata.get("total_spend")) > 0: - window_spend = money(metadata.get("total_spend")) +def accept_pi(message): + return str(message.get("provider") or "") == PROVIDER + + +def pi_session_roots(): + roots = [] + for agent_dir in pi_agent_dirs(): + root = agent_dir / "sessions" + if root.is_dir() and root.resolve() not in [r.resolve() for r in roots]: + roots.append(root) + return roots - return { - "todayPrompts": today_requests, - "todayTotalTokens": sum(today_by_model.values()), - "todayTokensByModel": today_by_model, - "recentDays": [{"date": day, "messageCount": recent[day]} for day in recent_dates], - "totalPrompts": total_requests, - "activeDays": len(active_dates), - "activeDates": sorted(active_dates), - "modelUsage": model_usage, - "windowSpend": window_spend, - } +def pi_parent_session(path): + # A fork names the session it was cut from in its header line. + try: + with open(path, encoding="utf-8", errors="replace") as handle: + header = json.loads(handle.readline()) + except (OSError, ValueError): + return None + parent = header.get("parentSession") if isinstance(header, dict) else None + return os.path.abspath(os.path.expanduser(parent)) if isinstance(parent, str) and parent else None -def summarize_model_series(payload: dict[str, Any], today: date) -> dict[str, Any]: - # /global/activity/model on LiteLLM 1.82 reports one total per model. The - # panel's bar is the sum of the four buckets, so the combined total sits in - # inputTokens. Prompt, cache, and output stay zero because this route does - # not split them. - model_usage: dict[str, dict[str, int]] = {} - today_by_model: dict[str, int] = {} - rows = payload.get("results") if isinstance(payload.get("results"), list) else [] - today_key = today.isoformat() - for raw in rows: - if not isinstance(raw, dict): + +def scan_pi_sessions(stats): + """Add the pi/omp assistant messages this collector counts. Returns whether + every root was read.""" + roots = pi_session_roots() + if not roots: + return True + rg = find_command("rg") or "rg" + try: + proc = subprocess.Popen( + # --no-ignore: these are data files, not a source tree. A $HOME that is + # a git checkout would otherwise hide every session behind its ignores. + [rg, "--json", "--no-ignore", "-e", pi_pattern(), *map(str, roots)], + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + text=True, + errors="replace", + env=ENV, + ) + except OSError: + return False + + counted = set() + owned = {} + today_owned = {} + assert proc.stdout is not None + for raw in proc.stdout: + try: + event = json.loads(raw) + if event.get("type") != "match": + continue + path = os.path.abspath(event["data"]["path"]["text"]) + line_number = event["data"].get("line_number") + entry = json.loads(event["data"]["lines"]["text"]) + except Exception: continue - name = str(raw.get("model") or "").strip() - if not name: + if not isinstance(entry, dict) or entry.get("type") != "message": continue - daily = raw.get("daily_data") if isinstance(raw.get("daily_data"), list) else [] - summed = 0 - today_tokens = 0 - for day_row in daily: - if not isinstance(day_row, dict): - continue - tokens = number(day_row.get("total_tokens")) - summed += tokens - if parse_activity_date(day_row.get("date"), today) == today_key: - today_tokens += tokens - total = number(raw.get("sum_total_tokens")) or summed - if total <= 0: + message = obj(entry.get("message")) + if message.get("role") != "assistant" or not accept_pi(message): + continue + usage = obj(message.get("usage")) + 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 not (input_tokens or output_tokens or cache_read or cache_write): + input_tokens = number(usage.get("totalTokens")) + timestamp = entry.get("timestamp") + day = local_day(timestamp) or local_day(message.get("timestamp")) + if day is None or not (input_tokens or output_tokens or cache_read or cache_write): continue - bucket = empty_bucket() - bucket["inputTokens"] = total - model_usage[name] = bucket - if today_tokens > 0: - today_by_model[name] = today_tokens - return {"modelUsage": model_usage, "todayTokensByModel": today_by_model} + # A fork repeats its parent's messages with their ids and timestamps, so + # the pair says two lines are one message. An id alone is only unique + # within its file, and a line without one always counts. + message_id = str(entry.get("id") or "") + if message_id and timestamp is not None: + key = ("entry", message_id, str(timestamp)) + elif message_id: + key = ("path", path, message_id) + else: + key = ("line", path, line_number) + owned.setdefault(path, set()).add(key) + if day == stats.today: + today_owned.setdefault(path, set()).add(key) + if key in counted: + continue + counted.add(key) + stats.add(None, day, model_name(message.get("model")), input_tokens, output_tokens, cache_read, cache_write) + + try: + code = proc.wait(timeout=5) + except subprocess.TimeoutExpired: + proc.kill() + return False + + # A message belongs to the oldest session holding it, so a fork is a + # session of its own only for what it added, whichever file rg read first. + parents = {} + for path, keys in owned.items(): + inherited = set() + visited = {path} + parent = parents.setdefault(path, pi_parent_session(path)) + while parent and parent not in visited: + visited.add(parent) + inherited |= owned.get(parent, set()) + parent = parents.setdefault(parent, pi_parent_session(parent)) + if keys - inherited: + stats.add_session(path, bool(today_owned.get(path, set()) - inherited)) + # rg exits 1 when nothing matched and 2 when a file couldn't be read. + return code in (0, 1) + + +# ---------------------------------------------------------------- opencode + + +def opencode_db(): + # OPENCODE_DB points opencode at another database, as the CLI itself reads + # it; a relative name lives beside the default one. Only a build whose + # channel was literally "v2" wrote opencode-v2.db. + root = data_home() / "opencode" + configured = os.environ.get("OPENCODE_DB", "").strip() + if configured: + path = Path(configured).expanduser() + return path if path.is_absolute() else root / path + normal = root / "opencode.db" + legacy = root / "opencode-v2.db" + return legacy if not normal.is_file() and legacy.is_file() else normal + + +def table_columns(conn, table): + return {row[1] for row in conn.execute(f"PRAGMA table_info({table})")} + + +def opencode_row(stats, seen, message_id, session_id, entry, row_type, created): + # The original `message` store keeps the provider and model at the top + # level; the V2 `session_message` store nests them under `model` and keeps + # the role in its `type` column, with a compaction's tokens spent like any + # reply's. + if not isinstance(entry, dict) or (entry.get("role") or row_type) not in ("assistant", "compaction"): + return + model = entry.get("model") + provider = entry.get("providerID") or (model.get("providerID") if isinstance(model, dict) else "") + # Exact match: opencode provider ids are free-form. + if str(provider or "") != PROVIDER: + return + tokens = obj(entry.get("tokens")) + cache = obj(tokens.get("cache")) + 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): + return + # A migrated message can sit in both stores under the same id. + if message_id in seen: + return + seen.add(message_id) + day = local_day(obj(entry.get("time")).get("created")) or local_day(created) + if day is None: + return + model_id = entry.get("modelID") or (model.get("id") or model.get("modelID") if isinstance(model, dict) else model) + stats.add("opencode:" + str(session_id), day, model_name(model_id), input_tokens, output_tokens, cache_read, cache_write) + + +def scan_opencode(stats): + """Add opencode messages on the litellm provider, from both of its + stores. Returns whether the database was read to the end.""" + db = opencode_db() + if not db.is_file(): + return True + try: + # Read-only, and not immutable: opencode writes through a WAL, and an + # immutable open would miss what hasn't been checkpointed yet. + conn = sqlite3.connect(db.resolve().as_uri() + "?mode=ro", uri=True, timeout=2) + except sqlite3.Error: + return False + gate = '%"providerID"%:%"' + PROVIDER + '"%' + seen = set() + try: + conn.execute("PRAGMA query_only = ON") + tables = {name for (name,) in conn.execute("SELECT name FROM sqlite_master WHERE type = 'table'")} + if "session_message" in tables and {"id", "session_id", "type", "data"} <= table_columns(conn, "session_message"): + columns = table_columns(conn, "session_message") + created = "m.time_created" if "time_created" in columns else "NULL" + source, fork_filter = "session_message m", "" + if "time_created" in columns and "session_v2" in tables and {"id", "fork_session_id", "time_created"} <= table_columns(conn, "session_v2"): + # A fork copies its source session's messages under new ids; only + # what came after the fork was spent in it. + source += " LEFT JOIN session_v2 s ON s.id = m.session_id" + fork_filter = " AND (s.fork_session_id IS NULL OR m.time_created >= s.time_created)" + for message_id, session_id, row_type, time_created, raw in conn.execute( + f"SELECT m.id, m.session_id, m.type, {created}, m.data FROM {source}" + " WHERE m.type IN ('assistant', 'compaction') AND m.data LIKE ?" + fork_filter, (gate,) + ): + try: + opencode_row(stats, seen, str(message_id), session_id, json.loads(raw), str(row_type or ""), time_created) + except (TypeError, ValueError): + continue + if "message" in tables and {"session_id", "data"} <= table_columns(conn, "message"): + message_id = "id" if "id" in table_columns(conn, "message") else "rowid" + for row_id, session_id, raw in conn.execute(f"SELECT {message_id}, session_id, data FROM message WHERE data LIKE ?", (gate,)): + try: + opencode_row(stats, seen, str(row_id), session_id, json.loads(raw), "", None) + except (TypeError, ValueError): + continue + except sqlite3.Error: + # Transient lock, schema migration, corruption: the numbers stop here, + # incomplete. + return False + finally: + conn.close() + return True + + +# -------------------------------------------------------------- scan cache + + +def scan_cache_paths(): + digest = hashlib.sha1("\n".join([str(Path.home()), str(opencode_db()), os.environ.get("PI_CODING_AGENT_DIR", "")]).encode("utf-8")).hexdigest()[:16] + root = cache_root() + return root / f"litellm-scan-{digest}.json", root / f"litellm-scan-{digest}.lock" + +def run_local_scans(): + stats = Stats() + complete = scan_pi_sessions(stats) + complete = scan_opencode(stats) and complete + return stats.record(), complete -def config_path() -> Path: - root = os.environ.get("XDG_CONFIG_HOME") or str(Path.home() / ".config") - return Path(root) / "omarchy" / "agents" / "litellm.json" +def read_cached_stats(cache_file, max_age): + cached = read_fresh_json(cache_file, max_age) + if not isinstance(cached, dict) or cached.get("schemaVersion") != 1: + return None + # today* fields only mean "today" on the day they were scanned. + if cached.get("scanDate") != datetime.now().strftime("%Y-%m-%d"): + return None + stats = cached.get("stats") + return stats if isinstance(stats, dict) and set(Stats().record()) <= set(stats) else None -def read_config() -> dict[str, Any]: - path = config_path() - if not path.is_file(): - return {} + +def cached_local_stats(max_age): + """Local stats, with the cache as a pure optimization: any cache-layer + failure degrades to a direct scan and a warning on stderr.""" try: - parsed = json.loads(path.read_text()) - except (OSError, json.JSONDecodeError): - return {} - return parsed if isinstance(parsed, dict) else {} + cache_file, lock_file = scan_cache_paths() + cached = read_cached_stats(cache_file, max_age) + if cached is not None: + return cached + with lock_file.open("w") as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + cached = read_cached_stats(cache_file, max_age) + if cached is not None: + return cached + stats, complete = run_local_scans() + # A scan that couldn't read everything serves this run, but caching it + # would hide the missing usage from every reader until it expires. + if complete: + write_json(cache_file, {"schemaVersion": 1, "scanDate": datetime.now().strftime("%Y-%m-%d"), "stats": stats}) + return stats + except Exception as exc: + print(f"omarchy-agent-usage-litellm: cache unavailable ({exc}); scanning directly", file=sys.stderr) + return run_local_scans()[0] + + +# ------------------------------------------------------------- credentials + + +def config_path(): + root = Path(os.environ.get("XDG_CONFIG_HOME") or (Path.home() / ".config")) + return root / "omarchy" / "agents" / "litellm.json" + + +def credentials(): + config = obj(read_json(config_path())) + base = str(os.environ.get("LITELLM_BASE_URL") or config.get("baseUrl") or "").strip() + key = str(os.environ.get("LITELLM_API_KEY") or os.environ.get("LITELLM_MASTER_KEY") or config.get("apiKey") or "").strip() + return base, key -def normalize_base_url(raw: str) -> str: +# ------------------------------------------------------------------ budget + + +class LiteLLMError(Exception): + def __init__(self, message, refused=False, transport=False): + super().__init__(message) + self.refused = refused + self.transport = transport + + +def normalize_base_url(raw): url = raw.strip().rstrip("/") if url.endswith("/v1"): url = url[: -len("/v1")] parsed = urllib.parse.urlparse(url) if parsed.scheme not in ("http", "https") or not parsed.hostname: - raise LiteLLMError("LiteLLM base URL must be http or https") + raise LiteLLMError("The LiteLLM base URL must be http or https.") if parsed.username or parsed.password: - raise LiteLLMError("LiteLLM base URL must not carry credentials") + raise LiteLLMError("The LiteLLM base URL must not carry credentials.") return urllib.parse.urlunparse((parsed.scheme, parsed.netloc, parsed.path.rstrip("/"), "", "", "")) -def credentials(config: dict[str, Any]) -> tuple[str, str, float | None]: - base = str(os.environ.get("LITELLM_BASE_URL") or config.get("baseUrl") or "").strip() - key = str( - os.environ.get("LITELLM_API_KEY") - or os.environ.get("LITELLM_MASTER_KEY") - or config.get("apiKey") - or "" - ).strip() - budget_raw = config.get("maxBudget") - budget = money(budget_raw) if budget_raw not in (None, "") else None - if budget == 0: - budget = None - return base, key, budget - - -class SameOriginRedirectHandler(urllib.request.HTTPRedirectHandler): - def __init__(self, host: str) -> None: +# The key travels in every request, so a proxy that redirects elsewhere must +# not take it along. +class SameHostRedirectHandler(urllib.request.HTTPRedirectHandler): + def __init__(self, host): super().__init__() self.host = host def redirect_request(self, req, fp, code, msg, headers, newurl): parsed = urllib.parse.urlparse(newurl) if parsed.scheme not in ("http", "https") or parsed.hostname != self.host: - raise urllib.error.HTTPError(newurl, code, "cross-origin redirect refused", headers, fp) + raise LiteLLMError("LiteLLM redirected off the configured host, so the key wasn't sent there.") return super().redirect_request(req, fp, code, msg, headers, newurl) -class LiteLLMClient: - def __init__(self, base_url: str, api_key: str) -> None: - self.base_url = normalize_base_url(base_url) - self.api_key = api_key - self.host = urllib.parse.urlparse(self.base_url).hostname or "" - - def get(self, path: str, query: dict[str, str] | None = None) -> dict[str, Any]: - url = self.base_url + path - if query: - url += "?" + urllib.parse.urlencode(query) - request = urllib.request.Request( - url, - headers={"Authorization": "Bearer " + self.api_key, "Accept": "application/json"}, - method="GET", - ) - opener = urllib.request.build_opener(SameOriginRedirectHandler(self.host)) - try: - with opener.open(request, timeout=20) as response: - final = getattr(response, "geturl", lambda: url)() - if urllib.parse.urlparse(final).hostname != self.host: - raise LiteLLMError("LiteLLM redirected off the configured host") - payload = json.loads(response.read().decode("utf-8", errors="replace")) - except urllib.error.HTTPError as error: - if error.code in (401, 403): - raise LiteLLMError("LiteLLM rejected the API key", error.code) from error - raise LiteLLMError(f"LiteLLM returned HTTP {error.code}", error.code) from error - except LiteLLMError: - raise - except Exception as error: - raise LiteLLMError("LiteLLM did not respond") from error - if isinstance(payload, list): - return {"results": payload} - if not isinstance(payload, dict): - raise LiteLLMError("LiteLLM returned an unexpected payload") - return payload - - def _paged(self, path: str, query: dict[str, str]) -> dict[str, Any]: - payload = self.get(path, query) - page = 1 - while page < 20: - meta = payload.get("metadata") if isinstance(payload.get("metadata"), dict) else {} - if not meta.get("has_more"): - break - page += 1 - nxt = self.get(path, {**query, "page": str(page)}) - rows = nxt.get("results") - if isinstance(rows, list): - payload.setdefault("results", []) - if isinstance(payload["results"], list): - payload["results"].extend(rows) - payload["metadata"] = nxt.get("metadata") if isinstance(nxt.get("metadata"), dict) else {} - return payload - - def _has_usage(self, payload: dict[str, Any]) -> bool: - if isinstance(payload.get("daily_data"), list): - return True - if isinstance(payload.get("results"), list) and payload["results"]: - return True - return False +def money(value): + value = finite(value) + return value if value is not None and value >= 0 else None - def activity(self, start_day: date, end_day: date) -> tuple[dict[str, Any], str]: - query = {"start_date": start_day.isoformat(), "end_date": end_day.isoformat()} - try: - payload = self._paged("/global/activity", query) - except LiteLLMError as error: - if error.status not in (403, 404): - raise - payload = None - if payload is not None and self._has_usage(payload): - return payload, "global" - return self._paged("/user/daily/activity", query), "user" - - def model_series(self, start_day: date, end_day: date) -> dict[str, Any] | None: - query = {"start_date": start_day.isoformat(), "end_date": end_day.isoformat()} - try: - return self.get("/global/activity/model", query) - except LiteLLMError: - return None - def key_budget(self) -> tuple[float, float] | None: - try: - payload = self.get("/key/info") - except LiteLLMError: - return None - info = payload.get("info") if isinstance(payload.get("info"), dict) else payload - if not isinstance(info, dict) or info.get("max_budget") in (None, ""): - return None - funded = money(info.get("max_budget")) - if funded <= 0: - return None - return money(info.get("spend")), funded - - -def balance_from(spent: float, funded: float | None, estimated: bool) -> dict[str, Any] | None: - if funded is None or funded <= 0: +def fetch_budget(base_url, key): + """The key's budget from /key/info as a balance, or None when it has none.""" + base = normalize_base_url(base_url) + host = urllib.parse.urlparse(base).hostname + request = urllib.request.Request(base + "/key/info", headers={"Authorization": "Bearer " + key, "Accept": "application/json"}) + opener = urllib.request.build_opener(SameHostRedirectHandler(host)) + try: + with opener.open(request, timeout=10) as response: + if urllib.parse.urlparse(response.geturl()).hostname != host: + raise LiteLLMError("LiteLLM redirected off the configured host, so the key wasn't sent there.") + payload = json.loads(response.read().decode("utf-8", errors="replace")) + except LiteLLMError: + raise + except urllib.error.HTTPError as error: + if error.code in (401, 403): + raise LiteLLMError("LiteLLM rejected the API key. Set a current one in LITELLM_API_KEY or ~/.config/omarchy/agents/litellm.json.", refused=True) + raise LiteLLMError(f"LiteLLM's /key/info returned status {error.code}. Local LiteLLM stats are still shown.") + except (urllib.error.URLError, TimeoutError, OSError): + # A transport failure reached no server at all: no route, no DNS. + raise LiteLLMError("Couldn't reach the LiteLLM proxy. Retrying shortly. Local LiteLLM stats are still shown.", transport=True) + except ValueError: + raise LiteLLMError("LiteLLM returned an unreadable /key/info answer. Local LiteLLM stats are still shown.") + + payload = obj(payload) + info = obj(payload.get("info")) or payload + funded = money(info.get("max_budget")) + if not funded: + # A key without a budget can spend without limit: there is no balance. return None + spent = money(info.get("spend")) or 0.0 return { "remaining": round(max(0.0, funded - spent), 4), "funded": round(funded, 4), "spent": round(spent, 4), "currency": "USD", - "estimated": estimated, + "estimated": False, } -def scan(config: dict[str, Any] | None = None) -> dict[str, Any]: - config = read_config() if config is None else config - base_url, api_key, configured_budget = credentials(config) - if not base_url or not api_key: - return base_record(authHelpText=AUTH_HELP) +# One cache per proxy and key, so another key never shows this one's budget. +def budget_cache_path(base_url, key): + digest = hashlib.sha256((base_url + "\n" + key).encode("utf-8")).hexdigest()[:16] + return cache_root() / f"litellm-limits-{digest}.json" + +def collect_budget(base_url, key, force): + result = {"balance": None, "usageStatusText": "", "authHelpText": "", "stale": False, "fetchedAtMs": 0, "retryAdvised": False} + if not (base_url and key): + return result try: - client = LiteLLMClient(base_url, api_key) - today = datetime.now().astimezone().date() - start = today - timedelta(days=WINDOW_DAYS - 1) - payload, scope_name = client.activity(start, today) + # Normalized first, so ".../v1/" and "..." share one cache. + base_url = normalize_base_url(base_url) except LiteLLMError as error: - return base_record(usageStatusText="LiteLLM unavailable", authHelpText=str(error)) - - stats = summarize_activity(payload, today) - if scope_name == "global" and not stats["modelUsage"]: - series = client.model_series(start, today) - if series: - models = summarize_model_series(series, today) - if models["modelUsage"]: - stats["modelUsage"] = models["modelUsage"] - stats["todayTokensByModel"] = models["todayTokensByModel"] - window_spend = float(stats.pop("windowSpend")) - record = base_record( - ready=True, - hasLocalStats=stats["todayTotalTokens"] > 0 or stats["totalPrompts"] > 0 or bool(stats["modelUsage"]), - tierLabel="Proxy" if scope_name == "global" else "This key", - **stats, - ) - - key_budget = client.key_budget() - if key_budget is not None: - spent, funded = key_budget - ledger = balance_from(spent, funded, estimated=False) - else: - ledger = balance_from(window_spend, configured_budget, estimated=True) - if ledger is not None: - record["balance"] = ledger - return record - - -def main() -> int: - parser = argparse.ArgumentParser(description="Print the LiteLLM usage record as JSON") - # One activity request is the whole record, so there is no cache to force - # past and no faster limits-only path. The flags exist so every collector - # accepts the same invocation. - parser.add_argument("--force", action="store_true") - parser.add_argument("--limits-only", action="store_true") - parser.parse_args() + result.update(usageStatusText="Balance unavailable", authHelpText=str(error)) + return result + + cache = budget_cache_path(base_url, key) + cached = obj(read_json(cache)) + kept = cached.get("balance") if isinstance(cached.get("balance"), dict) else None + fetched_at = cached.get("fetchedAtMs") + fetched_at = round(fetched_at) if epoch_seconds(fetched_at) else 0 + + # A panel opened and closed repeatedly must not become one request per + # flick; --force is a person asking for fresh numbers. + if fetched_at and not force and 0 <= time.time() - fetched_at / 1000 < PROBE_MIN_INTERVAL_SECONDS: + result.update(balance=kept, fetchedAtMs=fetched_at) + return result try: - record = scan() + balance = fetch_budget(base_url, key) except LiteLLMError as error: - record = base_record(usageStatusText="LiteLLM unavailable", authHelpText=str(error)) - except Exception: - record = base_record(usageStatusText="LiteLLM unavailable", authHelpText="LiteLLM usage scan failed") - print("omarchy-agent-usage-litellm: scan failed", file=sys.stderr) - print(json.dumps(record, separators=(",", ":"))) + # The last good budget, marked stale, beats an empty section. + result.update(balance=kept, stale=kept is not None, fetchedAtMs=fetched_at if kept else 0, retryAdvised=error.transport) + if error.refused: + result.update(usageStatusText="Sign-in expired", authHelpText=str(error)) + elif kept is None: + result.update(usageStatusText="Balance unavailable", authHelpText=str(error)) + return result + + fetched_at = round(time.time() * 1000) + try: + write_json(cache, {"fetchedAtMs": fetched_at, "balance": balance}) + except OSError as exc: + print(f"omarchy-agent-usage-litellm: could not write the budget cache ({exc})", file=sys.stderr) + result.update(balance=balance, fetchedAtMs=fetched_at) + return result + + +# ------------------------------------------------------------------ record + + +def main(): + parser = argparse.ArgumentParser(description="Print the LiteLLM usage record as JSON") + parser.add_argument("--force", action="store_true", help="rescan local sessions and re-check the budget, ignoring caches") + parser.add_argument("--limits-only", action="store_true", help="reuse a recent local scan; only the budget must be fresh") + args = parser.parse_args() + + max_age = 0 if args.force else (LIMITS_ONLY_REUSE_SECONDS if args.limits_only else SCAN_REUSE_SECONDS) + stats = cached_local_stats(max_age) + base_url, key = credentials() + budget = collect_budget(base_url, key, args.force) + in_use = number(stats.get("totalPrompts")) > 0 + + record = { + "schemaVersion": 1, + "id": AGENT_ID, + "name": AGENT_NAME, + "updatedAt": datetime.now(timezone.utc).isoformat(), + "ready": in_use or bool(base_url and key), + "hasLocalStats": True, + # A proxy names no plan; a key's budget is a cap, not credit bought. + "tierLabel": "", + "usageStatusText": budget["usageStatusText"], + "authHelpText": budget["authHelpText"], + "limits": [], + "limitsStale": budget["stale"], + "limitsFetchedAt": budget["fetchedAtMs"], + } + if budget["balance"] is not None: + record["balance"] = budget["balance"] + if in_use and not (base_url and key): + record.update(usageStatusText="Waiting for auth", authHelpText=AUTH_HELP) + if budget["retryAdvised"]: + record["retryAdvised"] = True + record.update(stats) + print(json.dumps(record, separators=(",", ":"), sort_keys=True)) return 0 diff --git a/test/shell.d/agent-usage-litellm-scanner-test.sh b/test/shell.d/agent-usage-litellm-scanner-test.sh index 16eadf2ae4c..010721e302d 100755 --- a/test/shell.d/agent-usage-litellm-scanner-test.sh +++ b/test/shell.d/agent-usage-litellm-scanner-test.sh @@ -4,311 +4,209 @@ source "$(dirname "$0")/base-test.sh" require_command jq require_command python3 +require_command rg TEST_HOME=$(mktemp -d) -trap 'rm -rf "$TEST_HOME"' EXIT - -mkdir -p "$TEST_HOME/.config/omarchy/agents" - -# Without a URL and key the collector must still print a full record. The -# update runner writes whatever valid JSON appears on stdout, and the panel -# hides a record that has no usage. -no_key=$(HOME="$TEST_HOME" XDG_CONFIG_HOME="$TEST_HOME/.config" \ - env -u LITELLM_BASE_URL -u LITELLM_API_KEY -u LITELLM_MASTER_KEY \ - "$ROOT/bin/omarchy-agent-usage-litellm") - -[[ $(jq -r '.id + ":" + (.ready | tostring) + ":" + (.hasPromptStats | tostring) + ":" + .scope' <<<"$no_key") == "litellm:false:false:account" ]] || - fail "LiteLLM collector prints a valid record without credentials" "$no_key" -pass "LiteLLM collector prints a valid record without credentials" - -[[ $(jq -r '.authHelpText | contains("LITELLM_BASE_URL")' <<<"$no_key") == "true" ]] || - fail "LiteLLM collector tells the user how to configure the proxy" "$no_key" -pass "LiteLLM collector tells the user how to configure the proxy" - -result=$(python3 - "$ROOT/bin/omarchy-agent-usage-litellm" <<'PY' -import importlib.machinery -import importlib.util -import json -import os -import sys -import time -from datetime import date -from pathlib import Path - -os.environ["TZ"] = "UTC" -time.tzset() -for name in ("LITELLM_BASE_URL", "LITELLM_API_KEY", "LITELLM_MASTER_KEY"): - os.environ.pop(name, None) - -collector_path = str(Path(sys.argv[1])) -loader = importlib.machinery.SourceFileLoader("litellm_collector", collector_path) -spec = importlib.util.spec_from_loader(loader.name, loader) -scanner = importlib.util.module_from_spec(spec) -loader.exec_module(scanner) - -payload = { - "results": [ - { - "date": "2026-07-31", - "metrics": { - "spend": 0.4, - "prompt_tokens": 100, - "completion_tokens": 20, - "cache_read_input_tokens": 40, - "cache_creation_input_tokens": 10, - "api_requests": 3, - }, - "breakdown": { - "models": { - "ollama/qwen2.5": { - "metrics": { - "prompt_tokens": 100, - "completion_tokens": 20, - "cache_read_input_tokens": 40, - "cache_creation_input_tokens": 10, - } - } - } - }, - }, - { - "date": "2026-07-30T00:00:00Z", - "metrics": { - "prompt_tokens": 50, - "completion_tokens": 10, - "api_requests": 1, - }, - "breakdown": { - "models": { - "gpt-4": {"prompt_tokens": 50, "completion_tokens": 10} - } - }, - }, - ], - "metadata": {"total_spend": 1.25}, +SERVER_PID="" +cleanup() { + [[ -n $SERVER_PID ]] && kill "$SERVER_PID" 2>/dev/null + rm -rf "$TEST_HOME" } - -summary = scanner.summarize_activity(payload, date(2026, 7, 31)) -legacy = scanner.summarize_activity({ - "daily_data": [ - {"date": "Jul 31", "api_requests": 4, "total_tokens": 90}, - {"date": "Jul 30", "api_requests": 1, "total_tokens": 10}, - {"date": "Jan 02", "api_requests": 1, "total_tokens": 5}, - ], - "sum_total_tokens": 105, -}, date(2026, 7, 31)) -summary["legacyToday"] = legacy["todayTotalTokens"] -summary["legacyRequests"] = legacy["totalPrompts"] -summary["legacyModels"] = legacy["modelUsage"] -summary["legacyYear"] = scanner.parse_activity_date("Dec 28", date(2026, 1, 5)) == "2025-12-28" -series = scanner.summarize_model_series({ - "results": [ - { - "model": "ollama/qwen", - "sum_total_tokens": 90, - "daily_data": [{"date": "Jul 31", "total_tokens": 70, "api_requests": 2}], - }, - {"model": "", "sum_total_tokens": 10, "daily_data": []}, - {"model": "quiet", "sum_total_tokens": 0, "daily_data": []}, - ] -}, date(2026, 7, 31)) -summary["seriesTotal"] = series["modelUsage"]["ollama/qwen"]["inputTokens"] -summary["seriesSplit"] = series["modelUsage"]["ollama/qwen"]["outputTokens"] -summary["seriesToday"] = series["todayTokensByModel"]["ollama/qwen"] -summary["seriesSkipped"] = list(series["modelUsage"]) == ["ollama/qwen"] -summary["uncached"] = summary["modelUsage"]["ollama/qwen2.5"]["inputTokens"] -summary["cacheRead"] = summary["modelUsage"]["ollama/qwen2.5"]["cacheReadInputTokens"] -summary["flatModel"] = summary["modelUsage"]["gpt-4"]["outputTokens"] - -calls = [] -RealClient = scanner.LiteLLMClient - -class GlobalClient: - def __init__(self, base_url, api_key): - calls.append(("init", base_url, api_key)) - - def activity(self, start_day, end_day): - calls.append(("activity", start_day.isoformat(), end_day.isoformat())) - return payload, "global" - - def key_budget(self): - return (4.0, 20.0) - -class ForbiddenGlobal: - def __init__(self, base_url, api_key): - pass - - def activity(self, start_day, end_day): - raise scanner.LiteLLMError("LiteLLM rejected the API key", 403) - -class UserFallback(ForbiddenGlobal): - def activity(self, start_day, end_day): - return payload, "user" - - def key_budget(self): - return None - -scanner.LiteLLMClient = GlobalClient -record = scanner.scan({"baseUrl": "http://proxy.example:4000/v1", "apiKey": "sk-test"}) -summary["record"] = { - "id": record["id"], - "ready": record["ready"], - "scope": record["scope"], - "tierLabel": record["tierLabel"], - "hasPromptStats": record["hasPromptStats"], - "limits": record["limits"], - "balance": record["balance"], - "todayTotalTokens": record["todayTotalTokens"], -} -summary["normalizedBase"] = scanner.normalize_base_url("http://proxy.example:4000/v1") -start = date.fromisoformat(calls[1][1]) -end = date.fromisoformat(calls[1][2]) -summary["windowDays"] = (end - start).days == 29 and end == date.today() - -scanner.LiteLLMClient = UserFallback -user_record = scanner.scan({"baseUrl": "https://proxy.example", "apiKey": "sk-user", "maxBudget": 10}) -summary["userTier"] = user_record["tierLabel"] -summary["estimatedBalance"] = user_record["balance"] - -class LegacyModelClient: - def __init__(self, base_url, api_key): +trap cleanup EXIT + +export HOME="$TEST_HOME" +export XDG_CACHE_HOME="$TEST_HOME/cache" +export XDG_DATA_HOME="$TEST_HOME/data" +export XDG_CONFIG_HOME="$TEST_HOME/config" +unset LITELLM_BASE_URL LITELLM_API_KEY LITELLM_MASTER_KEY PI_CODING_AGENT_DIR OPENCODE_DB + +collector="$ROOT/bin/omarchy-agent-usage-litellm" +STATE="$TEST_HOME/server" +mkdir -p "$STATE" +BUDGET_CACHE="$XDG_CACHE_HOME/omarchy/agent-usage" + +# A stand-in proxy: /key/info answers with $STATE/body and the status in +# $STATE/status, or redirects to $STATE/redirect when that is set. Every +# request and its key is noted, so a refused redirect can prove the key +# never left. +python3 - "$STATE" <<'PY' & +import http.server, pathlib, sys +state = pathlib.Path(sys.argv[1]) + +class Handler(http.server.BaseHTTPRequestHandler): + def do_GET(self): + with (state / "log").open("a") as log: + log.write(f"{self.headers.get('Host')} {self.path} {self.headers.get('Authorization', '')}\n") + redirect = state / "redirect" + if redirect.exists(): + self.send_response(302) + self.send_header("Location", redirect.read_text().strip()) + self.end_headers() + return + status = int((state / "status").read_text()) if (state / "status").exists() else 200 + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.end_headers() + self.wfile.write((state / "body").read_bytes() if (state / "body").exists() else b"{}") + + def log_message(self, *args): pass - def activity(self, start_day, end_day): - return {"daily_data": [{"date": "Jul 31", "api_requests": 4, "total_tokens": 90}]}, "global" - - def model_series(self, start_day, end_day): - return { - "results": [{ - "model": "ollama/qwen", - "sum_total_tokens": 90, - "daily_data": [{"date": "Jul 31", "total_tokens": 90}], - }] - } - - def key_budget(self): - return None - -scanner.LiteLLMClient = LegacyModelClient -legacy_record = scanner.scan({"baseUrl": "http://proxy.example:4000", "apiKey": "sk-old"}) -summary["legacyScanModel"] = legacy_record["modelUsage"]["ollama/qwen"]["inputTokens"] -summary["legacyScanReady"] = legacy_record["ready"] and legacy_record["tierLabel"] == "Proxy" - -refused = False -try: - scanner.normalize_base_url("http://user:secret@proxy.example:4000") -except scanner.LiteLLMError: - refused = True -summary["refusesCredentialsInUrl"] = refused - -# A 401 is the key itself, not a missing admin route, so activity() must not -# fall through to the user endpoint. -class RejectingGet: - def __init__(self): - self.paths = [] - - def get(self, path, query=None): - self.paths.append(path) - raise scanner.LiteLLMError("LiteLLM rejected the API key", 401) - -client = RealClient.__new__(RealClient) -rejecting = RejectingGet() -client.get = rejecting.get -try: - client.activity(date(2026, 7, 1), date(2026, 7, 31)) - summary["authStops"] = False -except scanner.LiteLLMError: - summary["authStops"] = rejecting.paths == ["/global/activity"] - -class FallbackGet: - def __init__(self): - self.paths = [] - - def get(self, path, query=None): - self.paths.append(path) - if path == "/global/activity": - raise scanner.LiteLLMError("LiteLLM rejected the API key", 403) - return payload - -client.get = FallbackGet().get -body, scope_name = client.activity(date(2026, 7, 1), date(2026, 7, 31)) -summary["fallsBack"] = scope_name == "user" and body["metadata"]["total_spend"] == 1.25 -print(json.dumps(summary, separators=(",", ":"))) +server = http.server.HTTPServer(("127.0.0.1", 0), Handler) +(state / "port").write_text(str(server.server_address[1])) +server.serve_forever() PY -) - -[[ $(jq -r '.todayTotalTokens' <<<"$result") == "120" ]] || - fail "LiteLLM collector totals today's uncached, cache, and output tokens once" "$result" -pass "LiteLLM collector totals today's uncached, cache, and output tokens once" - -[[ $(jq -r '.uncached' <<<"$result") == "50" ]] || - fail "LiteLLM collector splits cache tokens out of prompt tokens" "$result" -pass "LiteLLM collector splits cache tokens out of prompt tokens" - -[[ $(jq -r '.cacheRead' <<<"$result") == "40" ]] || - fail "LiteLLM collector keeps cache reads on their own counter" "$result" -pass "LiteLLM collector keeps cache reads on their own counter" - -[[ $(jq -r '.flatModel' <<<"$result") == "10" ]] || - fail "LiteLLM collector accepts model rows that are not nested under metrics" "$result" -pass "LiteLLM collector accepts model rows that are not nested under metrics" - -[[ $(jq -r '.recentDays[-1].messageCount' <<<"$result") == "120" ]] || - fail "LiteLLM collector builds the seven-day token series" "$result" -pass "LiteLLM collector builds the seven-day token series" - -[[ $(jq -r '.totalPrompts' <<<"$result") == "4" ]] || - fail "LiteLLM collector counts logged API requests" "$result" -pass "LiteLLM collector counts logged API requests" - -[[ $(jq -r '.windowSpend' <<<"$result") == "1.25" ]] || - fail "LiteLLM collector prefers the metadata spend total" "$result" -pass "LiteLLM collector prefers the metadata spend total" - -[[ $(jq -c '.record | {id, ready, scope, tierLabel, hasPromptStats, limits}' <<<"$result") == '{"id":"litellm","ready":true,"scope":"account","tierLabel":"Proxy","hasPromptStats":false,"limits":[]}' ]] || - fail "LiteLLM collector prints the display-ready record contract" "$result" -pass "LiteLLM collector prints the display-ready record contract" - -[[ $(jq -c '.record.balance' <<<"$result") == '{"remaining":16.0,"funded":20.0,"spent":4.0,"currency":"USD","estimated":false}' ]] || - fail "LiteLLM collector uses the key budget when the proxy reports one" "$result" -pass "LiteLLM collector uses the key budget when the proxy reports one" - -[[ $(jq -r '.normalizedBase' <<<"$result") == "http://proxy.example:4000" ]] || - fail "LiteLLM collector strips a trailing /v1 from the proxy URL" "$result" -pass "LiteLLM collector strips a trailing /v1 from the proxy URL" - -[[ $(jq -r '.windowDays' <<<"$result") == "true" ]] || - fail "LiteLLM collector asks for the last 30 days" "$result" -pass "LiteLLM collector asks for the last 30 days" - -[[ $(jq -r '.userTier' <<<"$result") == "This key" ]] || - fail "LiteLLM collector labels a non-admin key as its own usage" "$result" -pass "LiteLLM collector labels a non-admin key as its own usage" - -[[ $(jq -c '.estimatedBalance' <<<"$result") == '{"remaining":8.75,"funded":10.0,"spent":1.25,"currency":"USD","estimated":true}' ]] || - fail "LiteLLM collector estimates a balance from the configured budget" "$result" -pass "LiteLLM collector estimates a balance from the configured budget" - -[[ $(jq -r '.refusesCredentialsInUrl' <<<"$result") == "true" ]] || - fail "LiteLLM collector refuses a base URL that embeds a password" "$result" -pass "LiteLLM collector refuses a base URL that embeds a password" - -[[ $(jq -r '.authStops' <<<"$result") == "true" ]] || - fail "LiteLLM collector does not fall back when the key itself is rejected" "$result" -pass "LiteLLM collector does not fall back when the key itself is rejected" - -[[ $(jq -r '.fallsBack' <<<"$result") == "true" ]] || - fail "LiteLLM collector falls back to the user activity route" "$result" -pass "LiteLLM collector falls back to the user activity route" +SERVER_PID=$! +for _ in $(seq 50); do + [[ -s $STATE/port ]] && break + sleep 0.1 +done +[[ -s $STATE/port ]] || fail "LiteLLM test server starts" +PORT=$(cat "$STATE/port") +PROXY="http://127.0.0.1:$PORT" + +requests() { + [[ -f $STATE/log ]] && wc -l <"$STATE/log" || echo 0 +} -[[ $(jq -r '.legacyToday' <<<"$result") == "90" && $(jq -r '.legacyRequests' <<<"$result") == "6" && $(jq -r '.legacyYear' <<<"$result") == "true" && $(jq -c '.legacyModels' <<<"$result") == "{}" ]] || - fail "LiteLLM collector reads the older global activity day series" "$result" -pass "LiteLLM collector reads the older global activity day series" +serve() { + printf '%s' "$1" >"$STATE/status" + printf '%s' "${2-}" >"$STATE/body" +} -[[ $(jq -r '.seriesTotal' <<<"$result") == "90" && $(jq -r '.seriesSplit' <<<"$result") == "0" && $(jq -r '.seriesToday' <<<"$result") == "70" && $(jq -r '.seriesSkipped' <<<"$result") == "true" ]] || - fail "LiteLLM collector reads per-model totals from the model activity route" "$result" -pass "LiteLLM collector reads per-model totals from the model activity route" +serve 200 '{"key": "sk-hash", "info": {"key_alias": "laptop", "spend": 12.5, "max_budget": 50, "budget_reset_at": null}}' + +# ------------------------------------------------------------ unused machine + +result=$("$collector") +[[ $(jq -c '[.ready, .limits, has("balance"), .usageStatusText, .authHelpText, .tierLabel, .todayTotalTokens, (.recentDays | length), has("retryAdvised")]' <<<"$result") == '[false,[],false,"","","",0,7,false]' ]] || + fail "LiteLLM collector prints an empty, not-ready record on an unused machine" "$result" +[[ $(requests) == "0" ]] || + fail "LiteLLM collector asks nothing without a key" "$(requests)" +pass "LiteLLM collector prints an empty, not-ready record on an unused machine" + +# ----------------------------------------------------------------- budget + +mkdir -p "$XDG_CONFIG_HOME/omarchy/agents" +printf '{"baseUrl": "%s/v1/", "apiKey": "config-key"}\n' "$PROXY" >"$XDG_CONFIG_HOME/omarchy/agents/litellm.json" +result=$("$collector") +[[ $(jq -c '[.ready, (.balance == {"currency":"USD","estimated":false,"funded":50,"remaining":37.5,"spent":12.5}), .limitsStale, .usageStatusText, .authHelpText, .tierLabel, .hasPromptStats // true, has("scope")]' <<<"$result") == '[true,true,false,"","","",true,false]' ]] || + fail "LiteLLM collector reads the key's budget as a balance" "$result" +[[ $(tail -n 1 "$STATE/log") == "127.0.0.1:$PORT /key/info Bearer config-key" ]] || + fail "LiteLLM collector asks /key/info with the configured key" "$(cat "$STATE/log")" +pass "LiteLLM collector reads the key's budget from /key/info" + +[[ $(LITELLM_BASE_URL="$PROXY" LITELLM_API_KEY=env-key "$collector" --force >/dev/null && tail -n 1 "$STATE/log") == *"Bearer env-key" ]] || + fail "LiteLLM collector prefers the environment's key" "$(cat "$STATE/log")" +pass "LiteLLM collector prefers the environment's key" + +# A key without a budget spends without limit: no balance, nothing wrong. +serve 200 '{"info": {"spend": 3, "max_budget": null}}' +result=$("$collector" --force) +[[ $(jq -c '[.ready, has("balance"), .usageStatusText]' <<<"$result") == '[true,false,""]' ]] || + fail "LiteLLM collector shows no balance for a key without a budget" "$result" +pass "LiteLLM collector shows no balance for a key without a budget" + +# ------------------------------------------------------------- probe reuse + +serve 200 '{"info": {"spend": 10, "max_budget": 40}}' +"$collector" --force >/dev/null +before=$(requests) +result=$("$collector") +[[ $(requests) == "$before" && $(jq -r '.balance.remaining == 30' <<<"$result") == "true" ]] || + fail "LiteLLM collector reuses a check from seconds ago" "$result" +"$collector" --force >/dev/null +[[ $(requests) == "$((before + 1))" ]] || + fail "LiteLLM collector --force checks again" "$(requests)" +pass "LiteLLM collector reuses a recent check and --force checks again" + +# --------------------------------------------------------------- failures + +cache_file=$(ls "$BUDGET_CACHE"/litellm-limits-*.json | head -n 1) +for file in "$BUDGET_CACHE"/litellm-limits-*.json; do + jq -c --argjson t "$((($(date +%s) - 3600) * 1000))" '.fetchedAtMs = $t' "$file" >"$file.tmp" && mv "$file.tmp" "$file" +done + +serve 500 +result=$("$collector") +[[ $(jq -c '[(.balance.remaining == 30), .limitsStale, .usageStatusText, has("retryAdvised"), .limitsFetchedAt > 0]' <<<"$result") == '[true,true,"",false,true]' ]] || + fail "LiteLLM collector keeps the last budget stale after a failed check" "$result" +pass "LiteLLM collector keeps the last budget stale after a failed check" + +printf '{"baseUrl": "http://127.0.0.1:1", "apiKey": "config-key"}\n' >"$XDG_CONFIG_HOME/omarchy/agents/litellm.json" +result=$("$collector") +[[ $(jq -c '[.retryAdvised, .usageStatusText, has("balance")]' <<<"$result") == '[true,"Balance unavailable",false]' ]] || + fail "LiteLLM collector advises a retry when the proxy is unreachable" "$result" +pass "LiteLLM collector advises a retry only when the proxy is unreachable" +printf '{"baseUrl": "%s", "apiKey": "config-key"}\n' "$PROXY" >"$XDG_CONFIG_HOME/omarchy/agents/litellm.json" + +serve 401 +result=$("$collector") +[[ $(jq -c '[.usageStatusText, .limitsStale, (.balance.remaining == 30)]' <<<"$result") == '["Sign-in expired",true,true]' ]] || + fail "LiteLLM collector says a refused key needs replacing" "$result" +pass "LiteLLM collector says a refused key needs replacing" + +rm -f "$BUDGET_CACHE"/litellm-limits-*.json +serve 500 +result=$("$collector") +[[ $(jq -c '[.usageStatusText, has("balance"), .limitsFetchedAt]' <<<"$result") == '["Balance unavailable",false,0]' ]] || + fail "LiteLLM collector says the balance is unavailable with nothing kept" "$result" +[[ $(jq -r '.authHelpText' <<<"$result") == *"status 500"* ]] || + fail "LiteLLM collector says what failed" "$result" +pass "LiteLLM collector says the balance is unavailable with nothing kept" + +# A redirect to another host must not carry the key along. +serve 200 '{"info": {"spend": 1, "max_budget": 2}}' +printf 'http://localhost:%s/key/info' "$PORT" >"$STATE/redirect" +result=$("$collector" --force) +rm "$STATE/redirect" +[[ $(jq -r '.authHelpText' <<<"$result") == *"redirected off the configured host"* ]] || + fail "LiteLLM collector refuses a redirect to another host" "$result" +! grep -q "^localhost" "$STATE/log" || + fail "LiteLLM collector never sends the key to another host" "$(cat "$STATE/log")" +pass "LiteLLM collector refuses a redirect to another host" + +# A base URL carrying credentials is refused before any request. +printf '{"baseUrl": "http://user:pass@127.0.0.1:%s", "apiKey": "config-key"}\n' "$PORT" >"$XDG_CONFIG_HOME/omarchy/agents/litellm.json" +before=$(requests) +result=$("$collector" --force) +[[ $(requests) == "$before" && $(jq -r '.usageStatusText' <<<"$result") == "Balance unavailable" ]] || + fail "LiteLLM collector refuses a base URL with credentials in it" "$result" +pass "LiteLLM collector refuses a base URL with credentials in it" +rm "$XDG_CONFIG_HOME/omarchy/agents/litellm.json" + +# ------------------------------------------------------------- local stats + +today=$(date +%Y-%m-%d) +sessions="$TEST_HOME/.pi/agent/sessions/--work--" +mkdir -p "$sessions" "$XDG_DATA_HOME/opencode" +line() { + printf '{"type":"message","id":"%s","timestamp":"%sT12:00:00.000Z","message":{"role":"assistant","provider":"%s","model":"anthropic/claude-sonnet-4","usage":{"input":%s,"output":10,"cacheRead":5,"cacheWrite":0}}}\n' "$1" "$today" "$2" "$3" +} +{ line a litellm 100; line b anthropic 999; line c litellm false; } >"$sessions/one.jsonl" +printf '*\n' >"$TEST_HOME/.gitignore" +git -C "$TEST_HOME" init -q + +python3 - "$XDG_DATA_HOME/opencode/opencode.db" <<'PY' +import json, sqlite3, sys, time +conn = sqlite3.connect(sys.argv[1]); now = int(time.time() * 1000) +conn.execute("CREATE TABLE message (id TEXT PRIMARY KEY, session_id TEXT, time_created INTEGER, time_updated INTEGER, data TEXT)") +def row(mid, provider, inp): + return (mid, "s1", now, now, json.dumps({"role": "assistant", "providerID": provider, "modelID": "gpt-5", "tokens": {"input": inp, "output": 1, "reasoning": 1, "cache": {"read": 0, "write": 0}}, "time": {"created": now}})) +conn.executemany("INSERT INTO message VALUES (?, ?, ?, ?, ?)", [row("m1", "litellm", 20), row("m2", "openai", 999), row("m3", "litellm-proxy", 999)]) +conn.execute("INSERT INTO message VALUES ('m4', 's1', 0, 0, 'not json \"providerID\": \"litellm\"')") +conn.commit(); conn.close() +PY -[[ $(jq -r '.legacyScanModel' <<<"$result") == "90" && $(jq -r '.legacyScanReady' <<<"$result") == "true" ]] || - fail "LiteLLM collector attaches model rows when the day series has none" "$result" -pass "LiteLLM collector attaches model rows when the day series has none" +result=$("$collector" --force) +# pi: a 115 and c 15 (its boolean input reads as nothing); opencode m1 22. +# Other providers, a lookalike id and a broken row add nothing. +[[ $(jq -c '[.ready, .todayTotalTokens, .todayPrompts, .todaySessions, (.modelUsage | keys), .usageStatusText]' <<<"$result") == '[true,152,3,2,["claude-sonnet-4","gpt-5"],"Waiting for auth"]' ]] || + fail "LiteLLM collector counts pi and opencode usage on the litellm provider" "$result" +pass "LiteLLM collector counts pi and opencode usage on the litellm provider, waiting for auth without a key" + +# A scan from seconds ago is reused; --force reads again. +line d litellm 1000 >>"$sessions/one.jsonl" +[[ $(jq -r '.todayTotalTokens' <<<"$("$collector")") == "152" && $(jq -r '.todayTotalTokens' <<<"$("$collector" --force)") == "1167" ]] || + fail "LiteLLM collector reuses a recent scan and --force rescans" "$("$collector")" +pass "LiteLLM collector reuses a recent scan and --force rescans" From 02be06c942d6ee32637d318f621ed34c246bbac5 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:21:03 -0400 Subject: [PATCH 095/546] Keep the Local collector the exact complement of the provider collectors Count a repeated Codex token_count once, as the Codex collector does, decide a forked rollout by its first session_meta, and keep per-file rollout totals so a rescan only reads what changed. Find pi and omp sessions with rg --no-ignore across profiles and PI_CODING_AGENT_DIR, deduplicated by message id and timestamp. Open opencode's database without immutable=1, which missed rows still in its write-ahead log, and count compactions and the V2 store. Leave out Ollama's cloud models and any provider id another collector counts, even when local.json names it. Name the record Local models, with no placeholder tier, and print ready: false on a machine with no local usage. Co-authored-by: Nate Ashby Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-local | 1131 +++++++++++------ .../shell.d/agent-usage-local-scanner-test.sh | 346 ++--- 2 files changed, 934 insertions(+), 543 deletions(-) diff --git a/bin/omarchy-agent-usage-local b/bin/omarchy-agent-usage-local index 44a96685626..2d34e9344f8 100755 --- a/bin/omarchy-agent-usage-local +++ b/bin/omarchy-agent-usage-local @@ -4,22 +4,23 @@ # omarchy:hidden=true """Collect self-hosted model usage into one display-ready JSON record. -Every other collector answers "what did this subscription cost me". This one -answers the question none of them can: how much work ran on hardware you own. -It counts the turns the subscription collectors deliberately skip — Codex CLI -pointed at Ollama, opencode on a local gateway, pi on llama.cpp — whether the -endpoint is this machine or a box down the hall. - -An endpoint is self-hosted when its URL resolves to a loopback, private, or -link-local address, or to a bare/LAN hostname (`server.home`, `gpu-box`). -Provider ids with no URL to resolve fall back to a list of well-known local -runtimes. Neither is exhaustive, so `~/.config/omarchy/agents/local.json` can -name more: +Every other collector counts what a subscription or a provider bills. This +one counts the turns that ran on hardware you own: Codex CLI pointed at +Ollama, opencode on a local gateway, pi on llama.cpp, whether the endpoint is +this machine or a box down the hall. It is the exact complement of the +provider collectors: a provider id another collector counts is never counted +here, and the Codex collector's own `openai` rollouts never are either. + +A provider is self-hosted when it is a well-known local runtime, an id named +in ~/.config/omarchy/agents/local.json, or a Codex model_providers entry whose +base_url is on a loopback, private, link-local or LAN address: {"providers": ["homelab", "my-gateway"]} -There are no limits and no balance to report: the meter on self-hosted work is -the electricity bill. +Ollama's cloud models (a `:cloud` or `-cloud` tag) run on Ollama's servers +even when the local daemon fronts them, so they are never counted here. There +are no limits and no balance: the meter on self-hosted work is the +electricity bill. The agents panel only ever reads the JSON this prints. """ import argparse @@ -27,350 +28,245 @@ import fcntl import hashlib import ipaddress import json +import math import os +import re +import shutil import sqlite3 +import subprocess import sys import tempfile import time import tomllib -from datetime import datetime, timedelta +from datetime import datetime, timedelta, timezone from pathlib import Path from urllib.parse import urlparse AGENT_ID = "local" -AGENT_NAME = "Local" -TIER_LABEL = "Self-hosted" +AGENT_NAME = "Local models" -# Runtimes that serve models from your own hardware. Used when a provider id -# carries no URL to classify — opencode and pi record the id but not the -# endpoint, and Codex's built-in `--oss` provider has no config.toml entry. +# Runtimes that serve models from your own hardware, for provider ids with no +# URL to classify: opencode and pi record the id but not the endpoint, and +# Codex's built-in `--oss` providers have no config.toml entry. KNOWN_LOCAL_PROVIDERS = { "gpt4all", "jan", "koboldcpp", "llama-cpp", "llamacpp", "llamafile", "lm-studio", "lmstudio", "local", "localai", "mlx", "ollama", "openai-local", "oss", "sglang", "tabby", "text-generation-webui", "vllm", } -# Session files are scanned back this far, matching the Codex collector's -# window. Pi transcripts carry no rollover, so they are read in full. +# Provider ids the provider collectors count, from pi, omp and opencode. A +# local.json naming one of them would count its tokens twice, so they are +# never local, whatever the configuration says. +CLAIMED_PROVIDERS = { + "anthropic", "deepseek", "devin", "litellm", "meta", "minimax", "minimax-coding", + "minimax-token-plan", "ollama-cloud", "openai", "openai-codex", "opencode", + "opencode-go", "openrouter", "synthetic", "zai", "zhipu", +} + +# The Codex collector counts rollouts on the built-in provider, and ones too +# old to name a provider at all. +CODEX_NATIVE_PROVIDER = "openai" + +# Codex rollouts are read back this far, matching the Codex collector. NATIVE_WINDOW_DAYS = 30 -# Same reuse windows as the subscription collectors: a no-flag run reuses a -# scan only long enough to dedup the concurrent collectors the update command -# backgrounds, while --limits-only promises nothing about stats freshness and -# may reuse one for far longer. This collector has no limits to fetch, so -# --limits-only is purely a cheap way to redraw. SCAN_REUSE_SECONDS = 20 LIMITS_ONLY_REUSE_SECONDS = 900 -today = datetime.now().strftime("%Y-%m-%d") +# Per-file totals of Codex rollouts survive between runs, so a rescan only +# reads the files that changed. Bumping it retires every older record. +FILE_CACHE_SCHEMA = 1 +# Past this a token count or timestamp from disk is garbage, not usage. +MAX_TOKENS = 10_000_000_000 +EARLIEST_EPOCH = 946_684_800 # 2000-01-01 +LATEST_EPOCH = 4_102_444_800 # 2100-01-01 + + +# --------------------------------------------------------------- utilities -def local_day(value): - if value is None: - return datetime.now().strftime("%Y-%m-%d") - if isinstance(value, (int, float)): - # pi message timestamps are milliseconds; Codex timestamps are 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 obj(value): + return value if isinstance(value, dict) else {} + +# Token counts and timestamps on disk are untrusted: a boolean, NaN, infinite, +# negative, huge or non-numeric value reads as nothing. def number(value): - try: - return max(0, int(value or 0)) - except Exception: + if isinstance(value, bool) or not isinstance(value, (int, float)): return 0 + if not math.isfinite(value) or value < 0 or value > MAX_TOKENS: + return 0 + return int(value) -def model_name(value): - text = str(value or "").strip() - return text or "local" +def finite(value): + if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value): + return None + return float(value) -def is_self_hosted(url): - """True when a base URL points at hardware the user controls.""" - try: - host = urlparse(str(url or "")).hostname - except Exception: - return False - if not host: - return False - host = host.lower().rstrip(".") - if host == "localhost": - return True +def epoch_seconds(value): + if finite(value) is None: + return None + # Session timestamps are seconds or milliseconds. + seconds = value / 1000 if value > 100_000_000_000 else value + return seconds if EARLIEST_EPOCH <= seconds < LATEST_EPOCH else None + + +def local_day(value): + seconds = epoch_seconds(value) + if seconds is not None: + return datetime.fromtimestamp(seconds).strftime("%Y-%m-%d") + if not isinstance(value, str) or not value.strip(): + return None try: - ip = ipaddress.ip_address(host) - return ip.is_loopback or ip.is_private or ip.is_link_local + parsed = datetime.fromisoformat(value.strip().replace("Z", "+00:00")) except ValueError: - pass - # A bare hostname is a LAN name by construction; the rest are the suffixes - # home routers and mDNS hand out. - if "." not in host: - return True - return host.endswith((".home", ".lan", ".local", ".internal", ".localdomain", ".home.arpa")) + return None + if epoch_seconds(parsed.timestamp()) is None: + return None + return parsed.astimezone().strftime("%Y-%m-%d") if parsed.tzinfo else parsed.strftime("%Y-%m-%d") -def configured_providers(): - path = Path(os.environ.get("XDG_CONFIG_HOME") or (Path.home() / ".config")) / "omarchy" / "agents" / "local.json" +def model_name(raw): + return str(raw or "").strip() or AGENT_ID + + +def runtime_env(): + home = str(Path.home()) + path_parts = [ + os.environ.get("PATH", ""), + f"{home}/.local/bin", + f"{home}/.npm-global/bin", + f"{home}/.local/share/mise/shims", + ] + env = os.environ.copy() + env["PATH"] = os.pathsep.join(part for part in path_parts if part) + return env + + +ENV = runtime_env() + + +def find_command(name): + return shutil.which(name, path=ENV.get("PATH")) + + +def data_home(): + return Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) + + +def cache_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 read_json(path): try: - data = json.loads(path.read_text()) + return json.loads(path.read_text(encoding="utf-8")) except Exception: - return set() - names = data.get("providers") if isinstance(data, dict) else None - if not isinstance(names, list): - return set() - return {str(name).strip().lower() for name in names if str(name).strip()} + return None -def codex_endpoints(): - """provider id -> base_url, from the Codex CLI's own config.""" - codex_home = Path(os.environ.get("CODEX_HOME") or (Path.home() / ".codex")) +def read_fresh_json(path, max_age_seconds): + if max_age_seconds <= 0: + return None try: - with (codex_home / "config.toml").open("rb") as handle: - config = tomllib.load(handle) - except Exception: - return {} - providers = config.get("model_providers") - if not isinstance(providers, dict): - return {} - endpoints = {} - for name, entry in providers.items(): - if isinstance(entry, dict): - endpoints[str(name).lower()] = str(entry.get("base_url") or "") - return endpoints + # 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 + except OSError: + return None + return read_json(path) if 0 <= age < max_age_seconds else None -class Scan: +def write_json(path, payload): + # A temp name unique to this writer: 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 + + +# ------------------------------------------------------------------- stats + + +def empty_bucket(): + return {"inputTokens": 0, "outputTokens": 0, "cacheReadInputTokens": 0, "cacheCreationInputTokens": 0} + + +class Stats: + """Token and activity totals shared by every local source. Sessions and days + are sets, so a day active in two sources counts once toward activeDays.""" + def __init__(self): - self.extra_providers = configured_providers() - self.endpoints = codex_endpoints() - self.today = today - self.recent_days = [ - (datetime.now() - timedelta(days=offset)).strftime("%Y-%m-%d") for offset in range(6, -1, -1) - ] - self.recent = {day: 0 for day in self.recent_days} + 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() - self.today_tokens = 0 - self.today_by_model = {} + self.today_tokens_by_model = {} self.total_prompts = 0 self.total_sessions = set() self.active_days = set() self.model_usage = {} - self.providers_seen = set() - self.endpoints_seen = set() - self.seen_pi_messages = set() - self.complete = True - def provider_is_local(self, provider): - pid = str(provider or "").strip().lower() - if not pid: - return False - if pid in self.extra_providers or pid in KNOWN_LOCAL_PROVIDERS: - return True - url = self.endpoints.get(pid) - return bool(url) and is_self_hosted(url) - - def local_provider_ids(self): - """Every id this scan would accept, for pushing filters down into SQL.""" - ids = set(KNOWN_LOCAL_PROVIDERS) | set(self.extra_providers) - for pid, url in self.endpoints.items(): - if url and is_self_hosted(url): - ids.add(pid) - return ids - - def note_provider(self, provider): - pid = str(provider or "").strip().lower() - self.providers_seen.add(pid) - url = self.endpoints.get(pid) - if url: - host = urlparse(url).hostname - if host: - self.endpoints_seen.add(host.lower()) - - def add(self, day, session_key, model, input_tokens, output_tokens, cache_read, cache_write): + def add(self, session, day, model, input_tokens, output_tokens, cache_read, cache_write): total = input_tokens + output_tokens + cache_read + cache_write if total <= 0: return self.total_prompts += 1 - self.total_sessions.add(session_key) self.active_days.add(day) - if day in self.recent: - self.recent[day] += total - if day == self.today: - self.today_prompts += 1 - self.today_sessions.add(session_key) - self.today_tokens += total - self.today_by_model[model] = self.today_by_model.get(model, 0) + total - bucket = self.model_usage.setdefault( - model, - {"inputTokens": 0, "outputTokens": 0, "cacheReadInputTokens": 0, "cacheCreationInputTokens": 0}, - ) + if session is not None: + self.total_sessions.add(session) + bucket = self.model_usage.setdefault(model, empty_bucket()) bucket["inputTokens"] += input_tokens bucket["outputTokens"] += output_tokens bucket["cacheReadInputTokens"] += cache_read bucket["cacheCreationInputTokens"] += cache_write - - # ------------------------------------------------------------ native codex - def scan_codex(self): - codex_home = Path(os.environ.get("CODEX_HOME") or (Path.home() / ".codex")) - cutoff = time.time() - NATIVE_WINDOW_DAYS * 24 * 60 * 60 - for root in (codex_home / "sessions", codex_home / "archived_sessions"): - if not root.exists(): - continue - for path in root.rglob("*.jsonl"): - try: - if path.stat().st_mtime < cutoff: - continue - except OSError: - continue - self._scan_codex_file(path) - - def _scan_codex_file(self, path): - current_model = "local" - try: - with path.open(errors="replace") as handle: - for raw in handle: - try: - entry = json.loads(raw) - except Exception: - continue - if entry.get("type") == "session_meta": - provider = (entry.get("payload") or {}).get("model_provider") - # The subscription collector counts these; this one must not. - if not self.provider_is_local(provider): - return - self.note_provider(provider) - continue - if entry.get("type") == "turn_context": - payload = entry.get("payload") or {} - current_model = model_name(payload.get("model") or payload.get("model_slug") or current_model) - continue - payload = entry.get("payload") or entry - if entry.get("type") == "response_item" and isinstance(payload, dict): - payload = payload.get("payload") or payload - if not isinstance(payload, dict) or payload.get("type") != "token_count": - continue - usage = (payload.get("info") or {}).get("last_token_usage") or {} - cache_read = number(usage.get("cached_input_tokens")) - cache_write = number(usage.get("cache_write_input_tokens")) - input_tokens = max(0, number(usage.get("input_tokens")) - cache_read - cache_write) - output_tokens = number(usage.get("output_tokens")) - day = local_day(entry.get("timestamp") or path.stat().st_mtime) - self.add(day, str(path), current_model, input_tokens, output_tokens, cache_read, cache_write) - except Exception: + if day in self.recent: + # recentDays.messageCount holds tokens, despite the legacy name. + self.recent[day] += total + if day == self.today: + self.today_prompts += 1 + if session is not None: + self.today_sessions.add(session) + self.today_total_tokens += total + self.today_tokens_by_model[model] = self.today_tokens_by_model.get(model, 0) + total + + # A Codex rollout's totals arrive per day and model, prompts counted. + def add_totals(self, session, day, model, input_tokens, output_tokens, cache_read, cache_write, prompts): + if input_tokens + output_tokens + cache_read + cache_write <= 0: return + self.add(session, day, model, input_tokens, output_tokens, cache_read, cache_write) + self.total_prompts += prompts - 1 + if day == self.today: + self.today_prompts += prompts - 1 - # ----------------------------------------------------------------- pi / omp - def scan_pi(self): - for root in (Path.home() / ".pi" / "agent" / "sessions", Path.home() / ".omp" / "agent" / "sessions"): - if not root.exists(): - continue - for path in root.rglob("*.jsonl"): - try: - with path.open(errors="replace") as handle: - for raw in handle: - try: - entry = json.loads(raw) - except Exception: - continue - message = entry.get("message") or {} - if not isinstance(message, dict) or message.get("role") != "assistant": - continue - provider = message.get("provider") or message.get("api") - if not self.provider_is_local(provider): - continue - key = str(entry.get("id") or entry.get("uuid") or f"{path}:{raw[:64]}") - if key in self.seen_pi_messages: - continue - self.seen_pi_messages.add(key) - self.note_provider(provider) - usage = message.get("usage") or {} - day = local_day(entry.get("timestamp") or message.get("timestamp")) - self.add( - day, str(path), model_name(message.get("model")), - number(usage.get("input")), number(usage.get("output")), - number(usage.get("cacheRead")), number(usage.get("cacheWrite")), - ) - except Exception: - continue + def add_session(self, session, today): + self.total_sessions.add(session) + if today: + self.today_sessions.add(session) - # ----------------------------------------------------------------- opencode - def scan_opencode(self): - db = Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) / "opencode" / "opencode.db" - if not db.exists(): - return - try: - conn = sqlite3.connect(f"file:{db}?mode=ro&immutable=1", uri=True) - except Exception: - return - # OpenCode DBs hold every historical message with its JSON, so the - # provider filter belongs in SQL rather than in a Python loop over every - # row. The candidate ids are known before the query runs: the well-known - # runtimes, whatever the user configured, and any Codex provider whose - # endpoint resolves to hardware they host. json_valid() guards the - # extract so one malformed row cannot abort the scan. - candidates = sorted(self.local_provider_ids()) - if not candidates: - return - placeholders = ",".join("?" * len(candidates)) - try: - for session_id, raw in conn.execute( - "SELECT session_id, data FROM message" - " WHERE data LIKE '%\"role\"%:%\"assistant\"%'" - " 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" - f" IN ({placeholders})", - candidates, - ): - try: - entry = json.loads(raw) - except Exception: - continue - # Exact match again in Python: the SQL narrowed the rows, it did not - # vet them. - if not isinstance(entry, dict) or not self.provider_is_local(entry.get("providerID")): - continue - self.note_provider(entry.get("providerID")) - tokens = entry.get("tokens") or {} - cache = tokens.get("cache") or {} - created = (entry.get("time") or {}).get("created") - # opencode keeps thinking tokens out of output; both are generated. - self.add( - local_day(created), "opencode:" + str(session_id), model_name(entry.get("modelID")), - number(tokens.get("input")), number(tokens.get("output")) + number(tokens.get("reasoning")), - number(cache.get("read")), number(cache.get("write")), - ) - except Exception: - # A scan cut short by a database error must not be cached as the whole - # story, or the missing usage stays suppressed until the cache expires. - self.complete = False - return - finally: - try: - conn.close() - except Exception: - pass - - def stats(self): - """Snapshot the aggregated local usage into the record's stats dict.""" + def record(self): return { "todayPrompts": self.today_prompts, "todaySessions": len(self.today_sessions), - "todayTotalTokens": self.today_tokens, - "todayTokensByModel": self.today_by_model, - "recentDays": [{"date": day, "messageCount": self.recent[day]} for day in self.recent_days], + "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), @@ -379,155 +275,596 @@ class Scan: } -def cache_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 +# --------------------------------------------------------------- pi and omp +# --------------------------------------------------------------- pi and omp -def scan_cache_paths(): - codex_home = Path(os.environ.get("CODEX_HOME") or (Path.home() / ".codex")) - db = Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) / "opencode" / "opencode.db" - config = Path(os.environ.get("XDG_CONFIG_HOME") or (Path.home() / ".config")) / "omarchy" / "agents" / "local.json" - # The digest covers every path the scan reads: the codex session roots and - # its config.toml, the opencode DB, the provider config, and (via - # Path.home()) the pi/omp session roots. - digest = hashlib.sha1( - "\n".join(str(part) for part in (Path.home(), codex_home, db, config)).encode("utf-8") - ).hexdigest()[:16] - root = cache_root() - return root / f"local-scan-{digest}.json", root / f"local-scan-{digest}.lock" +def pi_agent_dirs(): + pi_agent = os.environ.get("PI_CODING_AGENT_DIR") + dirs = [Path(pi_agent).expanduser() if pi_agent else Path.home() / ".pi" / "agent", Path.home() / ".omp" / "agent"] + # `omp --profile=` (and pi's equivalent) moves the agent tree under + # /profiles//. + for base in (Path.home() / ".pi", Path.home() / ".omp"): + try: + dirs.extend(sorted(child / "agent" for child in (base / "profiles").iterdir() if child.is_dir())) + except OSError: + pass + return dirs -def read_fresh_json(path, max_age_seconds): - if max_age_seconds <= 0 or not path.exists(): - return None +def pi_session_roots(): + roots = [] + for agent_dir in pi_agent_dirs(): + root = agent_dir / "sessions" + if root.is_dir() and root.resolve() not in [r.resolve() for r in roots]: + roots.append(root) + return roots + + +def pi_parent_session(path): + # A fork names the session it was cut from in its header line. 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: + with open(path, encoding="utf-8", errors="replace") as handle: + header = json.loads(handle.readline()) + except (OSError, ValueError): return None - return None + parent = header.get("parentSession") if isinstance(header, dict) else None + return os.path.abspath(os.path.expanduser(parent)) if isinstance(parent, str) and parent else None -def write_json(path, payload): - # 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) +def scan_pi_sessions(stats): + """Add the pi/omp assistant messages this collector counts. Returns whether + every root was read.""" + roots = pi_session_roots() + if not roots: + return True + rg = find_command("rg") or "rg" try: - with os.fdopen(handle_fd, "w", encoding="utf-8") as handle: - handle.write(json.dumps(payload, separators=(",", ":")) + "\n") - # mkstemp opens at 0600; nothing in the cache is sensitive, so open it - # up to the usual 0644. - tmp.chmod(0o644) - tmp.replace(path) - except BaseException: - tmp.unlink(missing_ok=True) - raise + proc = subprocess.Popen( + # --no-ignore: these are data files, not a source tree. A $HOME that is + # a git checkout would otherwise hide every session behind its ignores. + [rg, "--json", "--no-ignore", "-e", pi_pattern(), *map(str, roots)], + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + text=True, + errors="replace", + env=ENV, + ) + except OSError: + return False + counted = set() + owned = {} + today_owned = {} + assert proc.stdout is not None + for raw in proc.stdout: + try: + event = json.loads(raw) + if event.get("type") != "match": + continue + path = os.path.abspath(event["data"]["path"]["text"]) + line_number = event["data"].get("line_number") + entry = json.loads(event["data"]["lines"]["text"]) + except Exception: + continue + if not isinstance(entry, dict) or entry.get("type") != "message": + continue + message = obj(entry.get("message")) + if message.get("role") != "assistant" or not accept_pi(message): + continue + usage = obj(message.get("usage")) + 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 not (input_tokens or output_tokens or cache_read or cache_write): + input_tokens = number(usage.get("totalTokens")) + timestamp = entry.get("timestamp") + day = local_day(timestamp) or local_day(message.get("timestamp")) + if day is None or not (input_tokens or output_tokens or cache_read or cache_write): + continue + # A fork repeats its parent's messages with their ids and timestamps, so + # the pair says two lines are one message. An id alone is only unique + # within its file, and a line without one always counts. + message_id = str(entry.get("id") or "") + if message_id and timestamp is not None: + key = ("entry", message_id, str(timestamp)) + elif message_id: + key = ("path", path, message_id) + else: + key = ("line", path, line_number) + owned.setdefault(path, set()).add(key) + if day == stats.today: + today_owned.setdefault(path, set()).add(key) + if key in counted: + continue + counted.add(key) + stats.add(None, day, model_name(message.get("model")), input_tokens, output_tokens, cache_read, cache_write) + + try: + code = proc.wait(timeout=5) + except subprocess.TimeoutExpired: + proc.kill() + return False + + # A message belongs to the oldest session holding it, so a fork is a + # session of its own only for what it added, whichever file rg read first. + parents = {} + for path, keys in owned.items(): + inherited = set() + visited = {path} + parent = parents.setdefault(path, pi_parent_session(path)) + while parent and parent not in visited: + visited.add(parent) + inherited |= owned.get(parent, set()) + parent = parents.setdefault(parent, pi_parent_session(parent)) + if keys - inherited: + stats.add_session(path, bool(today_owned.get(path, set()) - inherited)) + # rg exits 1 when nothing matched and 2 when a file couldn't be read. + return code in (0, 1) + + +# ---------------------------------------------------------------- opencode +# --------------------------------------------------------- classification -# The cache payload is a versioned envelope around the stats dict, so a -# corrupted or foreign-shaped file is a cache miss (rescan + rewrite) instead -# of a crash or a garbage record. -def read_cached_stats(cache_file, max_age_seconds): - 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, - # not merely old, whatever its mtime says. - 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 config_path(): + return Path(os.environ.get("XDG_CONFIG_HOME") or (Path.home() / ".config")) / "omarchy" / "agents" / "local.json" -def write_cached_stats(cache_file, stats): + +def configured_providers(): + names = obj(read_json(config_path())).get("providers") + if not isinstance(names, list): + return set() + return {str(name).strip().lower() for name in names if isinstance(name, str) and name.strip()} + + +def codex_home(): + return Path(os.environ.get("CODEX_HOME") or (Path.home() / ".codex")) + + +def codex_endpoints(): + """provider id -> base_url, from the Codex CLI's own config.""" try: - write_json(cache_file, {"schemaVersion": 1, "scanDate": today, "stats": stats}) - except Exception as exc: - print(f"omarchy-agent-usage-local: could not write usage cache ({exc})", file=sys.stderr) + with (codex_home() / "config.toml").open("rb") as handle: + config = tomllib.load(handle) + except (OSError, tomllib.TOMLDecodeError): + return {} + providers = config.get("model_providers") + if not isinstance(providers, dict): + return {} + return {str(name).lower(): str(entry.get("base_url") or "") for name, entry in providers.items() if isinstance(entry, dict)} + + +def is_self_hosted(url): + """True when a base URL points at hardware the user controls.""" + try: + host = urlparse(str(url or "")).hostname + except ValueError: + return False + if not host: + return False + host = host.lower().rstrip(".") + if host == "localhost": + return True + try: + ip = ipaddress.ip_address(host) + return ip.is_loopback or ip.is_private or ip.is_link_local + except ValueError: + pass + # A bare hostname is a LAN name by construction; the rest are the suffixes + # home routers and mDNS hand out. + if "." not in host: + return True + return host.endswith((".home", ".lan", ".local", ".internal", ".localdomain", ".home.arpa")) + + +def is_cloud_model(model): + # Ollama names its hosted models with a cloud tag ("gpt-oss:120b-cloud", + # "glm-4.6:cloud"); the local daemon only forwards them, so Ollama's + # servers did the work. + tag = str(model or "").lower().rpartition(":")[2] if ":" in str(model or "") else "" + return tag == "cloud" or tag.endswith("-cloud") + + +class Classifier: + def __init__(self): + self.configured = configured_providers() + self.endpoints = codex_endpoints() + ids = set(KNOWN_LOCAL_PROVIDERS) | self.configured + ids |= {pid for pid, url in self.endpoints.items() if url and is_self_hosted(url)} + self.local_ids = ids - CLAIMED_PROVIDERS + + def provider_is_local(self, provider): + return str(provider or "").strip().lower() in self.local_ids + + def accepts(self, provider, model): + return self.provider_is_local(provider) and not is_cloud_model(model) + + def fingerprint(self): + return sorted(self.local_ids) + + +CLASSIFIER = None -def run_local_scan(): - scan = Scan() - scan.scan_codex() - scan.scan_pi() - scan.scan_opencode() - return scan.stats(), scan.complete +def classifier(): + global CLASSIFIER + if CLASSIFIER is None: + CLASSIFIER = Classifier() + return CLASSIFIER + + +def pi_pattern(): + ids = "|".join(re.escape(pid) for pid in sorted(classifier().local_ids)) + return r'(?i)"provider"\s*:\s*"(?:' + ids + ')"' + + +def accept_pi(message): + # The Codex collector also counts any message on its openai-codex API, + # whatever the provider id says. + if str(message.get("api") or "").startswith("openai-codex"): + return False + return classifier().accepts(message.get("provider"), message.get("model")) + + +# ------------------------------------------------------------------- Codex + + +def file_cache_zone(): + # Records hold local days, so they are only true in the timezone that wrote + # them. + return [os.environ.get("TZ"), list(time.tzname), time.timezone, time.altzone] + + +def file_cache_path(): + return cache_root() / f"local-codex-files-{scan_digest()}.json" + + +def valid_file_record(value): + try: + counts = [totals for models in value["days"].values() for totals in models.values()] + return ( + type(value["mtime"]) in (int, float) and type(value["size"]) is int and isinstance(value["provider"], str) + and all(len(totals) == 5 and all(type(n) is int for n in totals) for totals in counts) + ) + except Exception: + return False -def cached_stats(max_age): - """Stats, with the cache as a pure optimization. +def read_file_cache(): + """Per-file totals of Codex rollouts from earlier scans, keyed by path. - The cache must never take the collector down: any cache-layer failure - (unwritable cache root, lock errors, disk full) degrades to a direct scan - and a warning on stderr. The JSON record is the contract; the cache is not. + Rollouts are append-only, so a file whose mtime and size both match its + record cannot have grown a turn since. A record keeps the rollout's + provider and every turn whatever that provider is, so a change to which + providers count as local needs no rereading. """ + payload = obj(read_json(file_cache_path())) + if payload.get("schemaVersion") != FILE_CACHE_SCHEMA or payload.get("zone") != file_cache_zone(): + return {} + files = obj(payload.get("files")) + return {path: record for path, record in files.items() if valid_file_record(record)} + + +def write_file_cache(files): + # Only the files this scan saw are written back, so deleted rollouts and + # ones past the window leave the cache. try: - return _cached_stats(max_age) - except Exception as exc: - print(f"omarchy-agent-usage-local: cache unavailable ({exc}); scanning directly", file=sys.stderr) - stats, _ = run_local_scan() - return stats + write_json(file_cache_path(), {"schemaVersion": FILE_CACHE_SCHEMA, "zone": file_cache_zone(), "files": files}) + except OSError as exc: + print(f"omarchy-agent-usage-local: could not write file cache ({exc})", file=sys.stderr) + + +def read_codex_rollout(path, stat_result): + """One rollout's provider and day -> model -> [input, output, cacheRead, + cacheWrite, prompts], counted the way the Codex collector counts its own.""" + record = {"mtime": stat_result.st_mtime, "size": stat_result.st_size, "provider": "", "days": {}} + current_model = "local" + seen_meta = False + previous_total_usage = None + try: + with path.open(errors="replace") as handle: + for raw in handle: + if '"token_count"' not in raw and '"turn_context"' not in raw and '"session_meta"' not in raw: + continue + try: + entry = json.loads(raw) + except ValueError: + continue + if not isinstance(entry, dict): + continue + if entry.get("type") == "session_meta": + # A forked rollout copies its parent's session_meta after its own, + # so only the first one, as Codex reads it, says who served this. + if seen_meta: + continue + seen_meta = True + record["provider"] = str(obj(entry.get("payload")).get("model_provider") or "").strip().lower() + # The built-in provider, or none recorded, is the Codex + # collector's: nothing here can count it. + if record["provider"] in ("", CODEX_NATIVE_PROVIDER): + break + continue + if not seen_meta: + continue + if entry.get("type") == "turn_context": + payload = obj(entry.get("payload")) + current_model = str(payload.get("model") or payload.get("model_slug") or current_model) + continue + payload = entry.get("payload") or entry + if entry.get("type") == "response_item" and isinstance(payload, dict): + payload = payload.get("payload") or payload + if not isinstance(payload, dict) or payload.get("type") != "token_count": + continue + info = obj(payload.get("info")) + # total_token_usage is cumulative for the session; the last turn is + # what this one spent. + usage = obj(info.get("last_token_usage")) + cache_read = number(usage.get("cached_input_tokens")) + cache_write = number(usage.get("cache_write_input_tokens")) + # Cached tokens are inside input_tokens, and reasoning inside + # output_tokens. + input_tokens = max(0, number(usage.get("input_tokens")) - cache_read - cache_write) + output_tokens = number(usage.get("output_tokens")) + if not (input_tokens or output_tokens or cache_read or cache_write): + continue + # Quota updates can repeat the previous request's last_token_usage. + # Only an unchanged cumulative snapshot proves this is a repeat: + # separate requests may spend identical counts. + total_usage = info.get("total_token_usage") + if isinstance(total_usage, dict) and total_usage: + if total_usage == previous_total_usage: + continue + previous_total_usage = total_usage + day = local_day(entry.get("timestamp")) or local_day(stat_result.st_mtime) + if day is None: + continue + totals = record["days"].setdefault(day, {}).setdefault(current_model, [0, 0, 0, 0, 0]) + totals[0] += input_tokens + totals[1] += output_tokens + totals[2] += cache_read + totals[3] += cache_write + totals[4] += 1 + except OSError: + # Keep what was read, but a partial record must not be cached as the + # file's whole story. + record["mtime"] = None + return record + + +def scan_codex(stats, use_file_cache): + cached_files = read_file_cache() if use_file_cache else {} + scanned_files = {} + cutoff = time.time() - NATIVE_WINDOW_DAYS * 24 * 60 * 60 + complete = True + for root in (codex_home() / "sessions", codex_home() / "archived_sessions"): + if not root.is_dir(): + continue + for path in root.rglob("*.jsonl"): + try: + stat_result = path.stat() + except OSError: + continue + if stat_result.st_mtime < cutoff: + continue + key = str(path) + record = cached_files.get(key) + if record is None or record["mtime"] != stat_result.st_mtime or record["size"] != stat_result.st_size: + record = read_codex_rollout(path, stat_result) + if record["mtime"] is None: + complete = False + else: + scanned_files[key] = record + if not classifier().provider_is_local(record["provider"]): + continue + for day, models in record["days"].items(): + for model, (input_tokens, output_tokens, cache_read, cache_write, prompts) in models.items(): + if not is_cloud_model(model): + stats.add_totals(key, day, model, input_tokens, output_tokens, cache_read, cache_write, prompts) + write_file_cache(scanned_files) + return complete + + +# ---------------------------------------------------------------- opencode + + +def opencode_db(): + # OPENCODE_DB points opencode at another database, as the CLI itself reads + # it; a relative name lives beside the default one. Only a build whose + # channel was literally "v2" wrote opencode-v2.db. + root = data_home() / "opencode" + configured = os.environ.get("OPENCODE_DB", "").strip() + if configured: + path = Path(configured).expanduser() + return path if path.is_absolute() else root / path + normal = root / "opencode.db" + legacy = root / "opencode-v2.db" + return legacy if not normal.is_file() and legacy.is_file() else normal + + +def table_columns(conn, table): + return {row[1] for row in conn.execute(f"PRAGMA table_info({table})")} + + +def opencode_row(stats, seen, message_id, session_id, entry, row_type, created): + # The original `message` store keeps the provider and model at the top + # level; the V2 `session_message` store nests them under `model` and keeps + # the role in its `type` column, with a compaction's tokens spent like any + # reply's. + if not isinstance(entry, dict) or (entry.get("role") or row_type) not in ("assistant", "compaction"): + return + model = entry.get("model") + provider = entry.get("providerID") or (model.get("providerID") if isinstance(model, dict) else "") + model_id = entry.get("modelID") or (model.get("id") or model.get("modelID") if isinstance(model, dict) else model) + if not classifier().accepts(provider, model_id): + return + tokens = obj(entry.get("tokens")) + cache = obj(tokens.get("cache")) + 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): + return + # A migrated message can sit in both stores under the same id. + if message_id in seen: + return + seen.add(message_id) + day = local_day(obj(entry.get("time")).get("created")) or local_day(created) + if day is None: + return + stats.add("opencode:" + str(session_id), day, model_name(model_id), input_tokens, output_tokens, cache_read, cache_write) + + +def scan_opencode(stats): + """Add opencode messages on a local provider, from both of its stores. + Returns whether the database was read to the end.""" + db = opencode_db() + ids = sorted(classifier().local_ids) + if not db.is_file() or not ids: + return True + try: + # Read-only, and not immutable: opencode writes through a WAL, and an + # immutable open would miss what hasn't been checkpointed yet. + conn = sqlite3.connect(db.resolve().as_uri() + "?mode=ro", uri=True, timeout=2) + except sqlite3.Error: + return False + # The LIKE gates skip rows whose JSON cannot name a local provider before + # any parse; the classifier stays the authority. LIKE ignores ASCII case, + # like the ids. + gates = ["%\"providerID\"%:%\"" + pid + "\"%" for pid in ids] + seen = set() + try: + conn.execute("PRAGMA query_only = ON") + tables = {name for (name,) in conn.execute("SELECT name FROM sqlite_master WHERE type = 'table'")} + if "session_message" in tables and {"id", "session_id", "type", "data"} <= table_columns(conn, "session_message"): + columns = table_columns(conn, "session_message") + created = "m.time_created" if "time_created" in columns else "NULL" + source, fork_filter = "session_message m", "" + if "time_created" in columns and "session_v2" in tables and {"id", "fork_session_id", "time_created"} <= table_columns(conn, "session_v2"): + # A fork copies its source session's messages under new ids; only + # what came after the fork was spent in it. + source += " LEFT JOIN session_v2 s ON s.id = m.session_id" + fork_filter = " AND (s.fork_session_id IS NULL OR m.time_created >= s.time_created)" + gate = " OR ".join(["m.data LIKE ?"] * len(gates)) + for message_id, session_id, row_type, time_created, raw in conn.execute( + f"SELECT m.id, m.session_id, m.type, {created}, m.data FROM {source}" + f" WHERE m.type IN ('assistant', 'compaction') AND ({gate})" + fork_filter, gates + ): + try: + opencode_row(stats, seen, str(message_id), session_id, json.loads(raw), str(row_type or ""), time_created) + except (TypeError, ValueError): + continue + if "message" in tables and {"session_id", "data"} <= table_columns(conn, "message"): + message_id = "id" if "id" in table_columns(conn, "message") else "rowid" + gate = " OR ".join(["data LIKE ?"] * len(gates)) + for row_id, session_id, raw in conn.execute(f"SELECT {message_id}, session_id, data FROM message WHERE {gate}", gates): + try: + opencode_row(stats, seen, str(row_id), session_id, json.loads(raw), "", None) + except (TypeError, ValueError): + continue + except sqlite3.Error: + # Transient lock, schema migration, corruption: the numbers stop here, + # incomplete. + return False + finally: + conn.close() + return True + + +# -------------------------------------------------------------- scan cache -def _cached_stats(max_age): - cache_file, lock_file = scan_cache_paths() +def scan_digest(): + # Every path the scan reads, and which providers it counts as local, so a + # changed local.json or Codex config is a new scan rather than a stale one. + parts = [str(Path.home()), str(codex_home()), str(opencode_db()), os.environ.get("PI_CODING_AGENT_DIR", ""), json.dumps(classifier().fingerprint())] + return hashlib.sha1("\n".join(parts).encode("utf-8")).hexdigest()[:16] - cached = read_cached_stats(cache_file, max_age) - if cached is not None: - return cached - with lock_file.open("w") as lock: - fcntl.flock(lock, fcntl.LOCK_EX) +def scan_cache_paths(): + root = cache_root() + digest = scan_digest() + return root / f"local-scan-{digest}.json", root / f"local-scan-{digest}.lock" + + +def run_local_scans(max_age): + stats = Stats() + complete = scan_codex(stats, use_file_cache=max_age > 0) + complete = scan_pi_sessions(stats) and complete + complete = scan_opencode(stats) and complete + return stats.record(), complete + + +def read_cached_stats(cache_file, max_age): + cached = read_fresh_json(cache_file, max_age) + if not isinstance(cached, dict) or cached.get("schemaVersion") != 2: + return None + # today* fields only mean "today" on the day they were scanned. + if cached.get("scanDate") != datetime.now().strftime("%Y-%m-%d"): + return None + stats = cached.get("stats") + return stats if isinstance(stats, dict) and set(Stats().record()) <= set(stats) else None + + +def cached_local_stats(max_age): + """Local stats, with the cache as a pure optimization: any cache-layer + failure degrades to a direct scan and a warning on stderr.""" + try: + cache_file, lock_file = scan_cache_paths() cached = read_cached_stats(cache_file, max_age) if cached is not None: return cached - stats, complete = run_local_scan() - # An interrupted scan still serves this run, but caching it would - # suppress the missing usage for every reader until the cache expires. - if complete: - write_cached_stats(cache_file, stats) - return stats + with lock_file.open("w") as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + cached = read_cached_stats(cache_file, max_age) + if cached is not None: + return cached + stats, complete = run_local_scans(max_age) + # A scan that couldn't read everything serves this run, but caching it + # would hide the missing usage from every reader until it expires. + if complete: + write_json(cache_file, {"schemaVersion": 2, "scanDate": datetime.now().strftime("%Y-%m-%d"), "stats": stats}) + return stats + except Exception as exc: + print(f"omarchy-agent-usage-local: cache unavailable ({exc}); scanning directly", file=sys.stderr) + return run_local_scans(0)[0] + + +# ------------------------------------------------------------------ record def main(): - parser = argparse.ArgumentParser() - # Accepted for parity with the other collectors: the update command passes - # them to everything it runs. There are no limits to fetch here, so - # --limits-only just means "a redraw may reuse an older scan". - parser.add_argument("--force", action="store_true") - parser.add_argument("--limits-only", action="store_true") + parser = argparse.ArgumentParser(description="Print the self-hosted model usage record as JSON") + # There are no limits to fetch, so --limits-only only means a redraw may + # reuse an older scan. + parser.add_argument("--force", action="store_true", help="rescan every session, ignoring caches") + parser.add_argument("--limits-only", action="store_true", help="reuse a recent scan") args = parser.parse_args() max_age = 0 if args.force else (LIMITS_ONLY_REUSE_SECONDS if args.limits_only else SCAN_REUSE_SECONDS) - + stats = cached_local_stats(max_age) record = { "schemaVersion": 1, "id": AGENT_ID, "name": AGENT_NAME, - "updatedAt": datetime.now().astimezone().isoformat(), - "ready": True, + "updatedAt": datetime.now(timezone.utc).isoformat(), + "ready": number(stats.get("totalPrompts")) > 0, "hasLocalStats": True, - "tierLabel": TIER_LABEL, + # Self-hosted work has no plan. + "tierLabel": "", "usageStatusText": "", "authHelpText": "", - # Nothing meters self-hosted work, so the panel draws no limit rows. "limits": [], + "limitsStale": False, + "limitsFetchedAt": 0, } - record.update(cached_stats(max_age)) - print(json.dumps(record, separators=(",", ":"))) + record.update(stats) + print(json.dumps(record, separators=(",", ":"), sort_keys=True)) + return 0 if __name__ == "__main__": - main() + raise SystemExit(main()) diff --git a/test/shell.d/agent-usage-local-scanner-test.sh b/test/shell.d/agent-usage-local-scanner-test.sh index 3c6940f47e5..f38009be2db 100755 --- a/test/shell.d/agent-usage-local-scanner-test.sh +++ b/test/shell.d/agent-usage-local-scanner-test.sh @@ -4,12 +4,54 @@ source "$(dirname "$0")/base-test.sh" require_command jq require_command python3 +require_command rg + +SCRATCH=$(mktemp -d) +trap 'rm -rf "$SCRATCH"' EXIT + +collector="$ROOT/bin/omarchy-agent-usage-local" +unset PI_CODING_AGENT_DIR OPENCODE_DB + +# A fixture home with everything the collector reads inside it. +use_home() { + export HOME="$1" + export CODEX_HOME="$1/.codex" + export XDG_CACHE_HOME="$1/.cache" + export XDG_DATA_HOME="$1/.local/share" + export XDG_CONFIG_HOME="$1/.config" + mkdir -p "$CODEX_HOME/sessions/$(date +%Y/%m/%d)" "$XDG_CONFIG_HOME/omarchy/agents" "$XDG_DATA_HOME/opencode" +} + +timestamp="$(date +%Y-%m-%d)T12:00:00Z" +rollouts="$(date +%Y/%m/%d)" + +# rollout [...]: a Codex rollout +# whose turns each spend 1000 input and 100 output tokens. +token_count() { + local total=$1 + printf '{"timestamp":"%s","type":"event_msg","payload":{"type":"token_count","info":{"last_token_usage":{"input_tokens":1000,"output_tokens":100,"cached_input_tokens":0},"total_token_usage":{"input_tokens":%s}}}}\n' "$timestamp" "$total" +} +rollout() { + local name=$1 provider=$2 model=$3 + { + printf '{"timestamp":"%s","type":"session_meta","payload":{"model_provider":"%s"}}\n' "$timestamp" "$provider" + printf '{"timestamp":"%s","type":"turn_context","payload":{"model":"%s"}}\n' "$timestamp" "$model" + token_count 1000 + } >"$CODEX_HOME/sessions/$rollouts/rollout-$name.jsonl" +} + +# ------------------------------------------------------------ unused machine + +use_home "$SCRATCH/empty" +result=$("$collector") +[[ $(jq -c '[.ready, .name, .tierLabel, .limits, .usageStatusText, .authHelpText, .todayTotalTokens, (.recentDays | length), .modelUsage, has("balance")]' <<<"$result") == '[false,"Local models","",[],"","",0,7,{},false]' ]] || + fail "Local collector prints an empty, not-ready record on an unused machine" "$result" +pass "Local collector prints an empty, not-ready record on an unused machine" -LOCAL_HOME=$(mktemp -d) -trap 'rm -rf "$LOCAL_HOME"' EXIT -mkdir -p "$LOCAL_HOME/.codex/sessions/$(date +%Y/%m/%d)" +# ------------------------------------------------------------------ Codex -cat >"$LOCAL_HOME/.codex/config.toml" <<'EOF' +use_home "$SCRATCH/codex" +cat >"$CODEX_HOME/config.toml" <<'EOF' [model_providers.homelab] base_url = "http://server.home:11434/v1" @@ -17,51 +59,32 @@ base_url = "http://server.home:11434/v1" base_url = "https://openrouter.ai/api/v1" EOF -timestamp="$(date +%Y-%m-%d)T12:00:00Z" -rollout() { - local name=$1 provider=$2 model=$3 - cat >"$LOCAL_HOME/.codex/sessions/$(date +%Y/%m/%d)/rollout-$name.jsonl" <>"$CODEX_HOME/sessions/$rollouts/rollout-ollama.jsonl" +# A fork copies its parent's session_meta after its own; the first decides. +{ + printf '{"timestamp":"%s","type":"session_meta","payload":{"model_provider":"openai"}}\n' "$timestamp" + printf '{"timestamp":"%s","type":"session_meta","payload":{"model_provider":"ollama"}}\n' "$timestamp" + printf '{"timestamp":"%s","type":"turn_context","payload":{"model":"qwen-forked"}}\n' "$timestamp" + token_count 1000 +} >"$CODEX_HOME/sessions/$rollouts/rollout-fork.jsonl" + +local_result=$("$collector") +[[ $(jq -c '[.modelUsage | keys[]]' <<<"$local_result") == '["qwen-here","qwen-net"]' ]] || fail "Local collector counts LAN endpoints and known runtimes, nothing else" "$local_result" -pass "Local collector counts self-hosted turns from this machine and the network" +[[ $(jq -c '[.todayTotalTokens, .modelUsage["qwen-here"].inputTokens, .todayPrompts, .ready]' <<<"$local_result") == "[3300,2000,3,true]" ]] || + fail "Local collector counts a repeated Codex token_count once" "$local_result" +pass "Local collector counts self-hosted Codex turns once, leaving Ollama Cloud and the subscription out" -[[ $(jq -r '.limits | length' <<<"$local_result") == "0" && $(jq -r '.tierLabel' <<<"$local_result") == "Self-hosted" ]] || - fail "Local collector reports no limits" "$local_result" -pass "Local collector reports no limits for unmetered work" - -# A custom gateway the heuristics cannot see is named in config. -mkdir -p "$LOCAL_HOME/.config/omarchy/agents" -echo '{"providers":["mystery"]}' >"$LOCAL_HOME/.config/omarchy/agents/local.json" -rollout mystery mystery qwen-custom -# --force: the config changed inside the scan-reuse window, so a no-flag run -# would legitimately serve the previous scan. -local_result=$(env "${env_args[@]}" "$ROOT/bin/omarchy-agent-usage-local" --force) -[[ $(jq -r '.modelUsage["qwen-custom"].inputTokens' <<<"$local_result") == "1000" ]] || - fail "Local collector honors configured provider ids" "$local_result" -pass "Local collector honors configured provider ids" - -# The local and Codex collectors partition the same sessions directory: every -# rollout is counted by exactly one of them, so the fix cannot lose a turn and -# the new tab cannot double-count one. -cat >"$LOCAL_HOME/codex" <<'EOF' +# The Local and Codex collectors partition the sessions directory: the +# subscription's turns are Codex's alone, and no model shows in both. +cat >"$SCRATCH/codex/codex" <<'EOF' #!/bin/bash while read -r request; do id=$(jq -r '.id // empty' <<<"$request") @@ -72,43 +95,100 @@ while read -r request; do esac done EOF -chmod +x "$LOCAL_HOME/codex" - -codex_result=$(env "${env_args[@]}" PATH="$LOCAL_HOME:$PATH" "$ROOT/bin/omarchy-agent-usage-codex" --force) +chmod +x "$SCRATCH/codex/codex" +codex_result=$(PATH="$SCRATCH/codex:$PATH" "$ROOT/bin/omarchy-agent-usage-codex" --force) +[[ $(jq -c '[.modelUsage | keys[]]' <<<"$codex_result") == '["gpt-sub","qwen-forked"]' ]] || + fail "Codex collector keeps the subscription's own turns" "$codex_result" +overlap=$(jq -n --argjson a "$(jq -c '.modelUsage | keys' <<<"$local_result")" --argjson b "$(jq -c '.modelUsage | keys' <<<"$codex_result")" '[$a[] | select(. as $m | $b | index($m))] | length') +[[ $overlap == "0" ]] || + fail "Local and Codex collectors never count the same turn" "$local_result$codex_result" +pass "Local and Codex collectors partition rollouts without overlap" + +# A reused rollout record is replayed, not reread: change a file's contents +# but not its size or mtime, and the old count stands. +file="$CODEX_HOME/sessions/$rollouts/rollout-homelab.jsonl" +touch -r "$file" "$SCRATCH/stamp" +sed -i 's/"output_tokens":100/"output_tokens":900/' "$file" +touch -r "$SCRATCH/stamp" "$file" +[[ $(jq -r '.modelUsage["qwen-net"].outputTokens' <<<"$("$collector" --limits-only)") == "100" ]] || + fail "Local collector replays an unchanged rollout from its file cache" "$("$collector" --limits-only)" +rm -f "$XDG_CACHE_HOME"/omarchy/agent-usage/local-scan-*.json +[[ $(jq -r '.modelUsage["qwen-net"].outputTokens' <<<"$("$collector")") == "100" ]] || + fail "Local collector replays an unchanged rollout after the scan cache expires" "$("$collector")" +[[ $(jq -r '.modelUsage["qwen-net"].outputTokens' <<<"$("$collector" --force)") == "900" ]] || + fail "Local collector --force rereads every rollout" "$("$collector" --force)" +pass "Local collector keeps a per-file Codex cache, and --force rereads" + +# A provider the heuristics can't see is named in local.json, and takes +# effect without --force, since the classification is part of the scan key. +echo '{"providers":["mystery"]}' >"$XDG_CONFIG_HOME/omarchy/agents/local.json" +rollout mystery mystery qwen-custom +[[ $(jq -r '.modelUsage["qwen-custom"].inputTokens' <<<"$("$collector")") == "1000" ]] || + fail "Local collector honors configured provider ids" "$("$collector")" +pass "Local collector honors configured provider ids" -[[ $(jq -r '.modelUsage | keys | join(",")' <<<"$codex_result") == "gpt-sub" ]] || - fail "Codex collector keeps only the subscription's own turns" "$codex_result" +# A provider collector's id is never local, whatever local.json says. +echo '{"providers":["openrouter","openai"]}' >"$XDG_CONFIG_HOME/omarchy/agents/local.json" +[[ $(jq -c '[.modelUsage | keys[]]' <<<"$("$collector")") == '["qwen-here","qwen-net"]' ]] || + fail "Local collector never counts a provider another collector counts" "$("$collector")" +pass "Local collector never counts a provider another collector counts" -overlap=$(jq -n --argjson a "$(jq -c '.modelUsage | keys' <<<"$local_result")" \ - --argjson b "$(jq -c '.modelUsage | keys' <<<"$codex_result")" \ - '$a - ($a - $b) | length') -[[ $overlap == "0" ]] || - fail "Local and Codex collectors must not count the same model twice" "$local_result$codex_result" -pass "Local and Codex collectors partition sessions without gaps or overlap" - -# opencode records the provider id but not the endpoint, so the well-known -# runtimes and configured ids carry the filter - pushed into SQL, since these -# databases keep every historical message with its JSON. -OC_HOME=$(mktemp -d) -trap 'rm -rf "$LOCAL_HOME" "$OC_HOME"' EXIT -mkdir -p "$OC_HOME/.config/omarchy/agents" -echo '{"providers":["homelab"]}' >"$OC_HOME/.config/omarchy/agents/local.json" - -python3 - "$OC_HOME/.local/share/opencode/opencode.db" <<'PY' -import json, sqlite3, sys, time -from pathlib import Path +# ------------------------------------------------------------- pi and omp -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) +use_home "$SCRATCH/pi" +sessions="$HOME/.pi/agent/sessions/--work--" +mkdir -p "$sessions" "$HOME/.omp/profiles/box/agent/sessions" +line() { + printf '{"type":"message","id":"%s","timestamp":"%s","message":{"role":"assistant","provider":"%s","model":"%s","usage":{"input":%s,"output":10,"cacheRead":0,"cacheWrite":0}}}\n' "$1" "$timestamp" "$2" "$3" "$4" +} +{ + line a ollama qwen3:30b 100 + line b LMStudio qwen3:8b 200 + line c ollama glm-4.6:cloud 999 + line d ollama-cloud qwen3:30b 999 + line e anthropic claude 999 + # The Codex collector counts its API whatever the provider id says. + printf '{"type":"message","id":"x","timestamp":"%s","message":{"role":"assistant","provider":"ollama","api":"openai-codex-responses","model":"gpt-5","usage":{"input":999}}}\n' "$timestamp" +} >"$sessions/one.jsonl" +{ + printf '{"type":"session","parentSession":"%s"}\n' "$sessions/one.jsonl" + line a ollama qwen3:30b 100 + line f vllm llama-70b 300 +} >"$HOME/.omp/profiles/box/agent/sessions/fork.jsonl" +printf '*\n' >"$HOME/.gitignore" +git -C "$HOME" init -q + +result=$("$collector") +# a 110, b 210, f 310; the fork's copy of a counts once, Ollama Cloud models +# and the ollama-cloud provider belong to the ollama collector, and Anthropic +# to claude. +[[ $(jq -c '[.todayTotalTokens, .todayPrompts, .todaySessions, ([.modelUsage | keys[]])]' <<<"$result") == '[630,3,2,["llama-70b","qwen3:30b","qwen3:8b"]]' ]] || + fail "Local collector counts pi and omp sessions on local providers once" "$result" +pass "Local collector counts pi and omp sessions on local providers once, profiles included" + +moved="$HOME/moved" +mkdir -p "$moved/sessions" +line g llamacpp gemma 40 >"$moved/sessions/moved.jsonl" +[[ $(jq -r '.modelUsage.gemma.inputTokens' <<<"$(PI_CODING_AGENT_DIR="$moved" "$collector" --force)") == "40" ]] || + fail "Local collector follows PI_CODING_AGENT_DIR" "$(PI_CODING_AGENT_DIR="$moved" "$collector" --force)" +pass "Local collector follows PI_CODING_AGENT_DIR" + +# ---------------------------------------------------------------- opencode + +use_home "$SCRATCH/opencode" +echo '{"providers":["homelab"]}' >"$XDG_CONFIG_HOME/omarchy/agents/local.json" +{ python3 - "$XDG_DATA_HOME/opencode/opencode.db" || true; } 2>/dev/null <<'PY' +import json, os, sqlite3, sys, time +conn = sqlite3.connect(sys.argv[1]); now = int(time.time() * 1000) +conn.execute("PRAGMA journal_mode=WAL") +conn.execute("CREATE TABLE message (id TEXT PRIMARY KEY, session_id TEXT NOT NULL, time_created INTEGER, time_updated INTEGER, data TEXT NOT NULL)") +conn.execute("CREATE TABLE session_message (id TEXT PRIMARY KEY, session_id TEXT, type TEXT, time_created INTEGER, data TEXT)") def message(mid, provider, model, role="assistant", inp=0, out=0, reasoning=0): - return (mid, "ses_1", now_ms, now_ms, json.dumps({ + return (mid, "ses_1", now, now, json.dumps({ "role": role, "providerID": provider, "modelID": model, "tokens": {"input": inp, "output": out, "reasoning": reasoning, "cache": {"read": 0, "write": 0}}, - "time": {"created": now_ms}, + "time": {"created": now}, })) conn.executemany("INSERT INTO message VALUES (?, ?, ?, ?, ?)", [ @@ -117,83 +197,57 @@ conn.executemany("INSERT INTO message VALUES (?, ?, ?, ?, ?)", [ message("m3", "anthropic", "claude-paid", inp=999, out=999), message("m4", "openai", "gpt-sub", inp=999, out=999), message("m5", "ollama", "qwen-oc", role="user", inp=999), + message("m6", "ollama", "kimi-k2:1t-cloud", inp=999), ]) -conn.execute("INSERT INTO message VALUES ('m6', 'ses_1', ?, ?, 'not json at all')", (now_ms, now_ms)) -conn.commit(); conn.close() +conn.execute("INSERT INTO message VALUES ('m7', 'ses_1', ?, ?, 'not json \"providerID\":\"ollama\"')", (now, now)) +conn.execute("INSERT INTO session_message VALUES ('c1', 'ses_2', 'compaction', ?, ?)", (now, json.dumps({ + "model": {"providerID": "ollama", "id": "qwen-oc"}, "tokens": {"input": 30, "output": 0, "cache": {"read": 0, "write": 0}}, "time": {"created": now}}))) +conn.commit() +# Killed with the rows still in the write-ahead log: an immutable open would +# not see them. +os.kill(os.getpid(), 9) PY -oc_env=(HOME="$OC_HOME" CODEX_HOME="$OC_HOME/.codex" XDG_CACHE_HOME="$OC_HOME/.cache" - XDG_DATA_HOME="$OC_HOME/.local/share" XDG_CONFIG_HOME="$OC_HOME/.config") -oc=$(env "${oc_env[@]}" "$ROOT/bin/omarchy-agent-usage-local") - -# 115 = 100 + 10 + 5 reasoning; 55 = 50 + 5. Paid providers and the user -# message contribute nothing, and the malformed row does not abort the scan. -[[ $(jq -r '.todayTotalTokens' <<<"$oc") == "170" ]] || - fail "Local collector counts opencode turns on self-hosted providers" "$oc" -[[ $(jq -r '.modelUsage | keys | sort | join(",")' <<<"$oc") == "qwen-cfg,qwen-oc" ]] || - fail "Local collector filters opencode rows to self-hosted providers" "$oc" -[[ $(jq -r '.modelUsage["qwen-oc"].outputTokens' <<<"$oc") == "15" ]] || - fail "Local collector counts opencode reasoning tokens as generated" "$oc" -pass "Local collector counts self-hosted opencode turns past paid and malformed rows" - -# The scan is cached like every other collector's: one envelope per data-path -# set, readable, versioned, and bypassed by --force. -cache_file=$(ls "$OC_HOME/.cache/omarchy/agent-usage/"local-scan-*.json 2>/dev/null | head -n 1) -[[ -n $cache_file && $(stat -c %a "$cache_file") == "644" ]] || - fail "Local collector leaves a readable cache file behind" "$oc" -[[ $(jq -r '.schemaVersion' "$cache_file") == "1" && $(jq -r '.stats.todayTotalTokens' "$cache_file") == "170" ]] || - fail "Local collector writes a versioned cache envelope" "$oc" -pass "Local collector caches its scan" - -python3 - "$OC_HOME/.local/share/opencode/opencode.db" <<'PY' +oc=$("$collector") +# 115 = 100 + 10 + 5 reasoning; 55 = 50 + 5; the compaction's 30. Paid +# providers, a user message, a cloud model and a malformed row add nothing. +[[ $(jq -c '[.todayTotalTokens, ([.modelUsage | keys[]]), .modelUsage["qwen-oc"].outputTokens, .todaySessions]' <<<"$oc") == '[200,["qwen-cfg","qwen-oc"],15,2]' ]] || + fail "Local collector counts self-hosted opencode turns, compactions included, from a live database" "$oc" +pass "Local collector counts self-hosted opencode turns, compactions included, from a live database" + +cache_file=$(ls "$XDG_CACHE_HOME"/omarchy/agent-usage/local-scan-*.json | head -n 1) +[[ $(stat -c %a "$cache_file") == "644" && $(jq -r '.stats.todayTotalTokens' "$cache_file") == "200" ]] || + fail "Local collector writes a readable scan cache" "$(cat "$cache_file")" + +python3 - "$XDG_DATA_HOME/opencode/opencode.db" <<'PY' import json, sqlite3, sys, time -conn = sqlite3.connect(sys.argv[1]); now_ms = int(time.time() * 1000) -conn.execute("INSERT INTO message VALUES (?, ?, ?, ?, ?)", ("m7", "ses_1", now_ms, now_ms, json.dumps({ +conn = sqlite3.connect(sys.argv[1]); now = int(time.time() * 1000) +conn.execute("INSERT INTO message VALUES ('m8', 'ses_1', ?, ?, ?)", (now, now, json.dumps({ "role": "assistant", "providerID": "ollama", "modelID": "qwen-oc", - "tokens": {"input": 30, "output": 0, "reasoning": 0, "cache": {"read": 0, "write": 0}}, - "time": {"created": now_ms}}))) + "tokens": {"input": 30, "output": 0, "reasoning": 0, "cache": {"read": 0, "write": 0}}, "time": {"created": now}}))) conn.commit(); conn.close() PY - -[[ $(jq -r '.todayTotalTokens' <<<"$(env "${oc_env[@]}" "$ROOT/bin/omarchy-agent-usage-local")") == "170" ]] || - fail "Local collector reuses a seconds-old scan" "$oc" -[[ $(jq -r '.todayTotalTokens' <<<"$(env "${oc_env[@]}" "$ROOT/bin/omarchy-agent-usage-local" --force)") == "200" ]] || - fail "Local collector --force rescans past the cache" "$oc" -pass "Local collector --force rescans past the cache" - -# A cache stamped with another local date holds another day's today* stats. -jq -c '.scanDate = "1999-01-01"' "$cache_file" >"$cache_file.tmp" && mv "$cache_file.tmp" "$cache_file" -[[ $(jq -r '.scanDate' <<<"$(cat "$cache_file")") == "1999-01-01" ]] && - [[ $(jq -r '.todayTotalTokens' <<<"$(env "${oc_env[@]}" "$ROOT/bin/omarchy-agent-usage-local" --limits-only)") == "200" ]] || - fail "Local collector treats a cache from another day as a miss" "$oc" -pass "Local collector treats a cache from another day as a miss" - -# A corrupt-but-parseable cache is a miss, not a garbage record. +[[ $(jq -r '.todayTotalTokens' <<<"$("$collector")") == "200" ]] || + fail "Local collector reuses a seconds-old scan" "$("$collector")" +[[ $(jq -r '.todayTotalTokens' <<<"$("$collector" --force)") == "230" ]] || + fail "Local collector --force rescans past the cache" "$("$collector" --force)" +pass "Local collector reuses a recent scan and --force rescans" + +jq -c '.scanDate = "1999-01-01" | .stats.todayTotalTokens = 1' "$cache_file" >"$cache_file.tmp" && mv "$cache_file.tmp" "$cache_file" +[[ $(jq -r '.todayTotalTokens' <<<"$("$collector" --limits-only)") == "230" ]] || + fail "Local collector treats a scan from another day as a miss" "$("$collector" --limits-only)" printf '[]' >"$cache_file" -[[ $(jq -r '.todayTotalTokens' <<<"$(env "${oc_env[@]}" "$ROOT/bin/omarchy-agent-usage-local" --limits-only)") == "200" ]] || - fail "Local collector recovers from a corrupt cache file" "$oc" -pass "Local collector recovers from a corrupt cache file" - -# An unwritable cache must not kill the collector: the record is the contract. -touch "$OC_HOME/not-a-dir" -broken=$(env "${oc_env[@]}" XDG_CACHE_HOME="$OC_HOME/not-a-dir" "$ROOT/bin/omarchy-agent-usage-local" 2>/dev/null) -[[ $(jq -r '.todayTotalTokens' <<<"$broken") == "200" ]] || - fail "Local collector still prints a record when the cache is unwritable" "$broken" -pass "Local collector survives an unwritable cache" - -# A scan cut short by a database error must not be cached as the whole story. -BROKEN_HOME=$(mktemp -d) -trap 'rm -rf "$LOCAL_HOME" "$OC_HOME" "$BROKEN_HOME"' EXIT -python3 - "$BROKEN_HOME/.local/share/opencode/opencode.db" <<'PY' -import sqlite3, sys -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 unrelated (id text PRIMARY KEY)") -conn.commit(); conn.close() -PY -env HOME="$BROKEN_HOME" CODEX_HOME="$BROKEN_HOME/.codex" XDG_CACHE_HOME="$BROKEN_HOME/.cache" \ - XDG_DATA_HOME="$BROKEN_HOME/.local/share" XDG_CONFIG_HOME="$BROKEN_HOME/.config" \ - "$ROOT/bin/omarchy-agent-usage-local" >/dev/null -[[ -z $(ls "$BROKEN_HOME/.cache/omarchy/agent-usage/"local-scan-*.json 2>/dev/null) ]] || - fail "Local collector must not cache an interrupted scan" "broken-db" +[[ $(jq -r '.todayTotalTokens' <<<"$("$collector" --limits-only)") == "230" ]] || + fail "Local collector recovers from a corrupt cache file" "$("$collector" --limits-only)" +touch "$SCRATCH/not-a-dir" +[[ $(jq -r '.todayTotalTokens' <<<"$(XDG_CACHE_HOME="$SCRATCH/not-a-dir" "$collector" 2>/dev/null)") == "230" ]] || + fail "Local collector still prints a record when the cache is unwritable" "unwritable" +pass "Local collector treats stale, corrupt and unwritable caches as misses" + +# A scan cut short by a database error is not cached as the whole story. +use_home "$SCRATCH/broken" +echo "not a database" >"$XDG_DATA_HOME/opencode/opencode.db" +"$collector" >/dev/null +[[ -z $(ls "$XDG_CACHE_HOME"/omarchy/agent-usage/local-scan-*.json 2>/dev/null) ]] || + fail "Local collector does not cache an interrupted opencode scan" "broken-db" pass "Local collector does not cache an interrupted opencode scan" From d319edd1943367c688526653e3e8072bd4dd9dfb Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:02:36 -0400 Subject: [PATCH 096/546] Add Qwen Code as a selectable default coding agent Qwen Code (qwen) joins the default agent picker, installed through mise as npm:@qwen-code/qwen-code with a lazy stub on fresh installs and a migration for existing ones. The launcher runs it with --yolo and seeds a prompt with --prompt-interactive, bound to the flag so a leading dash stays literal. From #11067. Co-authored-by: Caden DeNike Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent | 5 ++ bin/omarchy-default-agent | 5 +- bin/omarchy-remove-preinstalls | 6 ++ default/fonts/omarchy/README.md | 1 + default/fonts/omarchy/omarchy.ttf | Bin 9972 -> 10196 bytes default/i18n/en/menu.json | 3 + default/omarchy/omarchy-menu.jsonc | 1 + install/user/mise.sh | 1 + manual/17-ai.md | 1 + migrations/1792032538.sh | 5 ++ test/shell.d/default-agent-test.sh | 20 ++++++- test/shell.d/menu-test.sh | 5 +- test/shell.d/mise-provision-test.sh | 4 +- test/shell.d/qwen-agent-test.sh | 85 ++++++++++++++++++++++++++++ 14 files changed, 135 insertions(+), 7 deletions(-) create mode 100644 migrations/1792032538.sh create mode 100644 test/shell.d/qwen-agent-test.sh diff --git a/bin/omarchy-agent b/bin/omarchy-agent index 6d0e585d498..59cb198e743 100755 --- a/bin/omarchy-agent +++ b/bin/omarchy-agent @@ -138,6 +138,11 @@ kimi) command=(kimi --yolo) fi ;; +qwen) + # Bind the value so a prompt starting with a dash stays literal. + command=(qwen --yolo) + [[ -n ${prompt:-} ]] && command+=("--prompt-interactive=$prompt") + ;; *) echo "Unsupported default agent: $agent" >&2 exit 1 diff --git a/bin/omarchy-default-agent b/bin/omarchy-default-agent index 7675424ba80..519200a6d23 100755 --- a/bin/omarchy-default-agent +++ b/bin/omarchy-default-agent @@ -1,7 +1,7 @@ #!/bin/bash # omarchy:summary=Set and launch the default coding agent -# omarchy:args=[pi|omp|opencode|ori|claude|codex|grok|openclaw|agy|hermes|copilot|crush|cursor-agent|muse|kimi] +# omarchy:args=[pi|omp|opencode|ori|claude|codex|grok|openclaw|agy|hermes|copilot|crush|cursor-agent|muse|kimi|qwen] # omarchy:examples=omarchy default agent | omarchy default agent codex | omarchy default agent claude installing=false @@ -43,8 +43,9 @@ muse | muse-code | musecode) ;; cursor | cursor-agent) agent="cursor-agent"; name="Cursor CLI" ;; kimi | kimi-code) agent="kimi"; name="Kimi"; agent_package="npm:@moonshot-ai/kimi-code" ;; +qwen | qwen-code) agent="qwen"; name="Qwen Code"; agent_package="npm:@qwen-code/qwen-code" ;; *) - echo "Usage: omarchy-default-agent " + echo "Usage: omarchy-default-agent " exit 1 ;; esac diff --git a/bin/omarchy-remove-preinstalls b/bin/omarchy-remove-preinstalls index c764b9ddfa4..f0568427d22 100755 --- a/bin/omarchy-remove-preinstalls +++ b/bin/omarchy-remove-preinstalls @@ -36,6 +36,12 @@ if gum confirm --affirmative "$(omarchy-i18n common.yes)" --negative "$(omarchy- rm -f ~/.local/bin/muse fi + # Preserve a user-managed Qwen Code launcher at the same path. + if [[ -f ~/.local/bin/qwen && ! -L ~/.local/bin/qwen ]] && + grep -Eq '^mise use -g .*"npm:@qwen-code/qwen-code"' ~/.local/bin/qwen; then + rm -f ~/.local/bin/qwen + fi + omarchy-pkg-drop \ aether \ cliamp \ diff --git a/default/fonts/omarchy/README.md b/default/fonts/omarchy/README.md index df93aa555ae..544fbda5bf6 100644 --- a/default/fonts/omarchy/README.md +++ b/default/fonts/omarchy/README.md @@ -17,6 +17,7 @@ The private-use glyphs in `omarchy.ttf` are: - `U+E90C` — OpenClaw, traced from the lobster mascot the openclaw package ships as `dist/control-ui/favicon.svg`, since upstream publishes no monochrome SVG - `U+E90D` — Cursor, from - `U+E90E` — Claude, from +- `U+E90F` — Qwen, from The agent marks are monochrome so the menu can render them using the active theme's foreground and selection colors. diff --git a/default/fonts/omarchy/omarchy.ttf b/default/fonts/omarchy/omarchy.ttf index 3c70d8031f15323654523a51fe288297c945edd7..429a3e4cc342b8f40adfd18bcd2f20d86d639163 100644 GIT binary patch delta 456 zcmez3d&Pf(VSTQkI|Bn_2Ll6>S9(rm+Br>TEd~bWD+~-`4H>D4DSGRy`WP6bLKqkr z%rbxi9E>cIKzcVR3%EHR*#;WSdZ0zE~;>?HHY}t77O4xX+ zc-YKe@Bn#i<|S;{1QY&-Iy2s}{U>2-TT)(L^`*-81>={Jw(@e@FSd-+xU#rTaxpS; z#dBwJbH$(JM&mQumSkpS{+)R4+`p)-%O}sB%Vb=Z1#~|HBM2KCOQogXnA{`j&nPkZ tj${lA3j>4gWD6++E^Y^?8i~z$Qt!Ee(s62fOpJn)d(=`{SqjTj^8j!6Za4q{ delta 229 zcmccO|HXHLVSSFE8v_Gl2Ll6>S9(rmTAE$6J_7^u6$S>;H5sXiDSGRy`WP4_8yFZE z%rbxi9E>ayK>iLOUnL{Aq{4=QRRzeu1LVi#R8rmN3oV34u^@(pqmD+(BmnGAsZ z3?N@2FEKaOxcg2mkPozvdt*UzF_Q(P4cHh!QsSGvq~CJ`MXl8J LCTFRqG6Tf`*&H~U diff --git a/default/i18n/en/menu.json b/default/i18n/en/menu.json index 2ac7625d1d8..54d0ab3393c 100644 --- a/default/i18n/en/menu.json +++ b/default/i18n/en/menu.json @@ -527,6 +527,9 @@ "setup.default.agent.pi": { "label": "Pi" }, + "setup.default.agent.qwen": { + "label": "Qwen Code" + }, "setup.default.browser": { "label": "Browser", "title": "Default Browser" diff --git a/default/omarchy/omarchy-menu.jsonc b/default/omarchy/omarchy-menu.jsonc index 2a0a7bb2dc1..b2d4149e16a 100644 --- a/default/omarchy/omarchy-menu.jsonc +++ b/default/omarchy/omarchy-menu.jsonc @@ -205,6 +205,7 @@ "setup.default.agent.opencode": {"icon":"","iconFont":"omarchy","label":"OpenCode","checked":"[[ \"$(omarchy-default-agent)\" == \"opencode\" ]]","action":"omarchy-default-agent opencode"}, "setup.default.agent.ori": {"icon":"","iconFont":"omarchy","label":"Ori","checked":"[[ \"$(omarchy-default-agent)\" == \"ori\" ]]","action":"omarchy-default-agent ori"}, "setup.default.agent.pi": {"icon":"","iconFont":"omarchy","label":"Pi","checked":"[[ \"$(omarchy-default-agent)\" == \"pi\" ]]","action":"omarchy-default-agent pi"}, + "setup.default.agent.qwen": {"icon":"","iconFont":"omarchy","label":"Qwen Code","checked":"[[ \"$(omarchy-default-agent)\" == \"qwen\" ]]","action":"omarchy-default-agent qwen"}, "setup.default.browser": {"icon":"","label":"Browser","title":"Default Browser"}, "setup.default.browser.chromium": {"icon":"","label":"Chromium","checked":"[[ \"$(omarchy-default-browser)\" == \"chromium\" ]]","action":"omarchy-default-browser chromium"}, "setup.default.browser.chrome": {"icon":"󰊯","label":"Chrome","checked":"[[ \"$(omarchy-default-browser)\" == \"chrome\" ]]","action":"omarchy-default-browser chrome"}, diff --git a/install/user/mise.sh b/install/user/mise.sh index 7923c8227d9..0bf535ac85c 100644 --- a/install/user/mise.sh +++ b/install/user/mise.sh @@ -14,6 +14,7 @@ omarchy-mise-install pi omarchy-mise-install github:can1357/oh-my-pi omp omarchy-mise-install grok omarchy-mise-install npm:@moonshot-ai/kimi-code kimi +omarchy-mise-install npm:@qwen-code/qwen-code qwen omarchy-mise-install cursor-agent omarchy-mise-install npm:@kitlangton/ghui ghui omarchy-mise-install aqua:modem-dev/hunk hunk diff --git a/manual/17-ai.md b/manual/17-ai.md index 61de87494f8..ed33dea4e23 100644 --- a/manual/17-ai.md +++ b/manual/17-ai.md @@ -12,6 +12,7 @@ Omarchy treats AI coding agents as first-class citizens, but it doesn't pick a f | `crush` | [Crush](https://github.com/charmbracelet/crush) | | `grok` | [Grok Build](https://x.ai/build) from xAI | | `kimi` | [Kimi Code](https://www.kimi.com/code/docs/en/) | +| `qwen` | [Qwen Code](https://github.com/QwenLM/qwen-code), Alibaba's coding agent | | `pi` | [Mario Zechner's Pi](https://github.com/badlogic/pi-mono) | | `omp` | [Oh My Pi](https://github.com/can1357/oh-my-pi) | | `ori` | [Ori](https://openrouter.ai/docs/guides/ori/harness), OpenRouter's harness | diff --git a/migrations/1792032538.sh b/migrations/1792032538.sh new file mode 100644 index 00000000000..e5f1fe3baf5 --- /dev/null +++ b/migrations/1792032538.sh @@ -0,0 +1,5 @@ +echo "Install Qwen Code coding agent mise wrapper" + +if [[ ! -f $HOME/.local/state/omarchy/preinstalls-removed ]] && omarchy-cmd-missing qwen; then + omarchy-mise-install npm:@qwen-code/qwen-code qwen +fi diff --git a/test/shell.d/default-agent-test.sh b/test/shell.d/default-agent-test.sh index d3bdc792b9b..1e74d5f1e2d 100644 --- a/test/shell.d/default-agent-test.sh +++ b/test/shell.d/default-agent-test.sh @@ -132,6 +132,7 @@ crush_package="crush" agy_package="antigravity-cli" ori_package="github:OpenRouterLabs/ori-releases" kimi_package="npm:@moonshot-ai/kimi-code" +qwen_package="npm:@qwen-code/qwen-code" cursor_agent_package="cursor-agent" muse_package="http:muse[url=https://api.meta.ai/muse-launcher.sh,bin=muse,version_list_url=https://api.meta.ai/muse-code/channels/muse-stable,version_json_path=.version]" @@ -153,6 +154,7 @@ assert_lazy_stub "$omp_package" omp assert_lazy_stub "$crush_package" crush assert_lazy_stub "$ori_package" ori assert_lazy_stub "$kimi_package" kimi +assert_lazy_stub "$qwen_package" qwen assert_lazy_stub "$cursor_agent_package" cursor-agent assert_lazy_stub "$muse_package" muse pass "custom agent lazy stubs preserve their mise packages" @@ -166,6 +168,7 @@ grep -Fx "$crush_package" "$stub_log" >/dev/null || fail "user setup creates the grep -Fx "$ori_package ori" "$stub_log" >/dev/null || fail "user setup creates the Ori lazy stub" grep -Fx "$muse_package muse" "$stub_log" >/dev/null || fail "user setup creates the Muse lazy stub" grep -Fx "$kimi_package kimi" "$stub_log" >/dev/null || fail "user setup creates the Kimi lazy stub" +grep -Fx "$qwen_package qwen" "$stub_log" >/dev/null || fail "user setup creates the Qwen Code lazy stub" pass "user setup creates the custom agent lazy stubs" : >"$stub_log" @@ -173,6 +176,7 @@ source "$ROOT/install/user/mise.sh" grep -Fx "$cursor_agent_package" "$stub_log" >/dev/null || fail "user setup recreates the Cursor CLI lazy stub" grep -Fx "$muse_package muse" "$stub_log" >/dev/null || fail "user setup recreates the Muse lazy stub" grep -Fx "$kimi_package kimi" "$stub_log" >/dev/null || fail "user setup recreates the Kimi lazy stub" +grep -Fx "$qwen_package qwen" "$stub_log" >/dev/null || fail "user setup recreates the Qwen Code lazy stub" pass "user setup recreates wrappers even when the commands are present" : >"$stub_log" @@ -193,6 +197,10 @@ pass "Muse migration preserves existing installs and the preinstall opt-out" OMARCHY_TEST_MISSING_COMMAND=kimi source "$ROOT/migrations/1786967403.sh" >/dev/null grep -Fx "$kimi_package kimi" "$stub_log" >/dev/null || fail "Kimi migration creates its lazy stub" +: >"$stub_log" +OMARCHY_TEST_MISSING_COMMAND=qwen source "$ROOT/migrations/1792032538.sh" >/dev/null +grep -Fx "$qwen_package qwen" "$stub_log" >/dev/null || fail "Qwen Code migration creates its lazy stub" + : >"$stub_log" source "$ROOT/migrations/1785617047.sh" >/dev/null grep -Fx "$omp_package omp" "$stub_log" >/dev/null || fail "Oh My Pi migration creates a working lazy stub" @@ -296,6 +304,7 @@ source "$ROOT/migrations/1785846769.sh" >/dev/null source "$ROOT/migrations/1787342993.sh" >/dev/null OMARCHY_TEST_MISSING_COMMAND=cursor-agent source "$ROOT/migrations/1788577553.sh" >/dev/null OMARCHY_TEST_MISSING_COMMAND=kimi source "$ROOT/migrations/1786967403.sh" >/dev/null +OMARCHY_TEST_MISSING_COMMAND=qwen source "$ROOT/migrations/1792032538.sh" >/dev/null [[ ! -s $stub_log ]] || fail "agent migrations respect the preinstall opt-out" [[ ! -e $test_home/.local/bin/omp ]] || fail "agent migration removes the obsolete Oh My Pi wrapper after opt-out" @@ -429,9 +438,10 @@ pass "agent migrations install working wrappers without overriding the preinstal "$ROOT/bin/omarchy-mise-install" "$muse_package" muse "$ROOT/bin/omarchy-mise-install" "$kimi_package" kimi +"$ROOT/bin/omarchy-mise-install" "$qwen_package" qwen touch "$test_home/.local/bin/agy" "$test_home/.local/bin/ori" omarchy-remove-preinstalls >/dev/null -for command in agy omp ori grok crush cursor-agent muse kimi; do +for command in agy omp ori grok crush cursor-agent muse kimi qwen; do [[ ! -e $test_home/.local/bin/$command ]] || fail "Remove Preinstalls deletes the $command lazy stub" done pass "Remove Preinstalls deletes every optional agent lazy stub" @@ -514,6 +524,8 @@ declare -A expected_agents=( [github-copilot]="copilot" [kimi]="kimi" [kimi-code]="kimi" + [qwen]="qwen" + [qwen-code]="qwen" [cursor]="cursor-agent" [cursor-agent]="cursor-agent" [muse]="muse" @@ -533,6 +545,7 @@ declare -A expected_packages=( [agy]="$agy_package" [copilot]="copilot" [kimi]="$kimi_package" + [qwen]="$qwen_package" [cursor-agent]="$cursor_agent_package" [muse]="$muse_package" ) @@ -803,6 +816,10 @@ assert_launch hermes env -u HERMES_SESSION_SOURCE hermes chat --yolo --tui "--qu assert_launch agy agy --dangerously-skip-permissions --prompt-interactive "Review this project" assert_launch copilot copilot --allow-all --interactive "Review this project" assert_launch kimi kimi --prompt "Review this project" +assert_launch qwen qwen --yolo "--prompt-interactive=Review this project" +printf '%s\n' qwen >"$agent_file" +omarchy-agent-prompt '-n --help; echo literal' +assert_launched qwen "keeps an option-like prompt literal" qwen --yolo "--prompt-interactive=-n --help; echo literal" pass "agent launcher adapts initial prompts for every supported agent" literal_muse_prompt=$'--disable-sandbox !Crash {$(touch must-not-run)}\ntrailing\\ ' @@ -832,6 +849,7 @@ assert_bypass hermes hermes --yolo assert_bypass agy agy --dangerously-skip-permissions assert_bypass copilot copilot --allow-all assert_bypass kimi kimi --yolo +assert_bypass qwen qwen --yolo pass "agent launcher skips permission prompts for every supported agent" printf '%s\n' "opencode" >"$agent_file" diff --git a/test/shell.d/menu-test.sh b/test/shell.d/menu-test.sh index 15e2c134cd3..32de56500b7 100644 --- a/test/shell.d/menu-test.sh +++ b/test/shell.d/menu-test.sh @@ -253,6 +253,7 @@ const expectedAgents = { codex: { icon: '\ue905', iconFont: 'omarchy', label: 'Codex' }, grok: { icon: '\ue904', iconFont: 'omarchy', label: 'Grok' }, kimi: { icon: '\u{f06a9}', label: 'Kimi' }, + qwen: { icon: '\ue90f', iconFont: 'omarchy', label: 'Qwen Code' }, hermes: { icon: '\ue90a', iconFont: 'omarchy', label: 'Hermes' }, openclaw: { icon: '\ue90c', iconFont: 'omarchy', label: 'OpenClaw' }, copilot: { icon: '', label: 'Copilot' }, @@ -278,7 +279,7 @@ assertDeepEqual( defaultItems .filter(item => item.parent === 'setup.default.agent') .map(item => item.label), - ['Antigravity', 'Claude', 'Codex', 'Copilot', 'Crush', 'Cursor CLI', 'Grok', 'Hermes', 'Kimi', 'Muse Code', 'omp', 'OpenClaw', 'OpenCode', 'Ori', 'Pi'], + ['Antigravity', 'Claude', 'Codex', 'Copilot', 'Crush', 'Cursor CLI', 'Grok', 'Hermes', 'Kimi', 'Muse Code', 'omp', 'OpenClaw', 'OpenCode', 'Ori', 'Pi', 'Qwen Code'], 'menu sorts coding agents alphabetically' ) const expectedDefaults = { @@ -875,5 +876,5 @@ assert( JS font_charset=$(fc-query --format='%{charset}' "$ROOT/default/fonts/omarchy/omarchy.ttf") -[[ $font_charset == *"e900-e90e"* ]] || fail "Omarchy icon font includes every custom menu glyph" +[[ $font_charset == *"e900-e90f"* ]] || fail "Omarchy icon font includes every custom menu glyph" pass "Omarchy icon font includes the official agent marks" diff --git a/test/shell.d/mise-provision-test.sh b/test/shell.d/mise-provision-test.sh index f19a31eb6aa..168ce141b15 100644 --- a/test/shell.d/mise-provision-test.sh +++ b/test/shell.d/mise-provision-test.sh @@ -17,7 +17,7 @@ SH chmod +x "$mock_bin/mise" source "$ROOT/install/user/mise.sh" -for agent in kimi cursor-agent muse; do +for agent in kimi cursor-agent muse qwen; do [[ -x $HOME/.local/bin/$agent && ! -L $HOME/.local/bin/$agent ]] || fail "provisioning creates the missing $agent wrapper" cp "$HOME/.local/bin/$agent" "$test_tmp/$agent-wrapper" @@ -27,7 +27,7 @@ mapfile -d '' -t args <"$OMARCHY_TEST_MISE_LOG" fail "provisioning creates lazy wrappers without installing packages" pass "provisioning creates lazy Kimi, Cursor, and Muse wrappers" -for agent in kimi cursor-agent muse; do +for agent in kimi cursor-agent muse qwen; do launcher="$HOME/.local/bin/$agent" for kind in script binary symlink; do rm -f "$launcher" diff --git a/test/shell.d/qwen-agent-test.sh b/test/shell.d/qwen-agent-test.sh new file mode 100644 index 00000000000..0d3fb8240cc --- /dev/null +++ b/test/shell.d/qwen-agent-test.sh @@ -0,0 +1,85 @@ +#!/bin/bash + +set -euo pipefail +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" + +test_tmp=$(mktemp -d) +trap 'rm -rf "$test_tmp"' EXIT +mock_bin="$test_tmp/bin" +export HOME="$test_tmp/home with spaces" +export PATH="$mock_bin:$HOME/.local/bin:$ROOT/bin:/usr/local/bin:/usr/bin" +export OMARCHY_PATH="$ROOT" +export OMARCHY_TEST_MISE_LOG="$test_tmp/mise-log" +mkdir -p "$mock_bin" "$HOME/.local/bin" +qwen="$HOME/.local/bin/qwen" +package="npm:@qwen-code/qwen-code" + +cat >"$mock_bin/mise" <<'SH' +#!/bin/bash +printf '%s\0' "$@" >>"$OMARCHY_TEST_MISE_LOG" +SH +cat >"$mock_bin/noop" <<'SH' +#!/bin/bash +exit 0 +SH +chmod +x "$mock_bin/mise" "$mock_bin/noop" +for command in gum hyprctl omarchy-webapp-remove-all omarchy-tui-remove-all omarchy-pkg-drop; do + ln -s noop "$mock_bin/$command" +done + +apply_qwen() { + case $1 in + migration) source "$ROOT/migrations/1792032538.sh" >/dev/null ;; + removal) omarchy-remove-preinstalls >/dev/null ;; + esac +} + +apply_qwen migration +[[ -x $qwen && ! -L $qwen ]] || fail "Qwen Code migration creates a missing wrapper" +[[ ! -s $OMARCHY_TEST_MISE_LOG ]] || fail "Qwen Code migration installs the package before first launch" +cp "$qwen" "$test_tmp/wrapper" +apply_qwen migration +cmp -s "$qwen" "$test_tmp/wrapper" || fail "Qwen Code migration is idempotent" +"$qwen" --prompt 'Review this project' +mapfile -d '' -t args <"$OMARCHY_TEST_MISE_LOG" +[[ ${args[*]} == "use -g --quiet $package x $package -- qwen --prompt Review this project" ]] || + fail "Qwen Code migration wrapper installs the official package and forwards arguments" +pass "Qwen Code migration creates a lazy, repeatable wrapper" + +rm -f "$qwen" +mkdir -p "$HOME/.local/state/omarchy" +touch "$HOME/.local/state/omarchy/preinstalls-removed" +apply_qwen migration +[[ ! -e $qwen ]] || fail "Qwen Code migration respects removed preinstalls with a missing command" +rm "$HOME/.local/state/omarchy/preinstalls-removed" +pass "Qwen Code migration respects the preinstall opt-out" + +for operation in migration removal; do + for kind in script binary symlink; do + rm -f "$qwen" + case $kind in + script) printf '#!/bin/bash\necho user-managed qwen\n' >"$qwen" ;; + binary) cp /usr/bin/true "$qwen" ;; + symlink) ln -s /usr/bin/true "$qwen" ;; + esac + [[ $kind == "symlink" ]] || chmod +x "$qwen" + rm -f "$test_tmp/saved-$kind" + cp -P "$qwen" "$test_tmp/saved-$kind" + for attempt in 1 2; do + apply_qwen "$operation" + cmp -s "$qwen" "$test_tmp/saved-$kind" || fail "Qwen Code $operation preserves a user $kind" + if [[ $kind == "symlink" ]]; then + [[ -L $qwen && $(readlink "$qwen") == "/usr/bin/true" ]] || + fail "Qwen Code $operation preserves the user's symlink" + else + [[ ! -L $qwen ]] || fail "Qwen Code $operation keeps the user's $kind as a regular file" + fi + done + pass "Qwen Code $operation preserves a user $kind on repeated runs" + done +done + +omarchy-mise-install "$package" qwen +apply_qwen removal +[[ ! -e $qwen ]] || fail "Remove Preinstalls removes Omarchy's Qwen Code wrapper" +pass "Remove Preinstalls removes Omarchy's Qwen Code wrapper" From 9ae16c551ebac4ddbd3f961ac6aa703ff66cdaa1 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:05:57 -0400 Subject: [PATCH 097/546] Add DeepSeek Harness as a selectable default coding agent DeepSeek Harness (dsh) joins the default agent picker, installed through mise as npm:@deepseek-ai/dsh pinned to 0.2.0-rc.2, since every release so far is a prerelease that mise's npm backend would skip. The launcher opens its browser UI with the danger-full-access permission preset, and runs a supplied prompt once through the headless profile, which is the only way dsh accepts one. The manual names its session-log and feedback uploads and how to turn them off. From #14455, superseding #9777. Co-authored-by: john st Co-authored-by: T2816 Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent | 10 ++++ bin/omarchy-default-agent | 11 +++- bin/omarchy-remove-preinstalls | 6 ++ default/fonts/omarchy/README.md | 1 + default/fonts/omarchy/omarchy.ttf | Bin 10196 -> 10804 bytes default/i18n/en/menu.json | 3 + default/omarchy/omarchy-menu.jsonc | 1 + install/user/mise.sh | 2 + manual/17-ai.md | 3 + migrations/1792033356.sh | 5 ++ test/shell.d/default-agent-test.sh | 21 ++++++- test/shell.d/dsh-agent-test.sh | 85 ++++++++++++++++++++++++++++ test/shell.d/menu-test.sh | 5 +- test/shell.d/mise-provision-test.sh | 4 +- 14 files changed, 150 insertions(+), 7 deletions(-) create mode 100644 migrations/1792033356.sh create mode 100644 test/shell.d/dsh-agent-test.sh diff --git a/bin/omarchy-agent b/bin/omarchy-agent index 59cb198e743..ea3f4578153 100755 --- a/bin/omarchy-agent +++ b/bin/omarchy-agent @@ -143,6 +143,16 @@ qwen) command=(qwen --yolo) [[ -n ${prompt:-} ]] && command+=("--prompt-interactive=$prompt") ;; +dsh) + # The browser UI has no prompt option, so a prompt runs once through the + # headless profile. The launcher consumes one separator; the profile needs + # its own so a prompt starting with a dash stays literal. + if [[ -n ${prompt:-} ]]; then + command=(env DSH_PERMISSION_MODE=danger-full-access dsh headless -- -- "$prompt") + else + command=(env DSH_PERMISSION_MODE=danger-full-access dsh web) + fi + ;; *) echo "Unsupported default agent: $agent" >&2 exit 1 diff --git a/bin/omarchy-default-agent b/bin/omarchy-default-agent index 519200a6d23..758c87018aa 100755 --- a/bin/omarchy-default-agent +++ b/bin/omarchy-default-agent @@ -1,7 +1,7 @@ #!/bin/bash # omarchy:summary=Set and launch the default coding agent -# omarchy:args=[pi|omp|opencode|ori|claude|codex|grok|openclaw|agy|hermes|copilot|crush|cursor-agent|muse|kimi|qwen] +# omarchy:args=[pi|omp|opencode|ori|claude|codex|grok|openclaw|agy|hermes|copilot|crush|cursor-agent|muse|kimi|qwen|dsh] # omarchy:examples=omarchy default agent | omarchy default agent codex | omarchy default agent claude installing=false @@ -44,8 +44,15 @@ muse | muse-code | musecode) cursor | cursor-agent) agent="cursor-agent"; name="Cursor CLI" ;; kimi | kimi-code) agent="kimi"; name="Kimi"; agent_package="npm:@moonshot-ai/kimi-code" ;; qwen | qwen-code) agent="qwen"; name="Qwen Code"; agent_package="npm:@qwen-code/qwen-code" ;; +dsh | deepseek | deepseek-harness) + agent="dsh"; name="DeepSeek Harness" + # Every DSH release so far is a prerelease, and mise's npm backend lists only + # stable versions, so a bare package name resolves to nothing. Drop the pin + # once a stable release exists upstream. + agent_package="npm:@deepseek-ai/dsh@0.2.0-rc.2" + ;; *) - echo "Usage: omarchy-default-agent " + echo "Usage: omarchy-default-agent " exit 1 ;; esac diff --git a/bin/omarchy-remove-preinstalls b/bin/omarchy-remove-preinstalls index f0568427d22..01e35639c0d 100755 --- a/bin/omarchy-remove-preinstalls +++ b/bin/omarchy-remove-preinstalls @@ -42,6 +42,12 @@ if gum confirm --affirmative "$(omarchy-i18n common.yes)" --negative "$(omarchy- rm -f ~/.local/bin/qwen fi + # Preserve a user-managed DeepSeek Harness launcher at the same path. + if [[ -f ~/.local/bin/dsh && ! -L ~/.local/bin/dsh ]] && + grep -Eq '^mise use -g .*"npm:@deepseek-ai/dsh(@[^"]*)?"' ~/.local/bin/dsh; then + rm -f ~/.local/bin/dsh + fi + omarchy-pkg-drop \ aether \ cliamp \ diff --git a/default/fonts/omarchy/README.md b/default/fonts/omarchy/README.md index 544fbda5bf6..e36cab3be02 100644 --- a/default/fonts/omarchy/README.md +++ b/default/fonts/omarchy/README.md @@ -18,6 +18,7 @@ The private-use glyphs in `omarchy.ttf` are: - `U+E90D` — Cursor, from - `U+E90E` — Claude, from - `U+E90F` — Qwen, from +- `U+E910` — DeepSeek, from The agent marks are monochrome so the menu can render them using the active theme's foreground and selection colors. diff --git a/default/fonts/omarchy/omarchy.ttf b/default/fonts/omarchy/omarchy.ttf index 429a3e4cc342b8f40adfd18bcd2f20d86d639163..182d6927c840314acd8788e5f02a8420bbd79cc2 100644 GIT binary patch delta 842 zcmXw1O=}ZT6us|rX5O30SCg4(lbF~xiN>1NB%RQds!$iDP>a~2xUs=RW9=8TU=<{G zR}fUlLKnId+(}DaxDZ^paBHC}=|2z=!L8z(n9O13zH`qxm-ilXH2;3#`MK6wcm@FN z0^r#@D-UkJIAzQMkRt&1YqL?S7v5E$0wmr5fY($F9nk^Bj};$ou5N8#0`>|ZsVLM} z)^F8D-ji;#sd##=w%RB?Iew)02LQRhvA(%wKCDkmfYU3=@MZsF z5}{E08~R~{6pyKm@qfxGBb&WHG0+tzfTqltcDoONdgu+A^nCQV+Rv#*)q6AY$wZ9ALv=iv_azmqw z_mhbzW+R0;;yN*%8jK($WGYq762c^+h%m->D3_{;XxNTX(!vN)%7W6O*q3!HEF}?v z{NPDDKn>kA#q|pphn0vh;{j$`oMUX}EemPau0Eccm?$q=Two$KjdO}P6NYJ8x`?{2 zlMGQ23~8)zWjHMmkwGh>1@nW%#e^#S`Az6eYJ|q7s>Q_!WjK1{^85(lnqZtWD`*%7 z!hMmm-w8oNT_931?#W2N(6|}Xr67z@iU~4|h?IhJuI^j?sDS*pRGRsv?YEZ(=23rp rZ(xa1fXQ|!UewKI$I*W<8~-VlSj~-L6lpK#mW5t#G&VLHjeGD98Wea( delta 226 zcmdlIa>ajwVSTQkI|Bn_2Ll6>S9(rm+Br>TEd~bWD+~-`4H>D4DSGRy`WP6bLKqkr z%rbxi9E>cIKzb2o8G&Mq#>P@<={F{Cknv}fm?WDpIYrip10)W_5}W7AzUO9?nB1pn JF!_LH763iwHx>W@ diff --git a/default/i18n/en/menu.json b/default/i18n/en/menu.json index 54d0ab3393c..642d5c9016c 100644 --- a/default/i18n/en/menu.json +++ b/default/i18n/en/menu.json @@ -500,6 +500,9 @@ "setup.default.agent.cursor-agent": { "label": "Cursor CLI" }, + "setup.default.agent.dsh": { + "label": "DeepSeek Harness" + }, "setup.default.agent.grok": { "label": "Grok" }, diff --git a/default/omarchy/omarchy-menu.jsonc b/default/omarchy/omarchy-menu.jsonc index b2d4149e16a..9d8910837f1 100644 --- a/default/omarchy/omarchy-menu.jsonc +++ b/default/omarchy/omarchy-menu.jsonc @@ -196,6 +196,7 @@ "setup.default.agent.copilot": {"icon":"","label":"Copilot","checked":"[[ \"$(omarchy-default-agent)\" == \"copilot\" ]]","action":"omarchy-default-agent copilot"}, "setup.default.agent.crush": {"icon":"󰋑","label":"Crush","checked":"[[ \"$(omarchy-default-agent)\" == \"crush\" ]]","action":"omarchy-default-agent crush"}, "setup.default.agent.cursor-agent": {"icon":"","iconFont":"omarchy","label":"Cursor CLI","checked":"[[ \"$(omarchy-default-agent)\" == \"cursor-agent\" ]]","action":"omarchy-default-agent cursor-agent"}, + "setup.default.agent.dsh": {"icon":"","iconFont":"omarchy","label":"DeepSeek Harness","checked":"[[ \"$(omarchy-default-agent)\" == \"dsh\" ]]","action":"omarchy-default-agent dsh"}, "setup.default.agent.grok": {"icon":"","iconFont":"omarchy","label":"Grok","checked":"[[ \"$(omarchy-default-agent)\" == \"grok\" ]]","action":"omarchy-default-agent grok"}, "setup.default.agent.hermes": {"icon":"","iconFont":"omarchy","label":"Hermes","checked":"[[ \"$(omarchy-default-agent)\" == \"hermes\" ]]","action":"omarchy-default-agent hermes"}, "setup.default.agent.kimi": {"icon":"󰚩","label":"Kimi","checked":"[[ \"$(omarchy-default-agent)\" == \"kimi\" ]]","action":"omarchy-default-agent kimi"}, diff --git a/install/user/mise.sh b/install/user/mise.sh index 0bf535ac85c..a7af9bb7609 100644 --- a/install/user/mise.sh +++ b/install/user/mise.sh @@ -15,6 +15,8 @@ omarchy-mise-install github:can1357/oh-my-pi omp omarchy-mise-install grok omarchy-mise-install npm:@moonshot-ai/kimi-code kimi omarchy-mise-install npm:@qwen-code/qwen-code qwen +# Pinned: every DSH release so far is a prerelease, which mise's npm backend skips. +omarchy-mise-install npm:@deepseek-ai/dsh@0.2.0-rc.2 dsh omarchy-mise-install cursor-agent omarchy-mise-install npm:@kitlangton/ghui ghui omarchy-mise-install aqua:modem-dev/hunk hunk diff --git a/manual/17-ai.md b/manual/17-ai.md index ed33dea4e23..96703c9b30a 100644 --- a/manual/17-ai.md +++ b/manual/17-ai.md @@ -10,6 +10,7 @@ Omarchy treats AI coding agents as first-class citizens, but it doesn't pick a f | `agy` | [Google Antigravity CLI](https://github.com/google-antigravity/antigravity-cli) | | `copilot` | [GitHub Copilot CLI](https://github.com/github/copilot-cli) | | `crush` | [Crush](https://github.com/charmbracelet/crush) | +| `dsh` | [DeepSeek Harness](https://github.com/deepseek-ai/deepseek-harness) | | `grok` | [Grok Build](https://x.ai/build) from xAI | | `kimi` | [Kimi Code](https://www.kimi.com/code/docs/en/) | | `qwen` | [Qwen Code](https://github.com/QwenLM/qwen-code), Alibaba's coding agent | @@ -26,6 +27,8 @@ To wrap an additional CLI the same way, run `omarchy-mise-install [com `grok` is installed through mise's first-party `grok` package, the same way as `claude` and `codex`. Run `grok`, or pick it under _Setup > Defaults > Agent_. Skip the `curl | bash` installer on x.ai. That path puts a second copy in `~/.grok/bin`, prepends it on `PATH`, and overwrites Omarchy's wrapper. +`dsh` opens its browser UI on `127.0.0.1:3080`, with the launch directory as its default workspace. `omarchy agent prompt` runs the supplied task once through its headless profile. Omarchy starts it with DeepSeek's no-approval permission preset. The developer-preview release is pinned because mise lists only stable npm versions by default. DeepSeek Harness uploads session logs to DeepSeek when you use its official model API; turn that off under _Settings > General > Upload Session Log_ in its browser UI. Feedback you send it is uploaded with the session too, unless you start it with `DSH_TELEMETRY_MODE=DISABLED`. + ### The default agent Pick your default agent with `omarchy default agent ` or under _Setup > Defaults > Agent_ in the Omarchy Menu (`Super + Space`). If the agent isn't installed yet, picking it installs it first. A fresh Omarchy will invite you to make this choice with a one-time notification. diff --git a/migrations/1792033356.sh b/migrations/1792033356.sh new file mode 100644 index 00000000000..66367de151c --- /dev/null +++ b/migrations/1792033356.sh @@ -0,0 +1,5 @@ +echo "Install DeepSeek Harness coding agent mise wrapper" + +if [[ ! -f $HOME/.local/state/omarchy/preinstalls-removed ]] && omarchy-cmd-missing dsh; then + omarchy-mise-install npm:@deepseek-ai/dsh@0.2.0-rc.2 dsh +fi diff --git a/test/shell.d/default-agent-test.sh b/test/shell.d/default-agent-test.sh index 1e74d5f1e2d..3e543e50363 100644 --- a/test/shell.d/default-agent-test.sh +++ b/test/shell.d/default-agent-test.sh @@ -133,6 +133,7 @@ agy_package="antigravity-cli" ori_package="github:OpenRouterLabs/ori-releases" kimi_package="npm:@moonshot-ai/kimi-code" qwen_package="npm:@qwen-code/qwen-code" +dsh_package="npm:@deepseek-ai/dsh@0.2.0-rc.2" cursor_agent_package="cursor-agent" muse_package="http:muse[url=https://api.meta.ai/muse-launcher.sh,bin=muse,version_list_url=https://api.meta.ai/muse-code/channels/muse-stable,version_json_path=.version]" @@ -155,6 +156,7 @@ assert_lazy_stub "$crush_package" crush assert_lazy_stub "$ori_package" ori assert_lazy_stub "$kimi_package" kimi assert_lazy_stub "$qwen_package" qwen +assert_lazy_stub "$dsh_package" dsh assert_lazy_stub "$cursor_agent_package" cursor-agent assert_lazy_stub "$muse_package" muse pass "custom agent lazy stubs preserve their mise packages" @@ -169,6 +171,7 @@ grep -Fx "$ori_package ori" "$stub_log" >/dev/null || fail "user setup creates t grep -Fx "$muse_package muse" "$stub_log" >/dev/null || fail "user setup creates the Muse lazy stub" grep -Fx "$kimi_package kimi" "$stub_log" >/dev/null || fail "user setup creates the Kimi lazy stub" grep -Fx "$qwen_package qwen" "$stub_log" >/dev/null || fail "user setup creates the Qwen Code lazy stub" +grep -Fx "$dsh_package dsh" "$stub_log" >/dev/null || fail "user setup creates the DeepSeek Harness lazy stub" pass "user setup creates the custom agent lazy stubs" : >"$stub_log" @@ -177,6 +180,7 @@ grep -Fx "$cursor_agent_package" "$stub_log" >/dev/null || fail "user setup recr grep -Fx "$muse_package muse" "$stub_log" >/dev/null || fail "user setup recreates the Muse lazy stub" grep -Fx "$kimi_package kimi" "$stub_log" >/dev/null || fail "user setup recreates the Kimi lazy stub" grep -Fx "$qwen_package qwen" "$stub_log" >/dev/null || fail "user setup recreates the Qwen Code lazy stub" +grep -Fx "$dsh_package dsh" "$stub_log" >/dev/null || fail "user setup recreates the DeepSeek Harness lazy stub" pass "user setup recreates wrappers even when the commands are present" : >"$stub_log" @@ -201,6 +205,10 @@ grep -Fx "$kimi_package kimi" "$stub_log" >/dev/null || fail "Kimi migration cre OMARCHY_TEST_MISSING_COMMAND=qwen source "$ROOT/migrations/1792032538.sh" >/dev/null grep -Fx "$qwen_package qwen" "$stub_log" >/dev/null || fail "Qwen Code migration creates its lazy stub" +: >"$stub_log" +OMARCHY_TEST_MISSING_COMMAND=dsh source "$ROOT/migrations/1792033356.sh" >/dev/null +grep -Fx "$dsh_package dsh" "$stub_log" >/dev/null || fail "DeepSeek Harness migration creates its lazy stub" + : >"$stub_log" source "$ROOT/migrations/1785617047.sh" >/dev/null grep -Fx "$omp_package omp" "$stub_log" >/dev/null || fail "Oh My Pi migration creates a working lazy stub" @@ -305,6 +313,7 @@ source "$ROOT/migrations/1787342993.sh" >/dev/null OMARCHY_TEST_MISSING_COMMAND=cursor-agent source "$ROOT/migrations/1788577553.sh" >/dev/null OMARCHY_TEST_MISSING_COMMAND=kimi source "$ROOT/migrations/1786967403.sh" >/dev/null OMARCHY_TEST_MISSING_COMMAND=qwen source "$ROOT/migrations/1792032538.sh" >/dev/null +OMARCHY_TEST_MISSING_COMMAND=dsh source "$ROOT/migrations/1792033356.sh" >/dev/null [[ ! -s $stub_log ]] || fail "agent migrations respect the preinstall opt-out" [[ ! -e $test_home/.local/bin/omp ]] || fail "agent migration removes the obsolete Oh My Pi wrapper after opt-out" @@ -439,9 +448,10 @@ pass "agent migrations install working wrappers without overriding the preinstal "$ROOT/bin/omarchy-mise-install" "$muse_package" muse "$ROOT/bin/omarchy-mise-install" "$kimi_package" kimi "$ROOT/bin/omarchy-mise-install" "$qwen_package" qwen +"$ROOT/bin/omarchy-mise-install" "$dsh_package" dsh touch "$test_home/.local/bin/agy" "$test_home/.local/bin/ori" omarchy-remove-preinstalls >/dev/null -for command in agy omp ori grok crush cursor-agent muse kimi qwen; do +for command in agy omp ori grok crush cursor-agent muse kimi qwen dsh; do [[ ! -e $test_home/.local/bin/$command ]] || fail "Remove Preinstalls deletes the $command lazy stub" done pass "Remove Preinstalls deletes every optional agent lazy stub" @@ -526,6 +536,9 @@ declare -A expected_agents=( [kimi-code]="kimi" [qwen]="qwen" [qwen-code]="qwen" + [dsh]="dsh" + [deepseek]="dsh" + [deepseek-harness]="dsh" [cursor]="cursor-agent" [cursor-agent]="cursor-agent" [muse]="muse" @@ -546,6 +559,7 @@ declare -A expected_packages=( [copilot]="copilot" [kimi]="$kimi_package" [qwen]="$qwen_package" + [dsh]="$dsh_package" [cursor-agent]="$cursor_agent_package" [muse]="$muse_package" ) @@ -820,6 +834,10 @@ assert_launch qwen qwen --yolo "--prompt-interactive=Review this project" printf '%s\n' qwen >"$agent_file" omarchy-agent-prompt '-n --help; echo literal' assert_launched qwen "keeps an option-like prompt literal" qwen --yolo "--prompt-interactive=-n --help; echo literal" +assert_launch dsh env DSH_PERMISSION_MODE=danger-full-access dsh headless -- -- "Review this project" +printf '%s\n' dsh >"$agent_file" +omarchy-agent-prompt '-n --help; echo literal' +assert_launched dsh "keeps an option-like prompt literal" env DSH_PERMISSION_MODE=danger-full-access dsh headless -- -- '-n --help; echo literal' pass "agent launcher adapts initial prompts for every supported agent" literal_muse_prompt=$'--disable-sandbox !Crash {$(touch must-not-run)}\ntrailing\\ ' @@ -850,6 +868,7 @@ assert_bypass agy agy --dangerously-skip-permissions assert_bypass copilot copilot --allow-all assert_bypass kimi kimi --yolo assert_bypass qwen qwen --yolo +assert_bypass dsh env DSH_PERMISSION_MODE=danger-full-access dsh web pass "agent launcher skips permission prompts for every supported agent" printf '%s\n' "opencode" >"$agent_file" diff --git a/test/shell.d/dsh-agent-test.sh b/test/shell.d/dsh-agent-test.sh new file mode 100644 index 00000000000..e3fcc43bf68 --- /dev/null +++ b/test/shell.d/dsh-agent-test.sh @@ -0,0 +1,85 @@ +#!/bin/bash + +set -euo pipefail +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" + +test_tmp=$(mktemp -d) +trap 'rm -rf "$test_tmp"' EXIT +mock_bin="$test_tmp/bin" +export HOME="$test_tmp/home with spaces" +export PATH="$mock_bin:$HOME/.local/bin:$ROOT/bin:/usr/local/bin:/usr/bin" +export OMARCHY_PATH="$ROOT" +export OMARCHY_TEST_MISE_LOG="$test_tmp/mise-log" +mkdir -p "$mock_bin" "$HOME/.local/bin" +dsh="$HOME/.local/bin/dsh" +package="npm:@deepseek-ai/dsh@0.2.0-rc.2" + +cat >"$mock_bin/mise" <<'SH' +#!/bin/bash +printf '%s\0' "$@" >>"$OMARCHY_TEST_MISE_LOG" +SH +cat >"$mock_bin/noop" <<'SH' +#!/bin/bash +exit 0 +SH +chmod +x "$mock_bin/mise" "$mock_bin/noop" +for command in gum hyprctl omarchy-webapp-remove-all omarchy-tui-remove-all omarchy-pkg-drop; do + ln -s noop "$mock_bin/$command" +done + +apply_dsh() { + case $1 in + migration) source "$ROOT/migrations/1792033356.sh" >/dev/null ;; + removal) omarchy-remove-preinstalls >/dev/null ;; + esac +} + +apply_dsh migration +[[ -x $dsh && ! -L $dsh ]] || fail "DeepSeek Harness migration creates a missing wrapper" +[[ ! -s $OMARCHY_TEST_MISE_LOG ]] || fail "DeepSeek Harness migration installs the package before first launch" +cp "$dsh" "$test_tmp/wrapper" +apply_dsh migration +cmp -s "$dsh" "$test_tmp/wrapper" || fail "DeepSeek Harness migration is idempotent" +"$dsh" --prompt 'Review this project' +mapfile -d '' -t args <"$OMARCHY_TEST_MISE_LOG" +[[ ${args[*]} == "use -g --quiet $package x $package -- dsh --prompt Review this project" ]] || + fail "DeepSeek Harness migration wrapper installs the official package and forwards arguments" +pass "DeepSeek Harness migration creates a lazy, repeatable wrapper" + +rm -f "$dsh" +mkdir -p "$HOME/.local/state/omarchy" +touch "$HOME/.local/state/omarchy/preinstalls-removed" +apply_dsh migration +[[ ! -e $dsh ]] || fail "DeepSeek Harness migration respects removed preinstalls with a missing command" +rm "$HOME/.local/state/omarchy/preinstalls-removed" +pass "DeepSeek Harness migration respects the preinstall opt-out" + +for operation in migration removal; do + for kind in script binary symlink; do + rm -f "$dsh" + case $kind in + script) printf '#!/bin/bash\necho user-managed dsh\n' >"$dsh" ;; + binary) cp /usr/bin/true "$dsh" ;; + symlink) ln -s /usr/bin/true "$dsh" ;; + esac + [[ $kind == "symlink" ]] || chmod +x "$dsh" + rm -f "$test_tmp/saved-$kind" + cp -P "$dsh" "$test_tmp/saved-$kind" + for attempt in 1 2; do + apply_dsh "$operation" + cmp -s "$dsh" "$test_tmp/saved-$kind" || fail "DeepSeek Harness $operation preserves a user $kind" + if [[ $kind == "symlink" ]]; then + [[ -L $dsh && $(readlink "$dsh") == "/usr/bin/true" ]] || + fail "DeepSeek Harness $operation preserves the user's symlink" + else + [[ ! -L $dsh ]] || fail "DeepSeek Harness $operation keeps the user's $kind as a regular file" + fi + done + pass "DeepSeek Harness $operation preserves a user $kind on repeated runs" + done +done + +omarchy-mise-install "$package" dsh +apply_dsh removal +[[ ! -e $dsh ]] || fail "Remove Preinstalls removes Omarchy's DeepSeek Harness wrapper" +pass "Remove Preinstalls removes Omarchy's DeepSeek Harness wrapper" diff --git a/test/shell.d/menu-test.sh b/test/shell.d/menu-test.sh index 32de56500b7..addf7e676d8 100644 --- a/test/shell.d/menu-test.sh +++ b/test/shell.d/menu-test.sh @@ -254,6 +254,7 @@ const expectedAgents = { grok: { icon: '\ue904', iconFont: 'omarchy', label: 'Grok' }, kimi: { icon: '\u{f06a9}', label: 'Kimi' }, qwen: { icon: '\ue90f', iconFont: 'omarchy', label: 'Qwen Code' }, + dsh: { icon: '\ue910', iconFont: 'omarchy', label: 'DeepSeek Harness' }, hermes: { icon: '\ue90a', iconFont: 'omarchy', label: 'Hermes' }, openclaw: { icon: '\ue90c', iconFont: 'omarchy', label: 'OpenClaw' }, copilot: { icon: '', label: 'Copilot' }, @@ -279,7 +280,7 @@ assertDeepEqual( defaultItems .filter(item => item.parent === 'setup.default.agent') .map(item => item.label), - ['Antigravity', 'Claude', 'Codex', 'Copilot', 'Crush', 'Cursor CLI', 'Grok', 'Hermes', 'Kimi', 'Muse Code', 'omp', 'OpenClaw', 'OpenCode', 'Ori', 'Pi', 'Qwen Code'], + ['Antigravity', 'Claude', 'Codex', 'Copilot', 'Crush', 'Cursor CLI', 'DeepSeek Harness', 'Grok', 'Hermes', 'Kimi', 'Muse Code', 'omp', 'OpenClaw', 'OpenCode', 'Ori', 'Pi', 'Qwen Code'], 'menu sorts coding agents alphabetically' ) const expectedDefaults = { @@ -876,5 +877,5 @@ assert( JS font_charset=$(fc-query --format='%{charset}' "$ROOT/default/fonts/omarchy/omarchy.ttf") -[[ $font_charset == *"e900-e90f"* ]] || fail "Omarchy icon font includes every custom menu glyph" +[[ $font_charset == *"e900-e910"* ]] || fail "Omarchy icon font includes every custom menu glyph" pass "Omarchy icon font includes the official agent marks" diff --git a/test/shell.d/mise-provision-test.sh b/test/shell.d/mise-provision-test.sh index 168ce141b15..66d727f91bd 100644 --- a/test/shell.d/mise-provision-test.sh +++ b/test/shell.d/mise-provision-test.sh @@ -17,7 +17,7 @@ SH chmod +x "$mock_bin/mise" source "$ROOT/install/user/mise.sh" -for agent in kimi cursor-agent muse qwen; do +for agent in kimi cursor-agent muse qwen dsh; do [[ -x $HOME/.local/bin/$agent && ! -L $HOME/.local/bin/$agent ]] || fail "provisioning creates the missing $agent wrapper" cp "$HOME/.local/bin/$agent" "$test_tmp/$agent-wrapper" @@ -27,7 +27,7 @@ mapfile -d '' -t args <"$OMARCHY_TEST_MISE_LOG" fail "provisioning creates lazy wrappers without installing packages" pass "provisioning creates lazy Kimi, Cursor, and Muse wrappers" -for agent in kimi cursor-agent muse qwen; do +for agent in kimi cursor-agent muse qwen dsh; do launcher="$HOME/.local/bin/$agent" for kind in script binary symlink; do rm -f "$launcher" From fe3b82f9d483b326f0907c8acbcb2d90010d0da9 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:06:28 -0400 Subject: [PATCH 098/546] Add MiniMax Code as a selectable default coding agent MiniMax Code (mcode) joins the default agent picker, installed through mise as npm:@minimax-ai/code with its better-sqlite3 build allowed. Its TUI has no no-approval flag, since permission policies belong to exec mode, so the launcher runs plain mcode and passes a prompt after --. The launcher part of #11367; its usage collector and panel changes are left to the agents panel work. Co-authored-by: Josh LeBlanc Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent | 5 ++ bin/omarchy-default-agent | 5 +- bin/omarchy-remove-preinstalls | 6 ++ default/fonts/omarchy/README.md | 1 + default/fonts/omarchy/omarchy.ttf | Bin 10804 -> 11308 bytes default/i18n/en/menu.json | 3 + default/omarchy/omarchy-menu.jsonc | 1 + install/user/mise.sh | 1 + manual/17-ai.md | 1 + migrations/1792033557.sh | 5 ++ test/shell.d/default-agent-test.sh | 18 +++++- test/shell.d/mcode-agent-test.sh | 85 ++++++++++++++++++++++++++++ test/shell.d/menu-test.sh | 5 +- test/shell.d/mise-provision-test.sh | 4 +- 14 files changed, 133 insertions(+), 7 deletions(-) create mode 100644 migrations/1792033557.sh create mode 100644 test/shell.d/mcode-agent-test.sh diff --git a/bin/omarchy-agent b/bin/omarchy-agent index ea3f4578153..b67d7776ac1 100755 --- a/bin/omarchy-agent +++ b/bin/omarchy-agent @@ -153,6 +153,11 @@ dsh) command=(env DSH_PERMISSION_MODE=danger-full-access dsh web) fi ;; +mcode) + # The TUI has no no-approval flag; permission options belong to exec mode. + command=(mcode) + [[ -n ${prompt:-} ]] && command+=(-- "$prompt") + ;; *) echo "Unsupported default agent: $agent" >&2 exit 1 diff --git a/bin/omarchy-default-agent b/bin/omarchy-default-agent index 758c87018aa..2f69bb7b224 100755 --- a/bin/omarchy-default-agent +++ b/bin/omarchy-default-agent @@ -1,7 +1,7 @@ #!/bin/bash # omarchy:summary=Set and launch the default coding agent -# omarchy:args=[pi|omp|opencode|ori|claude|codex|grok|openclaw|agy|hermes|copilot|crush|cursor-agent|muse|kimi|qwen|dsh] +# omarchy:args=[pi|omp|opencode|ori|claude|codex|grok|openclaw|agy|hermes|copilot|crush|cursor-agent|muse|kimi|qwen|dsh|mcode] # omarchy:examples=omarchy default agent | omarchy default agent codex | omarchy default agent claude installing=false @@ -51,8 +51,9 @@ dsh | deepseek | deepseek-harness) # once a stable release exists upstream. agent_package="npm:@deepseek-ai/dsh@0.2.0-rc.2" ;; +mcode | minimax-code | minimax) agent="mcode"; name="MiniMax Code"; agent_package="npm:@minimax-ai/code[allow_builds=better-sqlite3]" ;; *) - echo "Usage: omarchy-default-agent " + echo "Usage: omarchy-default-agent " exit 1 ;; esac diff --git a/bin/omarchy-remove-preinstalls b/bin/omarchy-remove-preinstalls index 01e35639c0d..be3f3e9cf4c 100755 --- a/bin/omarchy-remove-preinstalls +++ b/bin/omarchy-remove-preinstalls @@ -48,6 +48,12 @@ if gum confirm --affirmative "$(omarchy-i18n common.yes)" --negative "$(omarchy- rm -f ~/.local/bin/dsh fi + # Preserve a user-managed MiniMax Code launcher at the same path. + if [[ -f ~/.local/bin/mcode && ! -L ~/.local/bin/mcode ]] && + grep -Eq '^mise use -g .*"npm:@minimax-ai/code\[allow_builds=better-sqlite3\]"' ~/.local/bin/mcode; then + rm -f ~/.local/bin/mcode + fi + omarchy-pkg-drop \ aether \ cliamp \ diff --git a/default/fonts/omarchy/README.md b/default/fonts/omarchy/README.md index e36cab3be02..23543c388ab 100644 --- a/default/fonts/omarchy/README.md +++ b/default/fonts/omarchy/README.md @@ -19,6 +19,7 @@ The private-use glyphs in `omarchy.ttf` are: - `U+E90E` — Claude, from - `U+E90F` — Qwen, from - `U+E910` — DeepSeek, from +- `U+E911` — MiniMax, from The agent marks are monochrome so the menu can render them using the active theme's foreground and selection colors. diff --git a/default/fonts/omarchy/omarchy.ttf b/default/fonts/omarchy/omarchy.ttf index 182d6927c840314acd8788e5f02a8420bbd79cc2..d11f8d92deda4cdc28f446645e16af8973e427af 100644 GIT binary patch delta 748 zcmYL{!E4h{9LImZq)YM|vn5SlE7kgn8>v$mx;7~c55to>!Ha`2gi+RM2kSc6lzNa} zJPd?U>|GFcQV~Xncu?rSP?#6#&D~Tyd6=-`?`3TbB=Gruzwi5dyqCiV@0Xf48a3w* z0PX@PcOP#&d$c;aas@y=0!SZL)=K4>4|96}m&k{(LV|&@3&b7bsY-RHejCPm#K**y zjoNBy4hv%d;-1V`OZ6=bmFvVWh_joe>RSF~?>X^4fO)W0+uk{M)mZ!paBYbK-Gk2= zM5px)e=&!c)5f%RSBYW!NbR0)C@>GZYDcZZeRxBjmr54H!0~f_PSZ8rq0i)Hf|0@p zN)X5dPo|Rj*3dJy$)E(kfNbLBMd5n-&_&G}ZgFd2BU*l{E&LSCrkc~Av>Qs`qTn8)$3=78V9XdVxQK3n zVian+amM0q|D9+ILW*}LJ)u3H$mB^vMiTw+!jm-imZs~PwJ>y) zhr**H;hyMES#qj{8F@%_uXOd8&(HkzzphN&`~zl;1c2k#my5SB{VG$AkucvXX9{s^ hS!U@?w_j&|o0x9xP0uRmwmQ?>re0m&T&I7K?0+mycDeun delta 236 zcmZ1zu_a`JVSS#U2Ll6R2Ll6>S9(rm+H?(0Uj_!|D+~-07cx>4Q}otZ^)WEWEMZ_^ zFv|c6a4@n+0r_`;e3gvck_s0F);SCevOq3FN=|-qqRw&?RvCAQ_{hF-@)S9LHc23#fo1YVCi%%6@(z=I/dev/null || fail "user setup creates grep -Fx "$kimi_package kimi" "$stub_log" >/dev/null || fail "user setup creates the Kimi lazy stub" grep -Fx "$qwen_package qwen" "$stub_log" >/dev/null || fail "user setup creates the Qwen Code lazy stub" grep -Fx "$dsh_package dsh" "$stub_log" >/dev/null || fail "user setup creates the DeepSeek Harness lazy stub" +grep -Fx "$mcode_package mcode" "$stub_log" >/dev/null || fail "user setup creates the MiniMax Code lazy stub" pass "user setup creates the custom agent lazy stubs" : >"$stub_log" @@ -181,6 +184,7 @@ grep -Fx "$muse_package muse" "$stub_log" >/dev/null || fail "user setup recreat grep -Fx "$kimi_package kimi" "$stub_log" >/dev/null || fail "user setup recreates the Kimi lazy stub" grep -Fx "$qwen_package qwen" "$stub_log" >/dev/null || fail "user setup recreates the Qwen Code lazy stub" grep -Fx "$dsh_package dsh" "$stub_log" >/dev/null || fail "user setup recreates the DeepSeek Harness lazy stub" +grep -Fx "$mcode_package mcode" "$stub_log" >/dev/null || fail "user setup recreates the MiniMax Code lazy stub" pass "user setup recreates wrappers even when the commands are present" : >"$stub_log" @@ -209,6 +213,10 @@ grep -Fx "$qwen_package qwen" "$stub_log" >/dev/null || fail "Qwen Code migratio OMARCHY_TEST_MISSING_COMMAND=dsh source "$ROOT/migrations/1792033356.sh" >/dev/null grep -Fx "$dsh_package dsh" "$stub_log" >/dev/null || fail "DeepSeek Harness migration creates its lazy stub" +: >"$stub_log" +OMARCHY_TEST_MISSING_COMMAND=mcode source "$ROOT/migrations/1792033557.sh" >/dev/null +grep -Fx "$mcode_package mcode" "$stub_log" >/dev/null || fail "MiniMax Code migration creates its lazy stub" + : >"$stub_log" source "$ROOT/migrations/1785617047.sh" >/dev/null grep -Fx "$omp_package omp" "$stub_log" >/dev/null || fail "Oh My Pi migration creates a working lazy stub" @@ -314,6 +322,7 @@ OMARCHY_TEST_MISSING_COMMAND=cursor-agent source "$ROOT/migrations/1788577553.sh OMARCHY_TEST_MISSING_COMMAND=kimi source "$ROOT/migrations/1786967403.sh" >/dev/null OMARCHY_TEST_MISSING_COMMAND=qwen source "$ROOT/migrations/1792032538.sh" >/dev/null OMARCHY_TEST_MISSING_COMMAND=dsh source "$ROOT/migrations/1792033356.sh" >/dev/null +OMARCHY_TEST_MISSING_COMMAND=mcode source "$ROOT/migrations/1792033557.sh" >/dev/null [[ ! -s $stub_log ]] || fail "agent migrations respect the preinstall opt-out" [[ ! -e $test_home/.local/bin/omp ]] || fail "agent migration removes the obsolete Oh My Pi wrapper after opt-out" @@ -449,9 +458,10 @@ pass "agent migrations install working wrappers without overriding the preinstal "$ROOT/bin/omarchy-mise-install" "$kimi_package" kimi "$ROOT/bin/omarchy-mise-install" "$qwen_package" qwen "$ROOT/bin/omarchy-mise-install" "$dsh_package" dsh +"$ROOT/bin/omarchy-mise-install" "$mcode_package" mcode touch "$test_home/.local/bin/agy" "$test_home/.local/bin/ori" omarchy-remove-preinstalls >/dev/null -for command in agy omp ori grok crush cursor-agent muse kimi qwen dsh; do +for command in agy omp ori grok crush cursor-agent muse kimi qwen dsh mcode; do [[ ! -e $test_home/.local/bin/$command ]] || fail "Remove Preinstalls deletes the $command lazy stub" done pass "Remove Preinstalls deletes every optional agent lazy stub" @@ -539,6 +549,9 @@ declare -A expected_agents=( [dsh]="dsh" [deepseek]="dsh" [deepseek-harness]="dsh" + [mcode]="mcode" + [minimax-code]="mcode" + [minimax]="mcode" [cursor]="cursor-agent" [cursor-agent]="cursor-agent" [muse]="muse" @@ -560,6 +573,7 @@ declare -A expected_packages=( [kimi]="$kimi_package" [qwen]="$qwen_package" [dsh]="$dsh_package" + [mcode]="$mcode_package" [cursor-agent]="$cursor_agent_package" [muse]="$muse_package" ) @@ -838,6 +852,7 @@ assert_launch dsh env DSH_PERMISSION_MODE=danger-full-access dsh headless -- -- printf '%s\n' dsh >"$agent_file" omarchy-agent-prompt '-n --help; echo literal' assert_launched dsh "keeps an option-like prompt literal" env DSH_PERMISSION_MODE=danger-full-access dsh headless -- -- '-n --help; echo literal' +assert_launch mcode mcode -- "Review this project" pass "agent launcher adapts initial prompts for every supported agent" literal_muse_prompt=$'--disable-sandbox !Crash {$(touch must-not-run)}\ntrailing\\ ' @@ -869,6 +884,7 @@ assert_bypass copilot copilot --allow-all assert_bypass kimi kimi --yolo assert_bypass qwen qwen --yolo assert_bypass dsh env DSH_PERMISSION_MODE=danger-full-access dsh web +assert_bypass mcode mcode pass "agent launcher skips permission prompts for every supported agent" printf '%s\n' "opencode" >"$agent_file" diff --git a/test/shell.d/mcode-agent-test.sh b/test/shell.d/mcode-agent-test.sh new file mode 100644 index 00000000000..dd8be30d0f6 --- /dev/null +++ b/test/shell.d/mcode-agent-test.sh @@ -0,0 +1,85 @@ +#!/bin/bash + +set -euo pipefail +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" + +test_tmp=$(mktemp -d) +trap 'rm -rf "$test_tmp"' EXIT +mock_bin="$test_tmp/bin" +export HOME="$test_tmp/home with spaces" +export PATH="$mock_bin:$HOME/.local/bin:$ROOT/bin:/usr/local/bin:/usr/bin" +export OMARCHY_PATH="$ROOT" +export OMARCHY_TEST_MISE_LOG="$test_tmp/mise-log" +mkdir -p "$mock_bin" "$HOME/.local/bin" +mcode="$HOME/.local/bin/mcode" +package="npm:@minimax-ai/code[allow_builds=better-sqlite3]" + +cat >"$mock_bin/mise" <<'SH' +#!/bin/bash +printf '%s\0' "$@" >>"$OMARCHY_TEST_MISE_LOG" +SH +cat >"$mock_bin/noop" <<'SH' +#!/bin/bash +exit 0 +SH +chmod +x "$mock_bin/mise" "$mock_bin/noop" +for command in gum hyprctl omarchy-webapp-remove-all omarchy-tui-remove-all omarchy-pkg-drop; do + ln -s noop "$mock_bin/$command" +done + +apply_mcode() { + case $1 in + migration) source "$ROOT/migrations/1792033557.sh" >/dev/null ;; + removal) omarchy-remove-preinstalls >/dev/null ;; + esac +} + +apply_mcode migration +[[ -x $mcode && ! -L $mcode ]] || fail "MiniMax Code migration creates a missing wrapper" +[[ ! -s $OMARCHY_TEST_MISE_LOG ]] || fail "MiniMax Code migration installs the package before first launch" +cp "$mcode" "$test_tmp/wrapper" +apply_mcode migration +cmp -s "$mcode" "$test_tmp/wrapper" || fail "MiniMax Code migration is idempotent" +"$mcode" --prompt 'Review this project' +mapfile -d '' -t args <"$OMARCHY_TEST_MISE_LOG" +[[ ${args[*]} == "use -g --quiet $package x $package -- mcode --prompt Review this project" ]] || + fail "MiniMax Code migration wrapper installs the official package and forwards arguments" +pass "MiniMax Code migration creates a lazy, repeatable wrapper" + +rm -f "$mcode" +mkdir -p "$HOME/.local/state/omarchy" +touch "$HOME/.local/state/omarchy/preinstalls-removed" +apply_mcode migration +[[ ! -e $mcode ]] || fail "MiniMax Code migration respects removed preinstalls with a missing command" +rm "$HOME/.local/state/omarchy/preinstalls-removed" +pass "MiniMax Code migration respects the preinstall opt-out" + +for operation in migration removal; do + for kind in script binary symlink; do + rm -f "$mcode" + case $kind in + script) printf '#!/bin/bash\necho user-managed mcode\n' >"$mcode" ;; + binary) cp /usr/bin/true "$mcode" ;; + symlink) ln -s /usr/bin/true "$mcode" ;; + esac + [[ $kind == "symlink" ]] || chmod +x "$mcode" + rm -f "$test_tmp/saved-$kind" + cp -P "$mcode" "$test_tmp/saved-$kind" + for attempt in 1 2; do + apply_mcode "$operation" + cmp -s "$mcode" "$test_tmp/saved-$kind" || fail "MiniMax Code $operation preserves a user $kind" + if [[ $kind == "symlink" ]]; then + [[ -L $mcode && $(readlink "$mcode") == "/usr/bin/true" ]] || + fail "MiniMax Code $operation preserves the user's symlink" + else + [[ ! -L $mcode ]] || fail "MiniMax Code $operation keeps the user's $kind as a regular file" + fi + done + pass "MiniMax Code $operation preserves a user $kind on repeated runs" + done +done + +omarchy-mise-install "$package" mcode +apply_mcode removal +[[ ! -e $mcode ]] || fail "Remove Preinstalls removes Omarchy's MiniMax Code wrapper" +pass "Remove Preinstalls removes Omarchy's MiniMax Code wrapper" diff --git a/test/shell.d/menu-test.sh b/test/shell.d/menu-test.sh index addf7e676d8..f43d136fe8c 100644 --- a/test/shell.d/menu-test.sh +++ b/test/shell.d/menu-test.sh @@ -255,6 +255,7 @@ const expectedAgents = { kimi: { icon: '\u{f06a9}', label: 'Kimi' }, qwen: { icon: '\ue90f', iconFont: 'omarchy', label: 'Qwen Code' }, dsh: { icon: '\ue910', iconFont: 'omarchy', label: 'DeepSeek Harness' }, + mcode: { icon: '\ue911', iconFont: 'omarchy', label: 'MiniMax Code' }, hermes: { icon: '\ue90a', iconFont: 'omarchy', label: 'Hermes' }, openclaw: { icon: '\ue90c', iconFont: 'omarchy', label: 'OpenClaw' }, copilot: { icon: '', label: 'Copilot' }, @@ -280,7 +281,7 @@ assertDeepEqual( defaultItems .filter(item => item.parent === 'setup.default.agent') .map(item => item.label), - ['Antigravity', 'Claude', 'Codex', 'Copilot', 'Crush', 'Cursor CLI', 'DeepSeek Harness', 'Grok', 'Hermes', 'Kimi', 'Muse Code', 'omp', 'OpenClaw', 'OpenCode', 'Ori', 'Pi', 'Qwen Code'], + ['Antigravity', 'Claude', 'Codex', 'Copilot', 'Crush', 'Cursor CLI', 'DeepSeek Harness', 'Grok', 'Hermes', 'Kimi', 'MiniMax Code', 'Muse Code', 'omp', 'OpenClaw', 'OpenCode', 'Ori', 'Pi', 'Qwen Code'], 'menu sorts coding agents alphabetically' ) const expectedDefaults = { @@ -877,5 +878,5 @@ assert( JS font_charset=$(fc-query --format='%{charset}' "$ROOT/default/fonts/omarchy/omarchy.ttf") -[[ $font_charset == *"e900-e910"* ]] || fail "Omarchy icon font includes every custom menu glyph" +[[ $font_charset == *"e900-e911"* ]] || fail "Omarchy icon font includes every custom menu glyph" pass "Omarchy icon font includes the official agent marks" diff --git a/test/shell.d/mise-provision-test.sh b/test/shell.d/mise-provision-test.sh index 66d727f91bd..1a6d6b55233 100644 --- a/test/shell.d/mise-provision-test.sh +++ b/test/shell.d/mise-provision-test.sh @@ -17,7 +17,7 @@ SH chmod +x "$mock_bin/mise" source "$ROOT/install/user/mise.sh" -for agent in kimi cursor-agent muse qwen dsh; do +for agent in kimi cursor-agent muse qwen dsh mcode; do [[ -x $HOME/.local/bin/$agent && ! -L $HOME/.local/bin/$agent ]] || fail "provisioning creates the missing $agent wrapper" cp "$HOME/.local/bin/$agent" "$test_tmp/$agent-wrapper" @@ -27,7 +27,7 @@ mapfile -d '' -t args <"$OMARCHY_TEST_MISE_LOG" fail "provisioning creates lazy wrappers without installing packages" pass "provisioning creates lazy Kimi, Cursor, and Muse wrappers" -for agent in kimi cursor-agent muse qwen dsh; do +for agent in kimi cursor-agent muse qwen dsh mcode; do launcher="$HOME/.local/bin/$agent" for kind in script binary symlink; do rm -f "$launcher" From 9f397fd6e27368c422343b424deb54dd72401e3a Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:09:48 -0400 Subject: [PATCH 099/546] Add Goose as a selectable default coding agent MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Goose (goose) joins the default agent picker, installed through mise as aqua:aaif-goose/goose, which picks the glibc release archive on x86_64 and aarch64 and checks its attestations. Goose has no bypass flag, so the launcher sets GOOSE_MODE=auto, and seeds a prompt with goose run --interactive --text= so the session stays open. From #14492. Co-authored-by: Kyle 🐆 Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent | 10 ++++ bin/omarchy-default-agent | 5 +- bin/omarchy-remove-preinstalls | 6 ++ default/fonts/omarchy/README.md | 1 + default/fonts/omarchy/omarchy.ttf | Bin 11308 -> 12000 bytes default/i18n/en/menu.json | 3 + default/omarchy/omarchy-menu.jsonc | 1 + install/user/mise.sh | 1 + manual/17-ai.md | 1 + migrations/1792033725.sh | 5 ++ test/shell.d/default-agent-test.sh | 16 +++++- test/shell.d/goose-agent-test.sh | 85 ++++++++++++++++++++++++++++ test/shell.d/menu-test.sh | 5 +- test/shell.d/mise-provision-test.sh | 4 +- 14 files changed, 136 insertions(+), 7 deletions(-) create mode 100644 migrations/1792033725.sh create mode 100644 test/shell.d/goose-agent-test.sh diff --git a/bin/omarchy-agent b/bin/omarchy-agent index b67d7776ac1..a53e52d4fd3 100755 --- a/bin/omarchy-agent +++ b/bin/omarchy-agent @@ -158,6 +158,16 @@ mcode) command=(mcode) [[ -n ${prompt:-} ]] && command+=(-- "$prompt") ;; +goose) + # Goose has no bypass flag; GOOSE_MODE outranks a configured approve mode. + # --text= binds a prompt that starts with a dash, and --interactive keeps + # the session open after the prompted turn. + if [[ -n ${prompt:-} ]]; then + command=(env GOOSE_MODE=auto goose run --interactive "--text=$prompt") + else + command=(env GOOSE_MODE=auto goose) + fi + ;; *) echo "Unsupported default agent: $agent" >&2 exit 1 diff --git a/bin/omarchy-default-agent b/bin/omarchy-default-agent index 2f69bb7b224..abaedb831d7 100755 --- a/bin/omarchy-default-agent +++ b/bin/omarchy-default-agent @@ -1,7 +1,7 @@ #!/bin/bash # omarchy:summary=Set and launch the default coding agent -# omarchy:args=[pi|omp|opencode|ori|claude|codex|grok|openclaw|agy|hermes|copilot|crush|cursor-agent|muse|kimi|qwen|dsh|mcode] +# omarchy:args=[pi|omp|opencode|ori|claude|codex|grok|openclaw|agy|hermes|copilot|crush|cursor-agent|muse|kimi|qwen|dsh|mcode|goose] # omarchy:examples=omarchy default agent | omarchy default agent codex | omarchy default agent claude installing=false @@ -52,8 +52,9 @@ dsh | deepseek | deepseek-harness) agent_package="npm:@deepseek-ai/dsh@0.2.0-rc.2" ;; mcode | minimax-code | minimax) agent="mcode"; name="MiniMax Code"; agent_package="npm:@minimax-ai/code[allow_builds=better-sqlite3]" ;; +goose) agent="goose"; name="Goose"; agent_package="aqua:aaif-goose/goose" ;; *) - echo "Usage: omarchy-default-agent " + echo "Usage: omarchy-default-agent " exit 1 ;; esac diff --git a/bin/omarchy-remove-preinstalls b/bin/omarchy-remove-preinstalls index be3f3e9cf4c..c8aea7ca9a7 100755 --- a/bin/omarchy-remove-preinstalls +++ b/bin/omarchy-remove-preinstalls @@ -54,6 +54,12 @@ if gum confirm --affirmative "$(omarchy-i18n common.yes)" --negative "$(omarchy- rm -f ~/.local/bin/mcode fi + # Preserve a user-managed Goose launcher at the same path. + if [[ -f ~/.local/bin/goose && ! -L ~/.local/bin/goose ]] && + grep -Eq '^mise use -g .*"aqua:aaif-goose/goose"' ~/.local/bin/goose; then + rm -f ~/.local/bin/goose + fi + omarchy-pkg-drop \ aether \ cliamp \ diff --git a/default/fonts/omarchy/README.md b/default/fonts/omarchy/README.md index 23543c388ab..8e2c4ddbec4 100644 --- a/default/fonts/omarchy/README.md +++ b/default/fonts/omarchy/README.md @@ -20,6 +20,7 @@ The private-use glyphs in `omarchy.ttf` are: - `U+E90F` — Qwen, from - `U+E910` — DeepSeek, from - `U+E911` — MiniMax, from +- `U+E912` — Goose, from The agent marks are monochrome so the menu can render them using the active theme's foreground and selection colors. diff --git a/default/fonts/omarchy/omarchy.ttf b/default/fonts/omarchy/omarchy.ttf index d11f8d92deda4cdc28f446645e16af8973e427af..266840c860f5bc16914bdae5098a17064acad305 100644 GIT binary patch delta 937 zcmX|NucY8Ec^PZYXQT)1*4f-6A~&rDR#F!y}teJ=OSodEjo)L1mNeJn#qrm`PMqWogy1se)&Uqr2k)I-OY;0Xy zJq2fn$Ul+iH&^d$Lyux^HOrs8wYs_1e)9PN@&!Vyz3r`?+je#MQV}x4G2r9ghj9Qs z_ypff9DhWsL=7AdJ@qn?S1Vx*CdB;)W%ZlCM+ka7>2 zGi@X-*=V+!je4)s>;9Ge$lQ5iqu#C>e+(V>o`C~`epfnri(Y{4)c->=usG^eIgFbkG# znt?}ok_e+Q$v_EL0h>C>`FL_jgAjdSlg9)d3$zJEo6GqV3zz1K1#Ou|(GLR86Jf}h zGy~g#WI`zbOp-!Mcg&VZjCsoBJWZwR`P%iQ1nQ+1JGJ?9Z-uam0A`@an3%3G5FEfy z4vlj@GsO%C)9K-UL-90W)YGT6s~wylKOq@zWP$6I@M8Qhh@)D)la|XVEfgHPVCoX5 zww53%vz#<#fcsNi-=NcV2EHT0Ahdwag;Ff`YMqtT#p}ri^jw*gIX9)pF-~)V7j8&g z*0MpF1>rb>p%@E7;dw#e`-WlZLb1EDFw?_#i_aI@?da?A_k$M`%Q8hKe6UX|1EaFY o@wM*^&R4pIyPJom!^f5H7NmoF(;W)oV0ZeyDQ<3U?W|$iALn0xTL1t6 delta 229 zcmaD5yC!0SVST=!Cj$dx2Ll6>S9(rmTC#Gy1_J~06$S>WEg7kaDSGRy`WP7GfbtAx z89)IJMiyxx{{WD$l95|d;laRq2grW{hbe2eidry!6o3UnCT z=8H@$jGM18Z/dev/null || fail "user setup creates grep -Fx "$qwen_package qwen" "$stub_log" >/dev/null || fail "user setup creates the Qwen Code lazy stub" grep -Fx "$dsh_package dsh" "$stub_log" >/dev/null || fail "user setup creates the DeepSeek Harness lazy stub" grep -Fx "$mcode_package mcode" "$stub_log" >/dev/null || fail "user setup creates the MiniMax Code lazy stub" +grep -Fx "$goose_package goose" "$stub_log" >/dev/null || fail "user setup creates the Goose lazy stub" pass "user setup creates the custom agent lazy stubs" : >"$stub_log" @@ -185,6 +188,7 @@ grep -Fx "$kimi_package kimi" "$stub_log" >/dev/null || fail "user setup recreat grep -Fx "$qwen_package qwen" "$stub_log" >/dev/null || fail "user setup recreates the Qwen Code lazy stub" grep -Fx "$dsh_package dsh" "$stub_log" >/dev/null || fail "user setup recreates the DeepSeek Harness lazy stub" grep -Fx "$mcode_package mcode" "$stub_log" >/dev/null || fail "user setup recreates the MiniMax Code lazy stub" +grep -Fx "$goose_package goose" "$stub_log" >/dev/null || fail "user setup recreates the Goose lazy stub" pass "user setup recreates wrappers even when the commands are present" : >"$stub_log" @@ -217,6 +221,10 @@ grep -Fx "$dsh_package dsh" "$stub_log" >/dev/null || fail "DeepSeek Harness mig OMARCHY_TEST_MISSING_COMMAND=mcode source "$ROOT/migrations/1792033557.sh" >/dev/null grep -Fx "$mcode_package mcode" "$stub_log" >/dev/null || fail "MiniMax Code migration creates its lazy stub" +: >"$stub_log" +OMARCHY_TEST_MISSING_COMMAND=goose source "$ROOT/migrations/1792033725.sh" >/dev/null +grep -Fx "$goose_package goose" "$stub_log" >/dev/null || fail "Goose migration creates its lazy stub" + : >"$stub_log" source "$ROOT/migrations/1785617047.sh" >/dev/null grep -Fx "$omp_package omp" "$stub_log" >/dev/null || fail "Oh My Pi migration creates a working lazy stub" @@ -323,6 +331,7 @@ OMARCHY_TEST_MISSING_COMMAND=kimi source "$ROOT/migrations/1786967403.sh" >/dev/ OMARCHY_TEST_MISSING_COMMAND=qwen source "$ROOT/migrations/1792032538.sh" >/dev/null OMARCHY_TEST_MISSING_COMMAND=dsh source "$ROOT/migrations/1792033356.sh" >/dev/null OMARCHY_TEST_MISSING_COMMAND=mcode source "$ROOT/migrations/1792033557.sh" >/dev/null +OMARCHY_TEST_MISSING_COMMAND=goose source "$ROOT/migrations/1792033725.sh" >/dev/null [[ ! -s $stub_log ]] || fail "agent migrations respect the preinstall opt-out" [[ ! -e $test_home/.local/bin/omp ]] || fail "agent migration removes the obsolete Oh My Pi wrapper after opt-out" @@ -459,9 +468,10 @@ pass "agent migrations install working wrappers without overriding the preinstal "$ROOT/bin/omarchy-mise-install" "$qwen_package" qwen "$ROOT/bin/omarchy-mise-install" "$dsh_package" dsh "$ROOT/bin/omarchy-mise-install" "$mcode_package" mcode +"$ROOT/bin/omarchy-mise-install" "$goose_package" goose touch "$test_home/.local/bin/agy" "$test_home/.local/bin/ori" omarchy-remove-preinstalls >/dev/null -for command in agy omp ori grok crush cursor-agent muse kimi qwen dsh mcode; do +for command in agy omp ori grok crush cursor-agent muse kimi qwen dsh mcode goose; do [[ ! -e $test_home/.local/bin/$command ]] || fail "Remove Preinstalls deletes the $command lazy stub" done pass "Remove Preinstalls deletes every optional agent lazy stub" @@ -552,6 +562,7 @@ declare -A expected_agents=( [mcode]="mcode" [minimax-code]="mcode" [minimax]="mcode" + [goose]="goose" [cursor]="cursor-agent" [cursor-agent]="cursor-agent" [muse]="muse" @@ -574,6 +585,7 @@ declare -A expected_packages=( [qwen]="$qwen_package" [dsh]="$dsh_package" [mcode]="$mcode_package" + [goose]="$goose_package" [cursor-agent]="$cursor_agent_package" [muse]="$muse_package" ) @@ -853,6 +865,7 @@ printf '%s\n' dsh >"$agent_file" omarchy-agent-prompt '-n --help; echo literal' assert_launched dsh "keeps an option-like prompt literal" env DSH_PERMISSION_MODE=danger-full-access dsh headless -- -- '-n --help; echo literal' assert_launch mcode mcode -- "Review this project" +assert_launch goose env GOOSE_MODE=auto goose run --interactive "--text=Review this project" pass "agent launcher adapts initial prompts for every supported agent" literal_muse_prompt=$'--disable-sandbox !Crash {$(touch must-not-run)}\ntrailing\\ ' @@ -885,6 +898,7 @@ assert_bypass kimi kimi --yolo assert_bypass qwen qwen --yolo assert_bypass dsh env DSH_PERMISSION_MODE=danger-full-access dsh web assert_bypass mcode mcode +assert_bypass goose env GOOSE_MODE=auto goose pass "agent launcher skips permission prompts for every supported agent" printf '%s\n' "opencode" >"$agent_file" diff --git a/test/shell.d/goose-agent-test.sh b/test/shell.d/goose-agent-test.sh new file mode 100644 index 00000000000..72b1f1018eb --- /dev/null +++ b/test/shell.d/goose-agent-test.sh @@ -0,0 +1,85 @@ +#!/bin/bash + +set -euo pipefail +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" + +test_tmp=$(mktemp -d) +trap 'rm -rf "$test_tmp"' EXIT +mock_bin="$test_tmp/bin" +export HOME="$test_tmp/home with spaces" +export PATH="$mock_bin:$HOME/.local/bin:$ROOT/bin:/usr/local/bin:/usr/bin" +export OMARCHY_PATH="$ROOT" +export OMARCHY_TEST_MISE_LOG="$test_tmp/mise-log" +mkdir -p "$mock_bin" "$HOME/.local/bin" +goose="$HOME/.local/bin/goose" +package="aqua:aaif-goose/goose" + +cat >"$mock_bin/mise" <<'SH' +#!/bin/bash +printf '%s\0' "$@" >>"$OMARCHY_TEST_MISE_LOG" +SH +cat >"$mock_bin/noop" <<'SH' +#!/bin/bash +exit 0 +SH +chmod +x "$mock_bin/mise" "$mock_bin/noop" +for command in gum hyprctl omarchy-webapp-remove-all omarchy-tui-remove-all omarchy-pkg-drop; do + ln -s noop "$mock_bin/$command" +done + +apply_goose() { + case $1 in + migration) source "$ROOT/migrations/1792033725.sh" >/dev/null ;; + removal) omarchy-remove-preinstalls >/dev/null ;; + esac +} + +apply_goose migration +[[ -x $goose && ! -L $goose ]] || fail "Goose migration creates a missing wrapper" +[[ ! -s $OMARCHY_TEST_MISE_LOG ]] || fail "Goose migration installs the package before first launch" +cp "$goose" "$test_tmp/wrapper" +apply_goose migration +cmp -s "$goose" "$test_tmp/wrapper" || fail "Goose migration is idempotent" +"$goose" --prompt 'Review this project' +mapfile -d '' -t args <"$OMARCHY_TEST_MISE_LOG" +[[ ${args[*]} == "use -g --quiet $package x $package -- goose --prompt Review this project" ]] || + fail "Goose migration wrapper installs the official package and forwards arguments" +pass "Goose migration creates a lazy, repeatable wrapper" + +rm -f "$goose" +mkdir -p "$HOME/.local/state/omarchy" +touch "$HOME/.local/state/omarchy/preinstalls-removed" +apply_goose migration +[[ ! -e $goose ]] || fail "Goose migration respects removed preinstalls with a missing command" +rm "$HOME/.local/state/omarchy/preinstalls-removed" +pass "Goose migration respects the preinstall opt-out" + +for operation in migration removal; do + for kind in script binary symlink; do + rm -f "$goose" + case $kind in + script) printf '#!/bin/bash\necho user-managed goose\n' >"$goose" ;; + binary) cp /usr/bin/true "$goose" ;; + symlink) ln -s /usr/bin/true "$goose" ;; + esac + [[ $kind == "symlink" ]] || chmod +x "$goose" + rm -f "$test_tmp/saved-$kind" + cp -P "$goose" "$test_tmp/saved-$kind" + for attempt in 1 2; do + apply_goose "$operation" + cmp -s "$goose" "$test_tmp/saved-$kind" || fail "Goose $operation preserves a user $kind" + if [[ $kind == "symlink" ]]; then + [[ -L $goose && $(readlink "$goose") == "/usr/bin/true" ]] || + fail "Goose $operation preserves the user's symlink" + else + [[ ! -L $goose ]] || fail "Goose $operation keeps the user's $kind as a regular file" + fi + done + pass "Goose $operation preserves a user $kind on repeated runs" + done +done + +omarchy-mise-install "$package" goose +apply_goose removal +[[ ! -e $goose ]] || fail "Remove Preinstalls removes Omarchy's Goose wrapper" +pass "Remove Preinstalls removes Omarchy's Goose wrapper" diff --git a/test/shell.d/menu-test.sh b/test/shell.d/menu-test.sh index f43d136fe8c..85f6e5c1f1e 100644 --- a/test/shell.d/menu-test.sh +++ b/test/shell.d/menu-test.sh @@ -256,6 +256,7 @@ const expectedAgents = { qwen: { icon: '\ue90f', iconFont: 'omarchy', label: 'Qwen Code' }, dsh: { icon: '\ue910', iconFont: 'omarchy', label: 'DeepSeek Harness' }, mcode: { icon: '\ue911', iconFont: 'omarchy', label: 'MiniMax Code' }, + goose: { icon: '\ue912', iconFont: 'omarchy', label: 'Goose' }, hermes: { icon: '\ue90a', iconFont: 'omarchy', label: 'Hermes' }, openclaw: { icon: '\ue90c', iconFont: 'omarchy', label: 'OpenClaw' }, copilot: { icon: '', label: 'Copilot' }, @@ -281,7 +282,7 @@ assertDeepEqual( defaultItems .filter(item => item.parent === 'setup.default.agent') .map(item => item.label), - ['Antigravity', 'Claude', 'Codex', 'Copilot', 'Crush', 'Cursor CLI', 'DeepSeek Harness', 'Grok', 'Hermes', 'Kimi', 'MiniMax Code', 'Muse Code', 'omp', 'OpenClaw', 'OpenCode', 'Ori', 'Pi', 'Qwen Code'], + ['Antigravity', 'Claude', 'Codex', 'Copilot', 'Crush', 'Cursor CLI', 'DeepSeek Harness', 'Goose', 'Grok', 'Hermes', 'Kimi', 'MiniMax Code', 'Muse Code', 'omp', 'OpenClaw', 'OpenCode', 'Ori', 'Pi', 'Qwen Code'], 'menu sorts coding agents alphabetically' ) const expectedDefaults = { @@ -878,5 +879,5 @@ assert( JS font_charset=$(fc-query --format='%{charset}' "$ROOT/default/fonts/omarchy/omarchy.ttf") -[[ $font_charset == *"e900-e911"* ]] || fail "Omarchy icon font includes every custom menu glyph" +[[ $font_charset == *"e900-e912"* ]] || fail "Omarchy icon font includes every custom menu glyph" pass "Omarchy icon font includes the official agent marks" diff --git a/test/shell.d/mise-provision-test.sh b/test/shell.d/mise-provision-test.sh index 1a6d6b55233..7fd2edff58d 100644 --- a/test/shell.d/mise-provision-test.sh +++ b/test/shell.d/mise-provision-test.sh @@ -17,7 +17,7 @@ SH chmod +x "$mock_bin/mise" source "$ROOT/install/user/mise.sh" -for agent in kimi cursor-agent muse qwen dsh mcode; do +for agent in kimi cursor-agent muse qwen dsh mcode goose; do [[ -x $HOME/.local/bin/$agent && ! -L $HOME/.local/bin/$agent ]] || fail "provisioning creates the missing $agent wrapper" cp "$HOME/.local/bin/$agent" "$test_tmp/$agent-wrapper" @@ -27,7 +27,7 @@ mapfile -d '' -t args <"$OMARCHY_TEST_MISE_LOG" fail "provisioning creates lazy wrappers without installing packages" pass "provisioning creates lazy Kimi, Cursor, and Muse wrappers" -for agent in kimi cursor-agent muse qwen dsh mcode; do +for agent in kimi cursor-agent muse qwen dsh mcode goose; do launcher="$HOME/.local/bin/$agent" for kind in script binary symlink; do rm -f "$launcher" From fd3352026ba58d282e4a2423291817bf7668cbc0 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:11:20 -0400 Subject: [PATCH 100/546] Add Kiro CLI as a selectable default coding agent Kiro CLI (kiro-cli) joins the default agent picker. Instead of piping Kiro's install script, mise's HTTP backend installs the per-platform release archive for x86_64 or aarch64, resolving the version from Kiro's stable manifest, with a lazy stub and migration like the other agents. The launcher starts kiro-cli chat with --trust-all-tools and seeds a prompt as its first input after --. Combines #15481 and #7693. Co-authored-by: Bradley Marks Co-authored-by: kenth-dev Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent | 6 ++ bin/omarchy-default-agent | 9 ++- bin/omarchy-remove-preinstalls | 6 ++ default/fonts/omarchy/README.md | 1 + default/fonts/omarchy/omarchy.ttf | Bin 12000 -> 12268 bytes default/i18n/en/menu.json | 3 + default/omarchy/omarchy-menu.jsonc | 1 + install/user/mise.sh | 1 + manual/17-ai.md | 1 + migrations/1792033795.sh | 5 ++ test/shell.d/default-agent-test.sh | 17 +++++- test/shell.d/kiro-cli-agent-test.sh | 85 ++++++++++++++++++++++++++++ test/shell.d/menu-test.sh | 5 +- test/shell.d/mise-provision-test.sh | 4 +- 14 files changed, 137 insertions(+), 7 deletions(-) create mode 100644 migrations/1792033795.sh create mode 100644 test/shell.d/kiro-cli-agent-test.sh diff --git a/bin/omarchy-agent b/bin/omarchy-agent index a53e52d4fd3..67fd3ffc62a 100755 --- a/bin/omarchy-agent +++ b/bin/omarchy-agent @@ -168,6 +168,12 @@ goose) command=(env GOOSE_MODE=auto goose) fi ;; +kiro-cli) + # A positional input seeds the interactive chat; -- keeps a prompt starting + # with a dash literal. + command=(kiro-cli chat --trust-all-tools) + [[ -n ${prompt:-} ]] && command+=(-- "$prompt") + ;; *) echo "Unsupported default agent: $agent" >&2 exit 1 diff --git a/bin/omarchy-default-agent b/bin/omarchy-default-agent index abaedb831d7..1c62689c849 100755 --- a/bin/omarchy-default-agent +++ b/bin/omarchy-default-agent @@ -1,7 +1,7 @@ #!/bin/bash # omarchy:summary=Set and launch the default coding agent -# omarchy:args=[pi|omp|opencode|ori|claude|codex|grok|openclaw|agy|hermes|copilot|crush|cursor-agent|muse|kimi|qwen|dsh|mcode|goose] +# omarchy:args=[pi|omp|opencode|ori|claude|codex|grok|openclaw|agy|hermes|copilot|crush|cursor-agent|muse|kimi|qwen|dsh|mcode|goose|kiro-cli] # omarchy:examples=omarchy default agent | omarchy default agent codex | omarchy default agent claude installing=false @@ -53,8 +53,13 @@ dsh | deepseek | deepseek-harness) ;; mcode | minimax-code | minimax) agent="mcode"; name="MiniMax Code"; agent_package="npm:@minimax-ai/code[allow_builds=better-sqlite3]" ;; goose) agent="goose"; name="Goose"; agent_package="aqua:aaif-goose/goose" ;; +kiro-cli | kiro) + agent="kiro-cli"; name="Kiro CLI" + # Kiro publishes per-platform archives behind a manifest of the stable release. + agent_package="http:kiro-cli[platforms_linux_x64_url=https://prod.download.cli.kiro.dev/stable/{{version}}/kirocli-x86_64-linux.tar.gz,platforms_linux_arm64_url=https://prod.download.cli.kiro.dev/stable/{{version}}/kirocli-aarch64-linux.tar.gz,version_list_url=https://prod.download.cli.kiro.dev/stable/latest/manifest.json,version_json_path=.version]" + ;; *) - echo "Usage: omarchy-default-agent " + echo "Usage: omarchy-default-agent " exit 1 ;; esac diff --git a/bin/omarchy-remove-preinstalls b/bin/omarchy-remove-preinstalls index c8aea7ca9a7..c31fd37db3e 100755 --- a/bin/omarchy-remove-preinstalls +++ b/bin/omarchy-remove-preinstalls @@ -60,6 +60,12 @@ if gum confirm --affirmative "$(omarchy-i18n common.yes)" --negative "$(omarchy- rm -f ~/.local/bin/goose fi + # Preserve a user-managed Kiro CLI launcher at the same path. + if [[ -f ~/.local/bin/kiro-cli && ! -L ~/.local/bin/kiro-cli ]] && + grep -Eq '^mise use -g .*"http:kiro-cli\[' ~/.local/bin/kiro-cli; then + rm -f ~/.local/bin/kiro-cli + fi + omarchy-pkg-drop \ aether \ cliamp \ diff --git a/default/fonts/omarchy/README.md b/default/fonts/omarchy/README.md index 8e2c4ddbec4..469011a7bed 100644 --- a/default/fonts/omarchy/README.md +++ b/default/fonts/omarchy/README.md @@ -21,6 +21,7 @@ The private-use glyphs in `omarchy.ttf` are: - `U+E910` — DeepSeek, from - `U+E911` — MiniMax, from - `U+E912` — Goose, from +- `U+E913` — Kiro, from the ghost in , with its tile dropped The agent marks are monochrome so the menu can render them using the active theme's foreground and selection colors. diff --git a/default/fonts/omarchy/omarchy.ttf b/default/fonts/omarchy/omarchy.ttf index 266840c860f5bc16914bdae5098a17064acad305..3f6a488be697bb7111405e8ce3578482688c85ea 100644 GIT binary patch delta 501 zcmaD5`zC&ZVSS;XHv2Ll6>S9(rm+DZ3}Zww5~R~Q)NVlq+_Q}otZ^)WCg889#~ zm}LM3I2c)Ef&2_0UnL{Aq#}TU?In=k1LPOvS3>N{-)^8Wz&R=J541&kIu}_DvwdZKt_g}n~{-|hm(Vqi-m=Sjg^H*J~${PC8R!7nT?%8 zNJv*qMOj^4QBlC$!dy#>msgRIMT?1#TSSPPhmnIrYCj_{kFYQon9ay7$jT}xBrYZ- z&M&|(B_bkWWT3BAX=`I+Qz2wv;OY_>7^y;46`QxP_*FU*8|A|vCL*OrT5%G VX*E4FCPvZ8F?x-xEZLbw`2a<7TWtUU delta 229 zcmaD8{~&gPVSRz17Xt%h2Ll6>S9(rm+Nb(|LJSPdR~Q&%V=__`Q}otZ^)WCg8Za<0 zm}LM3I2c)Efcy*~UnL{Aq{4@R%^%3`0rGQl@{<#-7@al)`9FYsi`>MD0!DKtpk5`A zeuccm+*IT4JGDT500RT}#)ACf5+0?^v)CCJ^jLrjZf@Ko%g87+`4;0_PM{+gM1c-t z+kBCUg>my0=2%&7MxYp@v9Z`Qo!ygdl>8Z`C-*4jO#Y)}!NvfRlHP2m{GJ;qs;6%@ LIZeNj87Kw-K_EDw diff --git a/default/i18n/en/menu.json b/default/i18n/en/menu.json index f836ffc394a..7eefcceb166 100644 --- a/default/i18n/en/menu.json +++ b/default/i18n/en/menu.json @@ -515,6 +515,9 @@ "setup.default.agent.kimi": { "label": "Kimi" }, + "setup.default.agent.kiro-cli": { + "label": "Kiro CLI" + }, "setup.default.agent.mcode": { "label": "MiniMax Code" }, diff --git a/default/omarchy/omarchy-menu.jsonc b/default/omarchy/omarchy-menu.jsonc index 2b74273f560..e773403e808 100644 --- a/default/omarchy/omarchy-menu.jsonc +++ b/default/omarchy/omarchy-menu.jsonc @@ -201,6 +201,7 @@ "setup.default.agent.grok": {"icon":"","iconFont":"omarchy","label":"Grok","checked":"[[ \"$(omarchy-default-agent)\" == \"grok\" ]]","action":"omarchy-default-agent grok"}, "setup.default.agent.hermes": {"icon":"","iconFont":"omarchy","label":"Hermes","checked":"[[ \"$(omarchy-default-agent)\" == \"hermes\" ]]","action":"omarchy-default-agent hermes"}, "setup.default.agent.kimi": {"icon":"󰚩","label":"Kimi","checked":"[[ \"$(omarchy-default-agent)\" == \"kimi\" ]]","action":"omarchy-default-agent kimi"}, + "setup.default.agent.kiro-cli": {"icon":"","iconFont":"omarchy","label":"Kiro CLI","checked":"[[ \"$(omarchy-default-agent)\" == \"kiro-cli\" ]]","action":"omarchy-default-agent kiro-cli"}, "setup.default.agent.mcode": {"icon":"","iconFont":"omarchy","label":"MiniMax Code","checked":"[[ \"$(omarchy-default-agent)\" == \"mcode\" ]]","action":"omarchy-default-agent mcode"}, "setup.default.agent.muse": {"icon":"󰛤","label":"Muse Code","checked":"[[ \"$(omarchy-default-agent)\" == \"muse\" ]]","action":"omarchy-default-agent muse"}, "setup.default.agent.omp": {"icon":"","iconFont":"omarchy","label":"omp","checked":"[[ \"$(omarchy-default-agent)\" == \"omp\" ]]","action":"omarchy-default-agent omp"}, diff --git a/install/user/mise.sh b/install/user/mise.sh index 3c7a5696fd4..c4152688761 100644 --- a/install/user/mise.sh +++ b/install/user/mise.sh @@ -19,6 +19,7 @@ omarchy-mise-install npm:@qwen-code/qwen-code qwen omarchy-mise-install npm:@deepseek-ai/dsh@0.2.0-rc.2 dsh omarchy-mise-install "npm:@minimax-ai/code[allow_builds=better-sqlite3]" mcode omarchy-mise-install aqua:aaif-goose/goose goose +omarchy-mise-install "http:kiro-cli[platforms_linux_x64_url=https://prod.download.cli.kiro.dev/stable/{{version}}/kirocli-x86_64-linux.tar.gz,platforms_linux_arm64_url=https://prod.download.cli.kiro.dev/stable/{{version}}/kirocli-aarch64-linux.tar.gz,version_list_url=https://prod.download.cli.kiro.dev/stable/latest/manifest.json,version_json_path=.version]" kiro-cli omarchy-mise-install cursor-agent omarchy-mise-install npm:@kitlangton/ghui ghui omarchy-mise-install aqua:modem-dev/hunk hunk diff --git a/manual/17-ai.md b/manual/17-ai.md index dac5e78b53f..9ce762ce2e1 100644 --- a/manual/17-ai.md +++ b/manual/17-ai.md @@ -14,6 +14,7 @@ Omarchy treats AI coding agents as first-class citizens, but it doesn't pick a f | `goose` | [Goose](https://github.com/aaif-goose/goose) | | `grok` | [Grok Build](https://x.ai/build) from xAI | | `kimi` | [Kimi Code](https://www.kimi.com/code/docs/en/) | +| `kiro-cli` | [Kiro CLI](https://kiro.dev/docs/cli/) | | `mcode` | [MiniMax Code](https://github.com/MiniMax-AI/minimax-code), MiniMax's coding agent | | `qwen` | [Qwen Code](https://github.com/QwenLM/qwen-code), Alibaba's coding agent | | `pi` | [Mario Zechner's Pi](https://github.com/badlogic/pi-mono) | diff --git a/migrations/1792033795.sh b/migrations/1792033795.sh new file mode 100644 index 00000000000..793336ffc75 --- /dev/null +++ b/migrations/1792033795.sh @@ -0,0 +1,5 @@ +echo "Install Kiro CLI coding agent mise wrapper" + +if [[ ! -f $HOME/.local/state/omarchy/preinstalls-removed ]] && omarchy-cmd-missing kiro-cli; then + omarchy-mise-install "http:kiro-cli[platforms_linux_x64_url=https://prod.download.cli.kiro.dev/stable/{{version}}/kirocli-x86_64-linux.tar.gz,platforms_linux_arm64_url=https://prod.download.cli.kiro.dev/stable/{{version}}/kirocli-aarch64-linux.tar.gz,version_list_url=https://prod.download.cli.kiro.dev/stable/latest/manifest.json,version_json_path=.version]" kiro-cli +fi diff --git a/test/shell.d/default-agent-test.sh b/test/shell.d/default-agent-test.sh index a1451ef5d54..0b929003e21 100644 --- a/test/shell.d/default-agent-test.sh +++ b/test/shell.d/default-agent-test.sh @@ -136,6 +136,7 @@ qwen_package="npm:@qwen-code/qwen-code" dsh_package="npm:@deepseek-ai/dsh@0.2.0-rc.2" mcode_package="npm:@minimax-ai/code[allow_builds=better-sqlite3]" goose_package="aqua:aaif-goose/goose" +kiro_cli_package="http:kiro-cli[platforms_linux_x64_url=https://prod.download.cli.kiro.dev/stable/{{version}}/kirocli-x86_64-linux.tar.gz,platforms_linux_arm64_url=https://prod.download.cli.kiro.dev/stable/{{version}}/kirocli-aarch64-linux.tar.gz,version_list_url=https://prod.download.cli.kiro.dev/stable/latest/manifest.json,version_json_path=.version]" cursor_agent_package="cursor-agent" muse_package="http:muse[url=https://api.meta.ai/muse-launcher.sh,bin=muse,version_list_url=https://api.meta.ai/muse-code/channels/muse-stable,version_json_path=.version]" @@ -161,6 +162,7 @@ assert_lazy_stub "$qwen_package" qwen assert_lazy_stub "$dsh_package" dsh assert_lazy_stub "$mcode_package" mcode assert_lazy_stub "$goose_package" goose +assert_lazy_stub "$kiro_cli_package" kiro-cli assert_lazy_stub "$cursor_agent_package" cursor-agent assert_lazy_stub "$muse_package" muse pass "custom agent lazy stubs preserve their mise packages" @@ -178,6 +180,7 @@ grep -Fx "$qwen_package qwen" "$stub_log" >/dev/null || fail "user setup creates grep -Fx "$dsh_package dsh" "$stub_log" >/dev/null || fail "user setup creates the DeepSeek Harness lazy stub" grep -Fx "$mcode_package mcode" "$stub_log" >/dev/null || fail "user setup creates the MiniMax Code lazy stub" grep -Fx "$goose_package goose" "$stub_log" >/dev/null || fail "user setup creates the Goose lazy stub" +grep -Fx "$kiro_cli_package kiro-cli" "$stub_log" >/dev/null || fail "user setup creates the Kiro CLI lazy stub" pass "user setup creates the custom agent lazy stubs" : >"$stub_log" @@ -189,6 +192,7 @@ grep -Fx "$qwen_package qwen" "$stub_log" >/dev/null || fail "user setup recreat grep -Fx "$dsh_package dsh" "$stub_log" >/dev/null || fail "user setup recreates the DeepSeek Harness lazy stub" grep -Fx "$mcode_package mcode" "$stub_log" >/dev/null || fail "user setup recreates the MiniMax Code lazy stub" grep -Fx "$goose_package goose" "$stub_log" >/dev/null || fail "user setup recreates the Goose lazy stub" +grep -Fx "$kiro_cli_package kiro-cli" "$stub_log" >/dev/null || fail "user setup recreates the Kiro CLI lazy stub" pass "user setup recreates wrappers even when the commands are present" : >"$stub_log" @@ -225,6 +229,10 @@ grep -Fx "$mcode_package mcode" "$stub_log" >/dev/null || fail "MiniMax Code mig OMARCHY_TEST_MISSING_COMMAND=goose source "$ROOT/migrations/1792033725.sh" >/dev/null grep -Fx "$goose_package goose" "$stub_log" >/dev/null || fail "Goose migration creates its lazy stub" +: >"$stub_log" +OMARCHY_TEST_MISSING_COMMAND=kiro-cli source "$ROOT/migrations/1792033795.sh" >/dev/null +grep -Fx "$kiro_cli_package kiro-cli" "$stub_log" >/dev/null || fail "Kiro CLI migration creates its lazy stub" + : >"$stub_log" source "$ROOT/migrations/1785617047.sh" >/dev/null grep -Fx "$omp_package omp" "$stub_log" >/dev/null || fail "Oh My Pi migration creates a working lazy stub" @@ -332,6 +340,7 @@ OMARCHY_TEST_MISSING_COMMAND=qwen source "$ROOT/migrations/1792032538.sh" >/dev/ OMARCHY_TEST_MISSING_COMMAND=dsh source "$ROOT/migrations/1792033356.sh" >/dev/null OMARCHY_TEST_MISSING_COMMAND=mcode source "$ROOT/migrations/1792033557.sh" >/dev/null OMARCHY_TEST_MISSING_COMMAND=goose source "$ROOT/migrations/1792033725.sh" >/dev/null +OMARCHY_TEST_MISSING_COMMAND=kiro-cli source "$ROOT/migrations/1792033795.sh" >/dev/null [[ ! -s $stub_log ]] || fail "agent migrations respect the preinstall opt-out" [[ ! -e $test_home/.local/bin/omp ]] || fail "agent migration removes the obsolete Oh My Pi wrapper after opt-out" @@ -469,9 +478,10 @@ pass "agent migrations install working wrappers without overriding the preinstal "$ROOT/bin/omarchy-mise-install" "$dsh_package" dsh "$ROOT/bin/omarchy-mise-install" "$mcode_package" mcode "$ROOT/bin/omarchy-mise-install" "$goose_package" goose +"$ROOT/bin/omarchy-mise-install" "$kiro_cli_package" kiro-cli touch "$test_home/.local/bin/agy" "$test_home/.local/bin/ori" omarchy-remove-preinstalls >/dev/null -for command in agy omp ori grok crush cursor-agent muse kimi qwen dsh mcode goose; do +for command in agy omp ori grok crush cursor-agent muse kimi qwen dsh mcode goose kiro-cli; do [[ ! -e $test_home/.local/bin/$command ]] || fail "Remove Preinstalls deletes the $command lazy stub" done pass "Remove Preinstalls deletes every optional agent lazy stub" @@ -563,6 +573,8 @@ declare -A expected_agents=( [minimax-code]="mcode" [minimax]="mcode" [goose]="goose" + [kiro-cli]="kiro-cli" + [kiro]="kiro-cli" [cursor]="cursor-agent" [cursor-agent]="cursor-agent" [muse]="muse" @@ -586,6 +598,7 @@ declare -A expected_packages=( [dsh]="$dsh_package" [mcode]="$mcode_package" [goose]="$goose_package" + [kiro-cli]="$kiro_cli_package" [cursor-agent]="$cursor_agent_package" [muse]="$muse_package" ) @@ -866,6 +879,7 @@ omarchy-agent-prompt '-n --help; echo literal' assert_launched dsh "keeps an option-like prompt literal" env DSH_PERMISSION_MODE=danger-full-access dsh headless -- -- '-n --help; echo literal' assert_launch mcode mcode -- "Review this project" assert_launch goose env GOOSE_MODE=auto goose run --interactive "--text=Review this project" +assert_launch kiro-cli kiro-cli chat --trust-all-tools -- "Review this project" pass "agent launcher adapts initial prompts for every supported agent" literal_muse_prompt=$'--disable-sandbox !Crash {$(touch must-not-run)}\ntrailing\\ ' @@ -899,6 +913,7 @@ assert_bypass qwen qwen --yolo assert_bypass dsh env DSH_PERMISSION_MODE=danger-full-access dsh web assert_bypass mcode mcode assert_bypass goose env GOOSE_MODE=auto goose +assert_bypass kiro-cli kiro-cli chat --trust-all-tools pass "agent launcher skips permission prompts for every supported agent" printf '%s\n' "opencode" >"$agent_file" diff --git a/test/shell.d/kiro-cli-agent-test.sh b/test/shell.d/kiro-cli-agent-test.sh new file mode 100644 index 00000000000..5547fb9ad62 --- /dev/null +++ b/test/shell.d/kiro-cli-agent-test.sh @@ -0,0 +1,85 @@ +#!/bin/bash + +set -euo pipefail +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" + +test_tmp=$(mktemp -d) +trap 'rm -rf "$test_tmp"' EXIT +mock_bin="$test_tmp/bin" +export HOME="$test_tmp/home with spaces" +export PATH="$mock_bin:$HOME/.local/bin:$ROOT/bin:/usr/local/bin:/usr/bin" +export OMARCHY_PATH="$ROOT" +export OMARCHY_TEST_MISE_LOG="$test_tmp/mise-log" +mkdir -p "$mock_bin" "$HOME/.local/bin" +kiro_cli="$HOME/.local/bin/kiro-cli" +package="http:kiro-cli[platforms_linux_x64_url=https://prod.download.cli.kiro.dev/stable/{{version}}/kirocli-x86_64-linux.tar.gz,platforms_linux_arm64_url=https://prod.download.cli.kiro.dev/stable/{{version}}/kirocli-aarch64-linux.tar.gz,version_list_url=https://prod.download.cli.kiro.dev/stable/latest/manifest.json,version_json_path=.version]" + +cat >"$mock_bin/mise" <<'SH' +#!/bin/bash +printf '%s\0' "$@" >>"$OMARCHY_TEST_MISE_LOG" +SH +cat >"$mock_bin/noop" <<'SH' +#!/bin/bash +exit 0 +SH +chmod +x "$mock_bin/mise" "$mock_bin/noop" +for command in gum hyprctl omarchy-webapp-remove-all omarchy-tui-remove-all omarchy-pkg-drop; do + ln -s noop "$mock_bin/$command" +done + +apply_kiro_cli() { + case $1 in + migration) source "$ROOT/migrations/1792033795.sh" >/dev/null ;; + removal) omarchy-remove-preinstalls >/dev/null ;; + esac +} + +apply_kiro_cli migration +[[ -x $kiro_cli && ! -L $kiro_cli ]] || fail "Kiro CLI migration creates a missing wrapper" +[[ ! -s $OMARCHY_TEST_MISE_LOG ]] || fail "Kiro CLI migration installs the package before first launch" +cp "$kiro_cli" "$test_tmp/wrapper" +apply_kiro_cli migration +cmp -s "$kiro_cli" "$test_tmp/wrapper" || fail "Kiro CLI migration is idempotent" +"$kiro_cli" --prompt 'Review this project' +mapfile -d '' -t args <"$OMARCHY_TEST_MISE_LOG" +[[ ${args[*]} == "use -g --quiet $package x $package -- kiro-cli --prompt Review this project" ]] || + fail "Kiro CLI migration wrapper installs the official package and forwards arguments" +pass "Kiro CLI migration creates a lazy, repeatable wrapper" + +rm -f "$kiro_cli" +mkdir -p "$HOME/.local/state/omarchy" +touch "$HOME/.local/state/omarchy/preinstalls-removed" +apply_kiro_cli migration +[[ ! -e $kiro_cli ]] || fail "Kiro CLI migration respects removed preinstalls with a missing command" +rm "$HOME/.local/state/omarchy/preinstalls-removed" +pass "Kiro CLI migration respects the preinstall opt-out" + +for operation in migration removal; do + for kind in script binary symlink; do + rm -f "$kiro_cli" + case $kind in + script) printf '#!/bin/bash\necho user-managed kiro-cli\n' >"$kiro_cli" ;; + binary) cp /usr/bin/true "$kiro_cli" ;; + symlink) ln -s /usr/bin/true "$kiro_cli" ;; + esac + [[ $kind == "symlink" ]] || chmod +x "$kiro_cli" + rm -f "$test_tmp/saved-$kind" + cp -P "$kiro_cli" "$test_tmp/saved-$kind" + for attempt in 1 2; do + apply_kiro_cli "$operation" + cmp -s "$kiro_cli" "$test_tmp/saved-$kind" || fail "Kiro CLI $operation preserves a user $kind" + if [[ $kind == "symlink" ]]; then + [[ -L $kiro_cli && $(readlink "$kiro_cli") == "/usr/bin/true" ]] || + fail "Kiro CLI $operation preserves the user's symlink" + else + [[ ! -L $kiro_cli ]] || fail "Kiro CLI $operation keeps the user's $kind as a regular file" + fi + done + pass "Kiro CLI $operation preserves a user $kind on repeated runs" + done +done + +omarchy-mise-install "$package" kiro-cli +apply_kiro_cli removal +[[ ! -e $kiro_cli ]] || fail "Remove Preinstalls removes Omarchy's Kiro CLI wrapper" +pass "Remove Preinstalls removes Omarchy's Kiro CLI wrapper" diff --git a/test/shell.d/menu-test.sh b/test/shell.d/menu-test.sh index 85f6e5c1f1e..b5192de6ce9 100644 --- a/test/shell.d/menu-test.sh +++ b/test/shell.d/menu-test.sh @@ -257,6 +257,7 @@ const expectedAgents = { dsh: { icon: '\ue910', iconFont: 'omarchy', label: 'DeepSeek Harness' }, mcode: { icon: '\ue911', iconFont: 'omarchy', label: 'MiniMax Code' }, goose: { icon: '\ue912', iconFont: 'omarchy', label: 'Goose' }, + 'kiro-cli': { icon: '\ue913', iconFont: 'omarchy', label: 'Kiro CLI' }, hermes: { icon: '\ue90a', iconFont: 'omarchy', label: 'Hermes' }, openclaw: { icon: '\ue90c', iconFont: 'omarchy', label: 'OpenClaw' }, copilot: { icon: '', label: 'Copilot' }, @@ -282,7 +283,7 @@ assertDeepEqual( defaultItems .filter(item => item.parent === 'setup.default.agent') .map(item => item.label), - ['Antigravity', 'Claude', 'Codex', 'Copilot', 'Crush', 'Cursor CLI', 'DeepSeek Harness', 'Goose', 'Grok', 'Hermes', 'Kimi', 'MiniMax Code', 'Muse Code', 'omp', 'OpenClaw', 'OpenCode', 'Ori', 'Pi', 'Qwen Code'], + ['Antigravity', 'Claude', 'Codex', 'Copilot', 'Crush', 'Cursor CLI', 'DeepSeek Harness', 'Goose', 'Grok', 'Hermes', 'Kimi', 'Kiro CLI', 'MiniMax Code', 'Muse Code', 'omp', 'OpenClaw', 'OpenCode', 'Ori', 'Pi', 'Qwen Code'], 'menu sorts coding agents alphabetically' ) const expectedDefaults = { @@ -879,5 +880,5 @@ assert( JS font_charset=$(fc-query --format='%{charset}' "$ROOT/default/fonts/omarchy/omarchy.ttf") -[[ $font_charset == *"e900-e912"* ]] || fail "Omarchy icon font includes every custom menu glyph" +[[ $font_charset == *"e900-e913"* ]] || fail "Omarchy icon font includes every custom menu glyph" pass "Omarchy icon font includes the official agent marks" diff --git a/test/shell.d/mise-provision-test.sh b/test/shell.d/mise-provision-test.sh index 7fd2edff58d..270bc06be80 100644 --- a/test/shell.d/mise-provision-test.sh +++ b/test/shell.d/mise-provision-test.sh @@ -17,7 +17,7 @@ SH chmod +x "$mock_bin/mise" source "$ROOT/install/user/mise.sh" -for agent in kimi cursor-agent muse qwen dsh mcode goose; do +for agent in kimi cursor-agent muse qwen dsh mcode goose kiro-cli; do [[ -x $HOME/.local/bin/$agent && ! -L $HOME/.local/bin/$agent ]] || fail "provisioning creates the missing $agent wrapper" cp "$HOME/.local/bin/$agent" "$test_tmp/$agent-wrapper" @@ -27,7 +27,7 @@ mapfile -d '' -t args <"$OMARCHY_TEST_MISE_LOG" fail "provisioning creates lazy wrappers without installing packages" pass "provisioning creates lazy Kimi, Cursor, and Muse wrappers" -for agent in kimi cursor-agent muse qwen dsh mcode goose; do +for agent in kimi cursor-agent muse qwen dsh mcode goose kiro-cli; do launcher="$HOME/.local/bin/$agent" for kind in script binary symlink; do rm -f "$launcher" From 3fa00968e12351a686b9ab9cecb5645298cfdc92 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:11:47 -0400 Subject: [PATCH 101/546] Add Mistral Vibe as a selectable default coding agent Mistral Vibe (vibe) joins the default agent picker. mise installs the standalone vibe-linux release archive for x86_64 or aarch64 from GitHub, since the pipx backend needs pipx or uv, which Omarchy does not ship. The launcher runs it with --auto-approve and seeds the interactive session with a positional prompt after --. Combines #14847 and #10838. Co-authored-by: Hans Co-authored-by: Josh Christopher Algdal Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent | 5 ++ bin/omarchy-default-agent | 5 +- bin/omarchy-remove-preinstalls | 6 ++ default/fonts/omarchy/README.md | 1 + default/fonts/omarchy/omarchy.ttf | Bin 12268 -> 12364 bytes default/i18n/en/menu.json | 3 + default/omarchy/omarchy-menu.jsonc | 1 + install/user/mise.sh | 1 + manual/17-ai.md | 1 + migrations/1792033880.sh | 5 ++ test/shell.d/default-agent-test.sh | 18 +++++- test/shell.d/menu-test.sh | 5 +- test/shell.d/mise-provision-test.sh | 4 +- test/shell.d/vibe-agent-test.sh | 85 ++++++++++++++++++++++++++++ 14 files changed, 133 insertions(+), 7 deletions(-) create mode 100644 migrations/1792033880.sh create mode 100644 test/shell.d/vibe-agent-test.sh diff --git a/bin/omarchy-agent b/bin/omarchy-agent index 67fd3ffc62a..2d74fd50346 100755 --- a/bin/omarchy-agent +++ b/bin/omarchy-agent @@ -174,6 +174,11 @@ kiro-cli) command=(kiro-cli chat --trust-all-tools) [[ -n ${prompt:-} ]] && command+=(-- "$prompt") ;; +vibe) + # A positional prompt seeds the interactive session; -p would run it headless. + command=(vibe --auto-approve) + [[ -n ${prompt:-} ]] && command+=(-- "$prompt") + ;; *) echo "Unsupported default agent: $agent" >&2 exit 1 diff --git a/bin/omarchy-default-agent b/bin/omarchy-default-agent index 1c62689c849..8aaa6186b09 100755 --- a/bin/omarchy-default-agent +++ b/bin/omarchy-default-agent @@ -1,7 +1,7 @@ #!/bin/bash # omarchy:summary=Set and launch the default coding agent -# omarchy:args=[pi|omp|opencode|ori|claude|codex|grok|openclaw|agy|hermes|copilot|crush|cursor-agent|muse|kimi|qwen|dsh|mcode|goose|kiro-cli] +# omarchy:args=[pi|omp|opencode|ori|claude|codex|grok|openclaw|agy|hermes|copilot|crush|cursor-agent|muse|kimi|qwen|dsh|mcode|goose|kiro-cli|vibe] # omarchy:examples=omarchy default agent | omarchy default agent codex | omarchy default agent claude installing=false @@ -58,8 +58,9 @@ kiro-cli | kiro) # Kiro publishes per-platform archives behind a manifest of the stable release. agent_package="http:kiro-cli[platforms_linux_x64_url=https://prod.download.cli.kiro.dev/stable/{{version}}/kirocli-x86_64-linux.tar.gz,platforms_linux_arm64_url=https://prod.download.cli.kiro.dev/stable/{{version}}/kirocli-aarch64-linux.tar.gz,version_list_url=https://prod.download.cli.kiro.dev/stable/latest/manifest.json,version_json_path=.version]" ;; +vibe | mistral | mistral-vibe) agent="vibe"; name="Mistral Vibe"; agent_package="github:mistralai/mistral-vibe[matching=vibe-linux]" ;; *) - echo "Usage: omarchy-default-agent " + echo "Usage: omarchy-default-agent " exit 1 ;; esac diff --git a/bin/omarchy-remove-preinstalls b/bin/omarchy-remove-preinstalls index c31fd37db3e..c50b149072c 100755 --- a/bin/omarchy-remove-preinstalls +++ b/bin/omarchy-remove-preinstalls @@ -66,6 +66,12 @@ if gum confirm --affirmative "$(omarchy-i18n common.yes)" --negative "$(omarchy- rm -f ~/.local/bin/kiro-cli fi + # Preserve a user-managed Mistral Vibe launcher at the same path. + if [[ -f ~/.local/bin/vibe && ! -L ~/.local/bin/vibe ]] && + grep -Eq '^mise use -g .*"github:mistralai/mistral-vibe\[matching=vibe-linux\]"' ~/.local/bin/vibe; then + rm -f ~/.local/bin/vibe + fi + omarchy-pkg-drop \ aether \ cliamp \ diff --git a/default/fonts/omarchy/README.md b/default/fonts/omarchy/README.md index 469011a7bed..c67eacf5e42 100644 --- a/default/fonts/omarchy/README.md +++ b/default/fonts/omarchy/README.md @@ -22,6 +22,7 @@ The private-use glyphs in `omarchy.ttf` are: - `U+E911` — MiniMax, from - `U+E912` — Goose, from - `U+E913` — Kiro, from the ghost in , with its tile dropped +- `U+E914` — Mistral AI, from The agent marks are monochrome so the menu can render them using the active theme's foreground and selection colors. diff --git a/default/fonts/omarchy/omarchy.ttf b/default/fonts/omarchy/omarchy.ttf index 3f6a488be697bb7111405e8ce3578482688c85ea..0cca4ccbee2f6cd0119a6ecfc9dc58764352894a 100644 GIT binary patch delta 324 zcmaD8eS9(rmTAKM|R|W>=D+~;BYcf(3Q}otZ^)WCgH83zR zm}LM3I2c*vfczamzDh=JNks?)yC#r-2gonU$xlx7)DPlhU{DSK@@;YxD+(AbnGAsZ z9w1*KFEKaOxcg2mkiUk3fqP>?esKv;qHI$t1A~47P{FH>dt@0IMJC^3e9I{UJFh~Q%1&u`&O^r>JMU{aV%m#`8AyACjq2XTx7(gjz zhkpzZfrbVqFbiln10x6<8@Je{WlXkF@@JHt+@n;$!otAdF!_&?6$gJ9ScU9nJLUJ> bKuJA)3noUf$!Yp+TuO delta 225 zcmX?;@FsqOVSS;XHv2Ll6>S9(rm+DZ3}Zww5~R~Q)NVlq+_Q}otZ^)WCg889#~ zm}LM3I2c)Ef&2_0UnL{Aq#}TU?In=k1LPOvS3>N{-)^8Wz&R=J541&kI8D%ElQ7V`mplrnfRLTg%GMgtUzvl*uH|bkU I-lN|J0MBzd^Z)<= diff --git a/default/i18n/en/menu.json b/default/i18n/en/menu.json index 7eefcceb166..34ebcf1578a 100644 --- a/default/i18n/en/menu.json +++ b/default/i18n/en/menu.json @@ -542,6 +542,9 @@ "setup.default.agent.qwen": { "label": "Qwen Code" }, + "setup.default.agent.vibe": { + "label": "Mistral Vibe" + }, "setup.default.browser": { "label": "Browser", "title": "Default Browser" diff --git a/default/omarchy/omarchy-menu.jsonc b/default/omarchy/omarchy-menu.jsonc index e773403e808..13e3efb32a3 100644 --- a/default/omarchy/omarchy-menu.jsonc +++ b/default/omarchy/omarchy-menu.jsonc @@ -203,6 +203,7 @@ "setup.default.agent.kimi": {"icon":"󰚩","label":"Kimi","checked":"[[ \"$(omarchy-default-agent)\" == \"kimi\" ]]","action":"omarchy-default-agent kimi"}, "setup.default.agent.kiro-cli": {"icon":"","iconFont":"omarchy","label":"Kiro CLI","checked":"[[ \"$(omarchy-default-agent)\" == \"kiro-cli\" ]]","action":"omarchy-default-agent kiro-cli"}, "setup.default.agent.mcode": {"icon":"","iconFont":"omarchy","label":"MiniMax Code","checked":"[[ \"$(omarchy-default-agent)\" == \"mcode\" ]]","action":"omarchy-default-agent mcode"}, + "setup.default.agent.vibe": {"icon":"","iconFont":"omarchy","label":"Mistral Vibe","checked":"[[ \"$(omarchy-default-agent)\" == \"vibe\" ]]","action":"omarchy-default-agent vibe"}, "setup.default.agent.muse": {"icon":"󰛤","label":"Muse Code","checked":"[[ \"$(omarchy-default-agent)\" == \"muse\" ]]","action":"omarchy-default-agent muse"}, "setup.default.agent.omp": {"icon":"","iconFont":"omarchy","label":"omp","checked":"[[ \"$(omarchy-default-agent)\" == \"omp\" ]]","action":"omarchy-default-agent omp"}, "setup.default.agent.openclaw": {"icon":"","iconFont":"omarchy","label":"OpenClaw","checked":"[[ \"$(omarchy-default-agent)\" == \"openclaw\" ]]","action":"omarchy-default-agent openclaw"}, diff --git a/install/user/mise.sh b/install/user/mise.sh index c4152688761..fbc3ec188ec 100644 --- a/install/user/mise.sh +++ b/install/user/mise.sh @@ -20,6 +20,7 @@ omarchy-mise-install npm:@deepseek-ai/dsh@0.2.0-rc.2 dsh omarchy-mise-install "npm:@minimax-ai/code[allow_builds=better-sqlite3]" mcode omarchy-mise-install aqua:aaif-goose/goose goose omarchy-mise-install "http:kiro-cli[platforms_linux_x64_url=https://prod.download.cli.kiro.dev/stable/{{version}}/kirocli-x86_64-linux.tar.gz,platforms_linux_arm64_url=https://prod.download.cli.kiro.dev/stable/{{version}}/kirocli-aarch64-linux.tar.gz,version_list_url=https://prod.download.cli.kiro.dev/stable/latest/manifest.json,version_json_path=.version]" kiro-cli +omarchy-mise-install "github:mistralai/mistral-vibe[matching=vibe-linux]" vibe omarchy-mise-install cursor-agent omarchy-mise-install npm:@kitlangton/ghui ghui omarchy-mise-install aqua:modem-dev/hunk hunk diff --git a/manual/17-ai.md b/manual/17-ai.md index 9ce762ce2e1..f54972af7e4 100644 --- a/manual/17-ai.md +++ b/manual/17-ai.md @@ -17,6 +17,7 @@ Omarchy treats AI coding agents as first-class citizens, but it doesn't pick a f | `kiro-cli` | [Kiro CLI](https://kiro.dev/docs/cli/) | | `mcode` | [MiniMax Code](https://github.com/MiniMax-AI/minimax-code), MiniMax's coding agent | | `qwen` | [Qwen Code](https://github.com/QwenLM/qwen-code), Alibaba's coding agent | +| `vibe` | [Mistral Vibe](https://docs.mistral.ai/vibe/code/overview), Mistral AI's coding agent | | `pi` | [Mario Zechner's Pi](https://github.com/badlogic/pi-mono) | | `omp` | [Oh My Pi](https://github.com/can1357/oh-my-pi) | | `ori` | [Ori](https://openrouter.ai/docs/guides/ori/harness), OpenRouter's harness | diff --git a/migrations/1792033880.sh b/migrations/1792033880.sh new file mode 100644 index 00000000000..bf1aec629c0 --- /dev/null +++ b/migrations/1792033880.sh @@ -0,0 +1,5 @@ +echo "Install Mistral Vibe coding agent mise wrapper" + +if [[ ! -f $HOME/.local/state/omarchy/preinstalls-removed ]] && omarchy-cmd-missing vibe; then + omarchy-mise-install "github:mistralai/mistral-vibe[matching=vibe-linux]" vibe +fi diff --git a/test/shell.d/default-agent-test.sh b/test/shell.d/default-agent-test.sh index 0b929003e21..f25f4351c21 100644 --- a/test/shell.d/default-agent-test.sh +++ b/test/shell.d/default-agent-test.sh @@ -137,6 +137,7 @@ dsh_package="npm:@deepseek-ai/dsh@0.2.0-rc.2" mcode_package="npm:@minimax-ai/code[allow_builds=better-sqlite3]" goose_package="aqua:aaif-goose/goose" kiro_cli_package="http:kiro-cli[platforms_linux_x64_url=https://prod.download.cli.kiro.dev/stable/{{version}}/kirocli-x86_64-linux.tar.gz,platforms_linux_arm64_url=https://prod.download.cli.kiro.dev/stable/{{version}}/kirocli-aarch64-linux.tar.gz,version_list_url=https://prod.download.cli.kiro.dev/stable/latest/manifest.json,version_json_path=.version]" +vibe_package="github:mistralai/mistral-vibe[matching=vibe-linux]" cursor_agent_package="cursor-agent" muse_package="http:muse[url=https://api.meta.ai/muse-launcher.sh,bin=muse,version_list_url=https://api.meta.ai/muse-code/channels/muse-stable,version_json_path=.version]" @@ -163,6 +164,7 @@ assert_lazy_stub "$dsh_package" dsh assert_lazy_stub "$mcode_package" mcode assert_lazy_stub "$goose_package" goose assert_lazy_stub "$kiro_cli_package" kiro-cli +assert_lazy_stub "$vibe_package" vibe assert_lazy_stub "$cursor_agent_package" cursor-agent assert_lazy_stub "$muse_package" muse pass "custom agent lazy stubs preserve their mise packages" @@ -181,6 +183,7 @@ grep -Fx "$dsh_package dsh" "$stub_log" >/dev/null || fail "user setup creates t grep -Fx "$mcode_package mcode" "$stub_log" >/dev/null || fail "user setup creates the MiniMax Code lazy stub" grep -Fx "$goose_package goose" "$stub_log" >/dev/null || fail "user setup creates the Goose lazy stub" grep -Fx "$kiro_cli_package kiro-cli" "$stub_log" >/dev/null || fail "user setup creates the Kiro CLI lazy stub" +grep -Fx "$vibe_package vibe" "$stub_log" >/dev/null || fail "user setup creates the Mistral Vibe lazy stub" pass "user setup creates the custom agent lazy stubs" : >"$stub_log" @@ -193,6 +196,7 @@ grep -Fx "$dsh_package dsh" "$stub_log" >/dev/null || fail "user setup recreates grep -Fx "$mcode_package mcode" "$stub_log" >/dev/null || fail "user setup recreates the MiniMax Code lazy stub" grep -Fx "$goose_package goose" "$stub_log" >/dev/null || fail "user setup recreates the Goose lazy stub" grep -Fx "$kiro_cli_package kiro-cli" "$stub_log" >/dev/null || fail "user setup recreates the Kiro CLI lazy stub" +grep -Fx "$vibe_package vibe" "$stub_log" >/dev/null || fail "user setup recreates the Mistral Vibe lazy stub" pass "user setup recreates wrappers even when the commands are present" : >"$stub_log" @@ -233,6 +237,10 @@ grep -Fx "$goose_package goose" "$stub_log" >/dev/null || fail "Goose migration OMARCHY_TEST_MISSING_COMMAND=kiro-cli source "$ROOT/migrations/1792033795.sh" >/dev/null grep -Fx "$kiro_cli_package kiro-cli" "$stub_log" >/dev/null || fail "Kiro CLI migration creates its lazy stub" +: >"$stub_log" +OMARCHY_TEST_MISSING_COMMAND=vibe source "$ROOT/migrations/1792033880.sh" >/dev/null +grep -Fx "$vibe_package vibe" "$stub_log" >/dev/null || fail "Mistral Vibe migration creates its lazy stub" + : >"$stub_log" source "$ROOT/migrations/1785617047.sh" >/dev/null grep -Fx "$omp_package omp" "$stub_log" >/dev/null || fail "Oh My Pi migration creates a working lazy stub" @@ -341,6 +349,7 @@ OMARCHY_TEST_MISSING_COMMAND=dsh source "$ROOT/migrations/1792033356.sh" >/dev/n OMARCHY_TEST_MISSING_COMMAND=mcode source "$ROOT/migrations/1792033557.sh" >/dev/null OMARCHY_TEST_MISSING_COMMAND=goose source "$ROOT/migrations/1792033725.sh" >/dev/null OMARCHY_TEST_MISSING_COMMAND=kiro-cli source "$ROOT/migrations/1792033795.sh" >/dev/null +OMARCHY_TEST_MISSING_COMMAND=vibe source "$ROOT/migrations/1792033880.sh" >/dev/null [[ ! -s $stub_log ]] || fail "agent migrations respect the preinstall opt-out" [[ ! -e $test_home/.local/bin/omp ]] || fail "agent migration removes the obsolete Oh My Pi wrapper after opt-out" @@ -479,9 +488,10 @@ pass "agent migrations install working wrappers without overriding the preinstal "$ROOT/bin/omarchy-mise-install" "$mcode_package" mcode "$ROOT/bin/omarchy-mise-install" "$goose_package" goose "$ROOT/bin/omarchy-mise-install" "$kiro_cli_package" kiro-cli +"$ROOT/bin/omarchy-mise-install" "$vibe_package" vibe touch "$test_home/.local/bin/agy" "$test_home/.local/bin/ori" omarchy-remove-preinstalls >/dev/null -for command in agy omp ori grok crush cursor-agent muse kimi qwen dsh mcode goose kiro-cli; do +for command in agy omp ori grok crush cursor-agent muse kimi qwen dsh mcode goose kiro-cli vibe; do [[ ! -e $test_home/.local/bin/$command ]] || fail "Remove Preinstalls deletes the $command lazy stub" done pass "Remove Preinstalls deletes every optional agent lazy stub" @@ -575,6 +585,9 @@ declare -A expected_agents=( [goose]="goose" [kiro-cli]="kiro-cli" [kiro]="kiro-cli" + [vibe]="vibe" + [mistral]="vibe" + [mistral-vibe]="vibe" [cursor]="cursor-agent" [cursor-agent]="cursor-agent" [muse]="muse" @@ -599,6 +612,7 @@ declare -A expected_packages=( [mcode]="$mcode_package" [goose]="$goose_package" [kiro-cli]="$kiro_cli_package" + [vibe]="$vibe_package" [cursor-agent]="$cursor_agent_package" [muse]="$muse_package" ) @@ -880,6 +894,7 @@ assert_launched dsh "keeps an option-like prompt literal" env DSH_PERMISSION_MOD assert_launch mcode mcode -- "Review this project" assert_launch goose env GOOSE_MODE=auto goose run --interactive "--text=Review this project" assert_launch kiro-cli kiro-cli chat --trust-all-tools -- "Review this project" +assert_launch vibe vibe --auto-approve -- "Review this project" pass "agent launcher adapts initial prompts for every supported agent" literal_muse_prompt=$'--disable-sandbox !Crash {$(touch must-not-run)}\ntrailing\\ ' @@ -914,6 +929,7 @@ assert_bypass dsh env DSH_PERMISSION_MODE=danger-full-access dsh web assert_bypass mcode mcode assert_bypass goose env GOOSE_MODE=auto goose assert_bypass kiro-cli kiro-cli chat --trust-all-tools +assert_bypass vibe vibe --auto-approve pass "agent launcher skips permission prompts for every supported agent" printf '%s\n' "opencode" >"$agent_file" diff --git a/test/shell.d/menu-test.sh b/test/shell.d/menu-test.sh index b5192de6ce9..49d3bf0f2ca 100644 --- a/test/shell.d/menu-test.sh +++ b/test/shell.d/menu-test.sh @@ -258,6 +258,7 @@ const expectedAgents = { mcode: { icon: '\ue911', iconFont: 'omarchy', label: 'MiniMax Code' }, goose: { icon: '\ue912', iconFont: 'omarchy', label: 'Goose' }, 'kiro-cli': { icon: '\ue913', iconFont: 'omarchy', label: 'Kiro CLI' }, + vibe: { icon: '\ue914', iconFont: 'omarchy', label: 'Mistral Vibe' }, hermes: { icon: '\ue90a', iconFont: 'omarchy', label: 'Hermes' }, openclaw: { icon: '\ue90c', iconFont: 'omarchy', label: 'OpenClaw' }, copilot: { icon: '', label: 'Copilot' }, @@ -283,7 +284,7 @@ assertDeepEqual( defaultItems .filter(item => item.parent === 'setup.default.agent') .map(item => item.label), - ['Antigravity', 'Claude', 'Codex', 'Copilot', 'Crush', 'Cursor CLI', 'DeepSeek Harness', 'Goose', 'Grok', 'Hermes', 'Kimi', 'Kiro CLI', 'MiniMax Code', 'Muse Code', 'omp', 'OpenClaw', 'OpenCode', 'Ori', 'Pi', 'Qwen Code'], + ['Antigravity', 'Claude', 'Codex', 'Copilot', 'Crush', 'Cursor CLI', 'DeepSeek Harness', 'Goose', 'Grok', 'Hermes', 'Kimi', 'Kiro CLI', 'MiniMax Code', 'Mistral Vibe', 'Muse Code', 'omp', 'OpenClaw', 'OpenCode', 'Ori', 'Pi', 'Qwen Code'], 'menu sorts coding agents alphabetically' ) const expectedDefaults = { @@ -880,5 +881,5 @@ assert( JS font_charset=$(fc-query --format='%{charset}' "$ROOT/default/fonts/omarchy/omarchy.ttf") -[[ $font_charset == *"e900-e913"* ]] || fail "Omarchy icon font includes every custom menu glyph" +[[ $font_charset == *"e900-e914"* ]] || fail "Omarchy icon font includes every custom menu glyph" pass "Omarchy icon font includes the official agent marks" diff --git a/test/shell.d/mise-provision-test.sh b/test/shell.d/mise-provision-test.sh index 270bc06be80..faa175d501f 100644 --- a/test/shell.d/mise-provision-test.sh +++ b/test/shell.d/mise-provision-test.sh @@ -17,7 +17,7 @@ SH chmod +x "$mock_bin/mise" source "$ROOT/install/user/mise.sh" -for agent in kimi cursor-agent muse qwen dsh mcode goose kiro-cli; do +for agent in kimi cursor-agent muse qwen dsh mcode goose kiro-cli vibe; do [[ -x $HOME/.local/bin/$agent && ! -L $HOME/.local/bin/$agent ]] || fail "provisioning creates the missing $agent wrapper" cp "$HOME/.local/bin/$agent" "$test_tmp/$agent-wrapper" @@ -27,7 +27,7 @@ mapfile -d '' -t args <"$OMARCHY_TEST_MISE_LOG" fail "provisioning creates lazy wrappers without installing packages" pass "provisioning creates lazy Kimi, Cursor, and Muse wrappers" -for agent in kimi cursor-agent muse qwen dsh mcode goose kiro-cli; do +for agent in kimi cursor-agent muse qwen dsh mcode goose kiro-cli vibe; do launcher="$HOME/.local/bin/$agent" for kind in script binary symlink; do rm -f "$launcher" diff --git a/test/shell.d/vibe-agent-test.sh b/test/shell.d/vibe-agent-test.sh new file mode 100644 index 00000000000..2cd37f7878f --- /dev/null +++ b/test/shell.d/vibe-agent-test.sh @@ -0,0 +1,85 @@ +#!/bin/bash + +set -euo pipefail +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" + +test_tmp=$(mktemp -d) +trap 'rm -rf "$test_tmp"' EXIT +mock_bin="$test_tmp/bin" +export HOME="$test_tmp/home with spaces" +export PATH="$mock_bin:$HOME/.local/bin:$ROOT/bin:/usr/local/bin:/usr/bin" +export OMARCHY_PATH="$ROOT" +export OMARCHY_TEST_MISE_LOG="$test_tmp/mise-log" +mkdir -p "$mock_bin" "$HOME/.local/bin" +vibe="$HOME/.local/bin/vibe" +package="github:mistralai/mistral-vibe[matching=vibe-linux]" + +cat >"$mock_bin/mise" <<'SH' +#!/bin/bash +printf '%s\0' "$@" >>"$OMARCHY_TEST_MISE_LOG" +SH +cat >"$mock_bin/noop" <<'SH' +#!/bin/bash +exit 0 +SH +chmod +x "$mock_bin/mise" "$mock_bin/noop" +for command in gum hyprctl omarchy-webapp-remove-all omarchy-tui-remove-all omarchy-pkg-drop; do + ln -s noop "$mock_bin/$command" +done + +apply_vibe() { + case $1 in + migration) source "$ROOT/migrations/1792033880.sh" >/dev/null ;; + removal) omarchy-remove-preinstalls >/dev/null ;; + esac +} + +apply_vibe migration +[[ -x $vibe && ! -L $vibe ]] || fail "Mistral Vibe migration creates a missing wrapper" +[[ ! -s $OMARCHY_TEST_MISE_LOG ]] || fail "Mistral Vibe migration installs the package before first launch" +cp "$vibe" "$test_tmp/wrapper" +apply_vibe migration +cmp -s "$vibe" "$test_tmp/wrapper" || fail "Mistral Vibe migration is idempotent" +"$vibe" --prompt 'Review this project' +mapfile -d '' -t args <"$OMARCHY_TEST_MISE_LOG" +[[ ${args[*]} == "use -g --quiet $package x $package -- vibe --prompt Review this project" ]] || + fail "Mistral Vibe migration wrapper installs the official package and forwards arguments" +pass "Mistral Vibe migration creates a lazy, repeatable wrapper" + +rm -f "$vibe" +mkdir -p "$HOME/.local/state/omarchy" +touch "$HOME/.local/state/omarchy/preinstalls-removed" +apply_vibe migration +[[ ! -e $vibe ]] || fail "Mistral Vibe migration respects removed preinstalls with a missing command" +rm "$HOME/.local/state/omarchy/preinstalls-removed" +pass "Mistral Vibe migration respects the preinstall opt-out" + +for operation in migration removal; do + for kind in script binary symlink; do + rm -f "$vibe" + case $kind in + script) printf '#!/bin/bash\necho user-managed vibe\n' >"$vibe" ;; + binary) cp /usr/bin/true "$vibe" ;; + symlink) ln -s /usr/bin/true "$vibe" ;; + esac + [[ $kind == "symlink" ]] || chmod +x "$vibe" + rm -f "$test_tmp/saved-$kind" + cp -P "$vibe" "$test_tmp/saved-$kind" + for attempt in 1 2; do + apply_vibe "$operation" + cmp -s "$vibe" "$test_tmp/saved-$kind" || fail "Mistral Vibe $operation preserves a user $kind" + if [[ $kind == "symlink" ]]; then + [[ -L $vibe && $(readlink "$vibe") == "/usr/bin/true" ]] || + fail "Mistral Vibe $operation preserves the user's symlink" + else + [[ ! -L $vibe ]] || fail "Mistral Vibe $operation keeps the user's $kind as a regular file" + fi + done + pass "Mistral Vibe $operation preserves a user $kind on repeated runs" + done +done + +omarchy-mise-install "$package" vibe +apply_vibe removal +[[ ! -e $vibe ]] || fail "Remove Preinstalls removes Omarchy's Mistral Vibe wrapper" +pass "Remove Preinstalls removes Omarchy's Mistral Vibe wrapper" From 278488b529107bc6a5b6600b06b152d5ad766ab8 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:14:18 -0400 Subject: [PATCH 102/546] Add Freebuff as a selectable default coding agent Freebuff (freebuff) joins the default agent picker, installed through mise as npm:freebuff with a lazy stub and migration like the other agents. It asks no tool permissions, so there is no bypass flag to pass. Freebuff takes no starting prompt either. Instead of refusing a prompted launch, which from the agents panel or a crash report would fail without a word, omarchy-agent opens Freebuff with the prompt on the clipboard and a notification saying so. Any agent that sets prompt_via_clipboard in its case gets the same handoff, and the panel's starter prompts keep working without knowing which agent is the default. The launcher part of #14019. Co-authored-by: Paul Lamb Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent | 15 +++++ bin/omarchy-default-agent | 5 +- bin/omarchy-remove-preinstalls | 6 ++ default/fonts/omarchy/README.md | 1 + default/fonts/omarchy/omarchy.ttf | Bin 12364 -> 12496 bytes default/i18n/en/cli.json | 2 + default/i18n/en/menu.json | 3 + default/omarchy/omarchy-menu.jsonc | 1 + install/user/mise.sh | 1 + manual/17-ai.md | 3 +- migrations/1792033907.sh | 5 ++ test/shell.d/default-agent-test.sh | 30 +++++++++- test/shell.d/freebuff-agent-test.sh | 85 ++++++++++++++++++++++++++++ test/shell.d/menu-test.sh | 5 +- test/shell.d/mise-provision-test.sh | 4 +- 15 files changed, 158 insertions(+), 8 deletions(-) create mode 100644 migrations/1792033907.sh create mode 100644 test/shell.d/freebuff-agent-test.sh diff --git a/bin/omarchy-agent b/bin/omarchy-agent index 2d74fd50346..f694b700414 100755 --- a/bin/omarchy-agent +++ b/bin/omarchy-agent @@ -179,12 +179,27 @@ vibe) command=(vibe --auto-approve) [[ -n ${prompt:-} ]] && command+=(-- "$prompt") ;; +freebuff) + # Freebuff takes no starting prompt and asks no tool permissions, so a + # prompt goes to the clipboard below instead. + command=(freebuff) + prompt_via_clipboard="Freebuff" + ;; *) echo "Unsupported default agent: $agent" >&2 exit 1 ;; esac +# An agent with no way to take a starting prompt still opens, with the prompt +# on the clipboard to paste as its first message, so a prompted launch from the +# agents panel or a crash report is not lost. +if [[ -n ${prompt:-} && -n ${prompt_via_clipboard:-} ]]; then + printf '%s' "$prompt" | wl-copy + omarchy-notification-send -g 󰅇 "$(omarchy-i18n agents.promptOnClipboard "$prompt_via_clipboard")" \ + "$(omarchy-i18n agents.promptOnClipboardBody)" +fi + # A new session uses the current selection rather than its parent's account # pin. Explicit provider homes still override the selection. Carry the new pin # in the command, since terminals launched through uwsm don't inherit our env. diff --git a/bin/omarchy-default-agent b/bin/omarchy-default-agent index 8aaa6186b09..ae194c129b9 100755 --- a/bin/omarchy-default-agent +++ b/bin/omarchy-default-agent @@ -1,7 +1,7 @@ #!/bin/bash # omarchy:summary=Set and launch the default coding agent -# omarchy:args=[pi|omp|opencode|ori|claude|codex|grok|openclaw|agy|hermes|copilot|crush|cursor-agent|muse|kimi|qwen|dsh|mcode|goose|kiro-cli|vibe] +# omarchy:args=[pi|omp|opencode|ori|claude|codex|grok|openclaw|agy|hermes|copilot|crush|cursor-agent|muse|kimi|qwen|dsh|mcode|goose|kiro-cli|vibe|freebuff] # omarchy:examples=omarchy default agent | omarchy default agent codex | omarchy default agent claude installing=false @@ -59,8 +59,9 @@ kiro-cli | kiro) agent_package="http:kiro-cli[platforms_linux_x64_url=https://prod.download.cli.kiro.dev/stable/{{version}}/kirocli-x86_64-linux.tar.gz,platforms_linux_arm64_url=https://prod.download.cli.kiro.dev/stable/{{version}}/kirocli-aarch64-linux.tar.gz,version_list_url=https://prod.download.cli.kiro.dev/stable/latest/manifest.json,version_json_path=.version]" ;; vibe | mistral | mistral-vibe) agent="vibe"; name="Mistral Vibe"; agent_package="github:mistralai/mistral-vibe[matching=vibe-linux]" ;; +freebuff) agent="freebuff"; name="Freebuff"; agent_package="npm:freebuff" ;; *) - echo "Usage: omarchy-default-agent " + echo "Usage: omarchy-default-agent " exit 1 ;; esac diff --git a/bin/omarchy-remove-preinstalls b/bin/omarchy-remove-preinstalls index c50b149072c..d07d3793d81 100755 --- a/bin/omarchy-remove-preinstalls +++ b/bin/omarchy-remove-preinstalls @@ -72,6 +72,12 @@ if gum confirm --affirmative "$(omarchy-i18n common.yes)" --negative "$(omarchy- rm -f ~/.local/bin/vibe fi + # Preserve a user-managed Freebuff launcher at the same path. + if [[ -f ~/.local/bin/freebuff && ! -L ~/.local/bin/freebuff ]] && + grep -Eq '^mise use -g .*"npm:freebuff"' ~/.local/bin/freebuff; then + rm -f ~/.local/bin/freebuff + fi + omarchy-pkg-drop \ aether \ cliamp \ diff --git a/default/fonts/omarchy/README.md b/default/fonts/omarchy/README.md index c67eacf5e42..c342fa4474b 100644 --- a/default/fonts/omarchy/README.md +++ b/default/fonts/omarchy/README.md @@ -23,6 +23,7 @@ The private-use glyphs in `omarchy.ttf` are: - `U+E912` — Goose, from - `U+E913` — Kiro, from the ghost in , with its tile dropped - `U+E914` — Mistral AI, from +- `U+E915` — Freebuff, from , its rects joined into one path The agent marks are monochrome so the menu can render them using the active theme's foreground and selection colors. diff --git a/default/fonts/omarchy/omarchy.ttf b/default/fonts/omarchy/omarchy.ttf index 0cca4ccbee2f6cd0119a6ecfc9dc58764352894a..cf881e5a8362a4386b5eb9590a4096beac30504b 100644 GIT binary patch delta 361 zcmX?;a3OJmVSTZnF9QQ(2Ll6>S9(rmT7l)Ai3|+PR~Q)NWinC|Q}otZ^)WCgePCc< zFv|c6a4@pSGcYJ?0Qo8zxg`}54D6Xez7LRJk&~aC7_m2q6UaXTI!4=2X++kBloLCKp@Zt@+alF2^G4qW`eZ~-fk+dN76JvUIMN#BZzQGD_q S{XQO!w4&71q|&rBpc(*_eOuxH delta 228 zcmcbRcqUS9(rmTAKM|R|W>=D+~;BYcf(3Q}otZ^)WCgH83zR zm}LM3I2c*vfczamzDh=JNks?)yC#r-2gonU$xlx7)DPlhU{DSK@@;YxD+(AbnGAsZ z9w1*KFEKaOxcg2mkiUk3fqP>?esKv;qHI$t1A~47P{FH>dt@0IMJC^3e9I{U"$mock_bin/omarchy-install-chromium-claude" <<'SH' @@ -36,6 +37,11 @@ cat >"$mock_bin/omarchy-notification-send" <<'SH' printf '%s\0' "$@" >>"$OMARCHY_TEST_NOTIFICATION_HISTORY" SH +cat >"$mock_bin/wl-copy" <<'SH' +#!/bin/bash +cat >"$OMARCHY_TEST_CLIPBOARD" +SH + cat >"$mock_bin/omarchy-cmd-missing" <<'SH' #!/bin/bash [[ $1 == ${OMARCHY_TEST_MISSING_COMMAND:-} ]] @@ -123,6 +129,7 @@ export OMARCHY_TEST_STUB_LOG="$stub_log" export OMARCHY_TEST_AGENT_TERMINAL_LOG="$terminal_log" export OMARCHY_TEST_AGENT_MENU_LOG="$menu_log" export OMARCHY_TEST_MUSE_LOGIN_LOG="$muse_login_log" +export OMARCHY_TEST_CLIPBOARD="$clipboard_log" export OMARCHY_PATH="$ROOT" grok_package="grok" @@ -138,6 +145,7 @@ mcode_package="npm:@minimax-ai/code[allow_builds=better-sqlite3]" goose_package="aqua:aaif-goose/goose" kiro_cli_package="http:kiro-cli[platforms_linux_x64_url=https://prod.download.cli.kiro.dev/stable/{{version}}/kirocli-x86_64-linux.tar.gz,platforms_linux_arm64_url=https://prod.download.cli.kiro.dev/stable/{{version}}/kirocli-aarch64-linux.tar.gz,version_list_url=https://prod.download.cli.kiro.dev/stable/latest/manifest.json,version_json_path=.version]" vibe_package="github:mistralai/mistral-vibe[matching=vibe-linux]" +freebuff_package="npm:freebuff" cursor_agent_package="cursor-agent" muse_package="http:muse[url=https://api.meta.ai/muse-launcher.sh,bin=muse,version_list_url=https://api.meta.ai/muse-code/channels/muse-stable,version_json_path=.version]" @@ -165,6 +173,7 @@ assert_lazy_stub "$mcode_package" mcode assert_lazy_stub "$goose_package" goose assert_lazy_stub "$kiro_cli_package" kiro-cli assert_lazy_stub "$vibe_package" vibe +assert_lazy_stub "$freebuff_package" freebuff assert_lazy_stub "$cursor_agent_package" cursor-agent assert_lazy_stub "$muse_package" muse pass "custom agent lazy stubs preserve their mise packages" @@ -184,6 +193,7 @@ grep -Fx "$mcode_package mcode" "$stub_log" >/dev/null || fail "user setup creat grep -Fx "$goose_package goose" "$stub_log" >/dev/null || fail "user setup creates the Goose lazy stub" grep -Fx "$kiro_cli_package kiro-cli" "$stub_log" >/dev/null || fail "user setup creates the Kiro CLI lazy stub" grep -Fx "$vibe_package vibe" "$stub_log" >/dev/null || fail "user setup creates the Mistral Vibe lazy stub" +grep -Fx "$freebuff_package freebuff" "$stub_log" >/dev/null || fail "user setup creates the Freebuff lazy stub" pass "user setup creates the custom agent lazy stubs" : >"$stub_log" @@ -197,6 +207,7 @@ grep -Fx "$mcode_package mcode" "$stub_log" >/dev/null || fail "user setup recre grep -Fx "$goose_package goose" "$stub_log" >/dev/null || fail "user setup recreates the Goose lazy stub" grep -Fx "$kiro_cli_package kiro-cli" "$stub_log" >/dev/null || fail "user setup recreates the Kiro CLI lazy stub" grep -Fx "$vibe_package vibe" "$stub_log" >/dev/null || fail "user setup recreates the Mistral Vibe lazy stub" +grep -Fx "$freebuff_package freebuff" "$stub_log" >/dev/null || fail "user setup recreates the Freebuff lazy stub" pass "user setup recreates wrappers even when the commands are present" : >"$stub_log" @@ -241,6 +252,10 @@ grep -Fx "$kiro_cli_package kiro-cli" "$stub_log" >/dev/null || fail "Kiro CLI m OMARCHY_TEST_MISSING_COMMAND=vibe source "$ROOT/migrations/1792033880.sh" >/dev/null grep -Fx "$vibe_package vibe" "$stub_log" >/dev/null || fail "Mistral Vibe migration creates its lazy stub" +: >"$stub_log" +OMARCHY_TEST_MISSING_COMMAND=freebuff source "$ROOT/migrations/1792033907.sh" >/dev/null +grep -Fx "$freebuff_package freebuff" "$stub_log" >/dev/null || fail "Freebuff migration creates its lazy stub" + : >"$stub_log" source "$ROOT/migrations/1785617047.sh" >/dev/null grep -Fx "$omp_package omp" "$stub_log" >/dev/null || fail "Oh My Pi migration creates a working lazy stub" @@ -350,6 +365,7 @@ OMARCHY_TEST_MISSING_COMMAND=mcode source "$ROOT/migrations/1792033557.sh" >/dev OMARCHY_TEST_MISSING_COMMAND=goose source "$ROOT/migrations/1792033725.sh" >/dev/null OMARCHY_TEST_MISSING_COMMAND=kiro-cli source "$ROOT/migrations/1792033795.sh" >/dev/null OMARCHY_TEST_MISSING_COMMAND=vibe source "$ROOT/migrations/1792033880.sh" >/dev/null +OMARCHY_TEST_MISSING_COMMAND=freebuff source "$ROOT/migrations/1792033907.sh" >/dev/null [[ ! -s $stub_log ]] || fail "agent migrations respect the preinstall opt-out" [[ ! -e $test_home/.local/bin/omp ]] || fail "agent migration removes the obsolete Oh My Pi wrapper after opt-out" @@ -489,9 +505,10 @@ pass "agent migrations install working wrappers without overriding the preinstal "$ROOT/bin/omarchy-mise-install" "$goose_package" goose "$ROOT/bin/omarchy-mise-install" "$kiro_cli_package" kiro-cli "$ROOT/bin/omarchy-mise-install" "$vibe_package" vibe +"$ROOT/bin/omarchy-mise-install" "$freebuff_package" freebuff touch "$test_home/.local/bin/agy" "$test_home/.local/bin/ori" omarchy-remove-preinstalls >/dev/null -for command in agy omp ori grok crush cursor-agent muse kimi qwen dsh mcode goose kiro-cli vibe; do +for command in agy omp ori grok crush cursor-agent muse kimi qwen dsh mcode goose kiro-cli vibe freebuff; do [[ ! -e $test_home/.local/bin/$command ]] || fail "Remove Preinstalls deletes the $command lazy stub" done pass "Remove Preinstalls deletes every optional agent lazy stub" @@ -588,6 +605,7 @@ declare -A expected_agents=( [vibe]="vibe" [mistral]="vibe" [mistral-vibe]="vibe" + [freebuff]="freebuff" [cursor]="cursor-agent" [cursor-agent]="cursor-agent" [muse]="muse" @@ -613,6 +631,7 @@ declare -A expected_packages=( [goose]="$goose_package" [kiro-cli]="$kiro_cli_package" [vibe]="$vibe_package" + [freebuff]="$freebuff_package" [cursor-agent]="$cursor_agent_package" [muse]="$muse_package" ) @@ -895,6 +914,11 @@ assert_launch mcode mcode -- "Review this project" assert_launch goose env GOOSE_MODE=auto goose run --interactive "--text=Review this project" assert_launch kiro-cli kiro-cli chat --trust-all-tools -- "Review this project" assert_launch vibe vibe --auto-approve -- "Review this project" +: >"$notification_history" +assert_launch freebuff freebuff +[[ $(<"$clipboard_log") == "Review this project" ]] || fail "Freebuff launcher puts the prompt on the clipboard" +mapfile -d '' -t notification_args <"$notification_history" +[[ ${notification_args[*]} == *"Freebuff takes no starting prompt"* ]] || fail "Freebuff launcher says where the prompt went" pass "agent launcher adapts initial prompts for every supported agent" literal_muse_prompt=$'--disable-sandbox !Crash {$(touch must-not-run)}\ntrailing\\ ' @@ -930,6 +954,10 @@ assert_bypass mcode mcode assert_bypass goose env GOOSE_MODE=auto goose assert_bypass kiro-cli kiro-cli chat --trust-all-tools assert_bypass vibe vibe --auto-approve +: >"$clipboard_log" +: >"$notification_history" +assert_bypass freebuff freebuff +[[ ! -s $clipboard_log && ! -s $notification_history ]] || fail "Freebuff launcher leaves the clipboard alone without a prompt" pass "agent launcher skips permission prompts for every supported agent" printf '%s\n' "opencode" >"$agent_file" diff --git a/test/shell.d/freebuff-agent-test.sh b/test/shell.d/freebuff-agent-test.sh new file mode 100644 index 00000000000..359876f82dd --- /dev/null +++ b/test/shell.d/freebuff-agent-test.sh @@ -0,0 +1,85 @@ +#!/bin/bash + +set -euo pipefail +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" + +test_tmp=$(mktemp -d) +trap 'rm -rf "$test_tmp"' EXIT +mock_bin="$test_tmp/bin" +export HOME="$test_tmp/home with spaces" +export PATH="$mock_bin:$HOME/.local/bin:$ROOT/bin:/usr/local/bin:/usr/bin" +export OMARCHY_PATH="$ROOT" +export OMARCHY_TEST_MISE_LOG="$test_tmp/mise-log" +mkdir -p "$mock_bin" "$HOME/.local/bin" +freebuff="$HOME/.local/bin/freebuff" +package="npm:freebuff" + +cat >"$mock_bin/mise" <<'SH' +#!/bin/bash +printf '%s\0' "$@" >>"$OMARCHY_TEST_MISE_LOG" +SH +cat >"$mock_bin/noop" <<'SH' +#!/bin/bash +exit 0 +SH +chmod +x "$mock_bin/mise" "$mock_bin/noop" +for command in gum hyprctl omarchy-webapp-remove-all omarchy-tui-remove-all omarchy-pkg-drop; do + ln -s noop "$mock_bin/$command" +done + +apply_freebuff() { + case $1 in + migration) source "$ROOT/migrations/1792033907.sh" >/dev/null ;; + removal) omarchy-remove-preinstalls >/dev/null ;; + esac +} + +apply_freebuff migration +[[ -x $freebuff && ! -L $freebuff ]] || fail "Freebuff migration creates a missing wrapper" +[[ ! -s $OMARCHY_TEST_MISE_LOG ]] || fail "Freebuff migration installs the package before first launch" +cp "$freebuff" "$test_tmp/wrapper" +apply_freebuff migration +cmp -s "$freebuff" "$test_tmp/wrapper" || fail "Freebuff migration is idempotent" +"$freebuff" --prompt 'Review this project' +mapfile -d '' -t args <"$OMARCHY_TEST_MISE_LOG" +[[ ${args[*]} == "use -g --quiet $package x $package -- freebuff --prompt Review this project" ]] || + fail "Freebuff migration wrapper installs the official package and forwards arguments" +pass "Freebuff migration creates a lazy, repeatable wrapper" + +rm -f "$freebuff" +mkdir -p "$HOME/.local/state/omarchy" +touch "$HOME/.local/state/omarchy/preinstalls-removed" +apply_freebuff migration +[[ ! -e $freebuff ]] || fail "Freebuff migration respects removed preinstalls with a missing command" +rm "$HOME/.local/state/omarchy/preinstalls-removed" +pass "Freebuff migration respects the preinstall opt-out" + +for operation in migration removal; do + for kind in script binary symlink; do + rm -f "$freebuff" + case $kind in + script) printf '#!/bin/bash\necho user-managed freebuff\n' >"$freebuff" ;; + binary) cp /usr/bin/true "$freebuff" ;; + symlink) ln -s /usr/bin/true "$freebuff" ;; + esac + [[ $kind == "symlink" ]] || chmod +x "$freebuff" + rm -f "$test_tmp/saved-$kind" + cp -P "$freebuff" "$test_tmp/saved-$kind" + for attempt in 1 2; do + apply_freebuff "$operation" + cmp -s "$freebuff" "$test_tmp/saved-$kind" || fail "Freebuff $operation preserves a user $kind" + if [[ $kind == "symlink" ]]; then + [[ -L $freebuff && $(readlink "$freebuff") == "/usr/bin/true" ]] || + fail "Freebuff $operation preserves the user's symlink" + else + [[ ! -L $freebuff ]] || fail "Freebuff $operation keeps the user's $kind as a regular file" + fi + done + pass "Freebuff $operation preserves a user $kind on repeated runs" + done +done + +omarchy-mise-install "$package" freebuff +apply_freebuff removal +[[ ! -e $freebuff ]] || fail "Remove Preinstalls removes Omarchy's Freebuff wrapper" +pass "Remove Preinstalls removes Omarchy's Freebuff wrapper" diff --git a/test/shell.d/menu-test.sh b/test/shell.d/menu-test.sh index 49d3bf0f2ca..50ba69abae8 100644 --- a/test/shell.d/menu-test.sh +++ b/test/shell.d/menu-test.sh @@ -259,6 +259,7 @@ const expectedAgents = { goose: { icon: '\ue912', iconFont: 'omarchy', label: 'Goose' }, 'kiro-cli': { icon: '\ue913', iconFont: 'omarchy', label: 'Kiro CLI' }, vibe: { icon: '\ue914', iconFont: 'omarchy', label: 'Mistral Vibe' }, + freebuff: { icon: '\ue915', iconFont: 'omarchy', label: 'Freebuff' }, hermes: { icon: '\ue90a', iconFont: 'omarchy', label: 'Hermes' }, openclaw: { icon: '\ue90c', iconFont: 'omarchy', label: 'OpenClaw' }, copilot: { icon: '', label: 'Copilot' }, @@ -284,7 +285,7 @@ assertDeepEqual( defaultItems .filter(item => item.parent === 'setup.default.agent') .map(item => item.label), - ['Antigravity', 'Claude', 'Codex', 'Copilot', 'Crush', 'Cursor CLI', 'DeepSeek Harness', 'Goose', 'Grok', 'Hermes', 'Kimi', 'Kiro CLI', 'MiniMax Code', 'Mistral Vibe', 'Muse Code', 'omp', 'OpenClaw', 'OpenCode', 'Ori', 'Pi', 'Qwen Code'], + ['Antigravity', 'Claude', 'Codex', 'Copilot', 'Crush', 'Cursor CLI', 'DeepSeek Harness', 'Freebuff', 'Goose', 'Grok', 'Hermes', 'Kimi', 'Kiro CLI', 'MiniMax Code', 'Mistral Vibe', 'Muse Code', 'omp', 'OpenClaw', 'OpenCode', 'Ori', 'Pi', 'Qwen Code'], 'menu sorts coding agents alphabetically' ) const expectedDefaults = { @@ -881,5 +882,5 @@ assert( JS font_charset=$(fc-query --format='%{charset}' "$ROOT/default/fonts/omarchy/omarchy.ttf") -[[ $font_charset == *"e900-e914"* ]] || fail "Omarchy icon font includes every custom menu glyph" +[[ $font_charset == *"e900-e915"* ]] || fail "Omarchy icon font includes every custom menu glyph" pass "Omarchy icon font includes the official agent marks" diff --git a/test/shell.d/mise-provision-test.sh b/test/shell.d/mise-provision-test.sh index faa175d501f..4c641df3575 100644 --- a/test/shell.d/mise-provision-test.sh +++ b/test/shell.d/mise-provision-test.sh @@ -17,7 +17,7 @@ SH chmod +x "$mock_bin/mise" source "$ROOT/install/user/mise.sh" -for agent in kimi cursor-agent muse qwen dsh mcode goose kiro-cli vibe; do +for agent in kimi cursor-agent muse qwen dsh mcode goose kiro-cli vibe freebuff; do [[ -x $HOME/.local/bin/$agent && ! -L $HOME/.local/bin/$agent ]] || fail "provisioning creates the missing $agent wrapper" cp "$HOME/.local/bin/$agent" "$test_tmp/$agent-wrapper" @@ -27,7 +27,7 @@ mapfile -d '' -t args <"$OMARCHY_TEST_MISE_LOG" fail "provisioning creates lazy wrappers without installing packages" pass "provisioning creates lazy Kimi, Cursor, and Muse wrappers" -for agent in kimi cursor-agent muse qwen dsh mcode goose kiro-cli vibe; do +for agent in kimi cursor-agent muse qwen dsh mcode goose kiro-cli vibe freebuff; do launcher="$HOME/.local/bin/$agent" for kind in script binary symlink; do rm -f "$launcher" From 15fd7f29d76729a72f3fc263431a9d73713cc41a Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:15:51 -0400 Subject: [PATCH 103/546] Add ZeroClaw as a selectable default coding agent ZeroClaw (zeroclaw) joins the default agent picker. mise installs its GitHub release for x86_64 or aarch64 in place of the PR's install-script wrapper, which would have put the binary in ~/.cargo/bin and edited the shell profile. The archive also carries zerocode, ZeroClaw's terminal UI, which is what the launcher opens: a bare zeroclaw only prints usage. zerocode takes no starting prompt, so a prompt goes to the clipboard the way Freebuff's does. ZeroClaw publishes no monochrome mark, so the menu uses the Nerd Font Z box that echoes its app icon. From #13993. Co-authored-by: Jonathan Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent | 6 ++ bin/omarchy-default-agent | 5 +- bin/omarchy-remove-preinstalls | 6 ++ default/i18n/en/menu.json | 3 + default/omarchy/omarchy-menu.jsonc | 1 + install/user/mise.sh | 1 + manual/17-ai.md | 3 +- migrations/1792034058.sh | 5 ++ test/shell.d/default-agent-test.sh | 16 +++++- test/shell.d/menu-test.sh | 3 +- test/shell.d/mise-provision-test.sh | 4 +- test/shell.d/zeroclaw-agent-test.sh | 85 +++++++++++++++++++++++++++++ 12 files changed, 131 insertions(+), 7 deletions(-) create mode 100644 migrations/1792034058.sh create mode 100644 test/shell.d/zeroclaw-agent-test.sh diff --git a/bin/omarchy-agent b/bin/omarchy-agent index f694b700414..898d5d229e6 100755 --- a/bin/omarchy-agent +++ b/bin/omarchy-agent @@ -185,6 +185,12 @@ freebuff) command=(freebuff) prompt_via_clipboard="Freebuff" ;; +zeroclaw) + # The terminal UI is zerocode, shipped in the same release. It takes no + # starting prompt, and autonomy is ZeroClaw configuration, not a flag. + command=(zerocode) + prompt_via_clipboard="ZeroClaw" + ;; *) echo "Unsupported default agent: $agent" >&2 exit 1 diff --git a/bin/omarchy-default-agent b/bin/omarchy-default-agent index ae194c129b9..e2df25347dc 100755 --- a/bin/omarchy-default-agent +++ b/bin/omarchy-default-agent @@ -1,7 +1,7 @@ #!/bin/bash # omarchy:summary=Set and launch the default coding agent -# omarchy:args=[pi|omp|opencode|ori|claude|codex|grok|openclaw|agy|hermes|copilot|crush|cursor-agent|muse|kimi|qwen|dsh|mcode|goose|kiro-cli|vibe|freebuff] +# omarchy:args=[pi|omp|opencode|ori|claude|codex|grok|openclaw|agy|hermes|copilot|crush|cursor-agent|muse|kimi|qwen|dsh|mcode|goose|kiro-cli|vibe|freebuff|zeroclaw] # omarchy:examples=omarchy default agent | omarchy default agent codex | omarchy default agent claude installing=false @@ -60,8 +60,9 @@ kiro-cli | kiro) ;; vibe | mistral | mistral-vibe) agent="vibe"; name="Mistral Vibe"; agent_package="github:mistralai/mistral-vibe[matching=vibe-linux]" ;; freebuff) agent="freebuff"; name="Freebuff"; agent_package="npm:freebuff" ;; +zeroclaw) agent="zeroclaw"; name="ZeroClaw"; agent_package="github:zeroclaw-labs/zeroclaw" ;; *) - echo "Usage: omarchy-default-agent " + echo "Usage: omarchy-default-agent " exit 1 ;; esac diff --git a/bin/omarchy-remove-preinstalls b/bin/omarchy-remove-preinstalls index d07d3793d81..aeab51ab8ce 100755 --- a/bin/omarchy-remove-preinstalls +++ b/bin/omarchy-remove-preinstalls @@ -78,6 +78,12 @@ if gum confirm --affirmative "$(omarchy-i18n common.yes)" --negative "$(omarchy- rm -f ~/.local/bin/freebuff fi + # Preserve a user-managed ZeroClaw launcher at the same path. + if [[ -f ~/.local/bin/zeroclaw && ! -L ~/.local/bin/zeroclaw ]] && + grep -Eq '^mise use -g .*"github:zeroclaw-labs/zeroclaw"' ~/.local/bin/zeroclaw; then + rm -f ~/.local/bin/zeroclaw + fi + omarchy-pkg-drop \ aether \ cliamp \ diff --git a/default/i18n/en/menu.json b/default/i18n/en/menu.json index a62f499ea00..08ebf4968e2 100644 --- a/default/i18n/en/menu.json +++ b/default/i18n/en/menu.json @@ -548,6 +548,9 @@ "setup.default.agent.vibe": { "label": "Mistral Vibe" }, + "setup.default.agent.zeroclaw": { + "label": "ZeroClaw" + }, "setup.default.browser": { "label": "Browser", "title": "Default Browser" diff --git a/default/omarchy/omarchy-menu.jsonc b/default/omarchy/omarchy-menu.jsonc index 725bca39cb1..9bcc04c4e2f 100644 --- a/default/omarchy/omarchy-menu.jsonc +++ b/default/omarchy/omarchy-menu.jsonc @@ -212,6 +212,7 @@ "setup.default.agent.ori": {"icon":"","iconFont":"omarchy","label":"Ori","checked":"[[ \"$(omarchy-default-agent)\" == \"ori\" ]]","action":"omarchy-default-agent ori"}, "setup.default.agent.pi": {"icon":"","iconFont":"omarchy","label":"Pi","checked":"[[ \"$(omarchy-default-agent)\" == \"pi\" ]]","action":"omarchy-default-agent pi"}, "setup.default.agent.qwen": {"icon":"","iconFont":"omarchy","label":"Qwen Code","checked":"[[ \"$(omarchy-default-agent)\" == \"qwen\" ]]","action":"omarchy-default-agent qwen"}, + "setup.default.agent.zeroclaw": {"icon":"󰬡","label":"ZeroClaw","checked":"[[ \"$(omarchy-default-agent)\" == \"zeroclaw\" ]]","action":"omarchy-default-agent zeroclaw"}, "setup.default.browser": {"icon":"","label":"Browser","title":"Default Browser"}, "setup.default.browser.chromium": {"icon":"","label":"Chromium","checked":"[[ \"$(omarchy-default-browser)\" == \"chromium\" ]]","action":"omarchy-default-browser chromium"}, "setup.default.browser.chrome": {"icon":"󰊯","label":"Chrome","checked":"[[ \"$(omarchy-default-browser)\" == \"chrome\" ]]","action":"omarchy-default-browser chrome"}, diff --git a/install/user/mise.sh b/install/user/mise.sh index a191affa383..5256597ece8 100644 --- a/install/user/mise.sh +++ b/install/user/mise.sh @@ -22,6 +22,7 @@ omarchy-mise-install aqua:aaif-goose/goose goose omarchy-mise-install "http:kiro-cli[platforms_linux_x64_url=https://prod.download.cli.kiro.dev/stable/{{version}}/kirocli-x86_64-linux.tar.gz,platforms_linux_arm64_url=https://prod.download.cli.kiro.dev/stable/{{version}}/kirocli-aarch64-linux.tar.gz,version_list_url=https://prod.download.cli.kiro.dev/stable/latest/manifest.json,version_json_path=.version]" kiro-cli omarchy-mise-install "github:mistralai/mistral-vibe[matching=vibe-linux]" vibe omarchy-mise-install npm:freebuff freebuff +omarchy-mise-install github:zeroclaw-labs/zeroclaw zeroclaw omarchy-mise-install cursor-agent omarchy-mise-install npm:@kitlangton/ghui ghui omarchy-mise-install aqua:modem-dev/hunk hunk diff --git a/manual/17-ai.md b/manual/17-ai.md index 45229cc855c..df0d5d9eb07 100644 --- a/manual/17-ai.md +++ b/manual/17-ai.md @@ -19,6 +19,7 @@ Omarchy treats AI coding agents as first-class citizens, but it doesn't pick a f | `mcode` | [MiniMax Code](https://github.com/MiniMax-AI/minimax-code), MiniMax's coding agent | | `qwen` | [Qwen Code](https://github.com/QwenLM/qwen-code), Alibaba's coding agent | | `vibe` | [Mistral Vibe](https://docs.mistral.ai/vibe/code/overview), Mistral AI's coding agent | +| `zeroclaw` | [ZeroClaw](https://github.com/zeroclaw-labs/zeroclaw), opened in its `zerocode` terminal UI | | `pi` | [Mario Zechner's Pi](https://github.com/badlogic/pi-mono) | | `omp` | [Oh My Pi](https://github.com/can1357/oh-my-pi) | | `ori` | [Ori](https://openrouter.ai/docs/guides/ori/harness), OpenRouter's harness | @@ -42,7 +43,7 @@ Choosing Claude also attempts to set up its browser extension for Chromium, Chro [Muse Code](https://dev.meta.ai) — Meta's `muse` — uses a preinstalled mise stub like the other agents. Picking it as the default installs Meta's official launcher through mise's HTTP backend. The launcher verifies and updates the native binary for your machine. -Once you've chosen, `Super + Shift + Ctrl + A` launches the default agent in a dedicated terminal window (or brings up the picker if you haven't chosen yet). You can also launch it straight into a task with `omarchy agent prompt "Review this project"`. Freebuff takes no starting prompt, so it opens with the prompt on your clipboard to paste instead. Agents launched this way run unattended in their respective don't-stop-to-ask modes, so be ready for them to actually do things! And since agents refuse to remember trust for your home directory, launches from `$HOME` start in `~/Work` instead. +Once you've chosen, `Super + Shift + Ctrl + A` launches the default agent in a dedicated terminal window (or brings up the picker if you haven't chosen yet). You can also launch it straight into a task with `omarchy agent prompt "Review this project"`. Freebuff and ZeroClaw take no starting prompt, so they open with the prompt on your clipboard to paste instead. Agents launched this way run unattended in their respective don't-stop-to-ask modes, so be ready for them to actually do things! And since agents refuse to remember trust for your home directory, launches from `$HOME` start in `~/Work` instead. There are terminal shortcuts too: `a` runs the default agent inline in the current terminal, while `c`, `cx`, and `cy` start OpenCode, Claude Code, and Codex directly (again in their auto-approving modes). Theme changes sync to the agents as well: Claude Code, Pi, OpenCode, and Hermes (once Omarchy has installed it) all follow along when you switch the Omarchy theme. diff --git a/migrations/1792034058.sh b/migrations/1792034058.sh new file mode 100644 index 00000000000..7bb022b44f3 --- /dev/null +++ b/migrations/1792034058.sh @@ -0,0 +1,5 @@ +echo "Install ZeroClaw coding agent mise wrapper" + +if [[ ! -f $HOME/.local/state/omarchy/preinstalls-removed ]] && omarchy-cmd-missing zeroclaw; then + omarchy-mise-install github:zeroclaw-labs/zeroclaw zeroclaw +fi diff --git a/test/shell.d/default-agent-test.sh b/test/shell.d/default-agent-test.sh index 8dbd4ef9947..a0a8910bb91 100644 --- a/test/shell.d/default-agent-test.sh +++ b/test/shell.d/default-agent-test.sh @@ -146,6 +146,7 @@ goose_package="aqua:aaif-goose/goose" kiro_cli_package="http:kiro-cli[platforms_linux_x64_url=https://prod.download.cli.kiro.dev/stable/{{version}}/kirocli-x86_64-linux.tar.gz,platforms_linux_arm64_url=https://prod.download.cli.kiro.dev/stable/{{version}}/kirocli-aarch64-linux.tar.gz,version_list_url=https://prod.download.cli.kiro.dev/stable/latest/manifest.json,version_json_path=.version]" vibe_package="github:mistralai/mistral-vibe[matching=vibe-linux]" freebuff_package="npm:freebuff" +zeroclaw_package="github:zeroclaw-labs/zeroclaw" cursor_agent_package="cursor-agent" muse_package="http:muse[url=https://api.meta.ai/muse-launcher.sh,bin=muse,version_list_url=https://api.meta.ai/muse-code/channels/muse-stable,version_json_path=.version]" @@ -174,6 +175,7 @@ assert_lazy_stub "$goose_package" goose assert_lazy_stub "$kiro_cli_package" kiro-cli assert_lazy_stub "$vibe_package" vibe assert_lazy_stub "$freebuff_package" freebuff +assert_lazy_stub "$zeroclaw_package" zeroclaw assert_lazy_stub "$cursor_agent_package" cursor-agent assert_lazy_stub "$muse_package" muse pass "custom agent lazy stubs preserve their mise packages" @@ -194,6 +196,7 @@ grep -Fx "$goose_package goose" "$stub_log" >/dev/null || fail "user setup creat grep -Fx "$kiro_cli_package kiro-cli" "$stub_log" >/dev/null || fail "user setup creates the Kiro CLI lazy stub" grep -Fx "$vibe_package vibe" "$stub_log" >/dev/null || fail "user setup creates the Mistral Vibe lazy stub" grep -Fx "$freebuff_package freebuff" "$stub_log" >/dev/null || fail "user setup creates the Freebuff lazy stub" +grep -Fx "$zeroclaw_package zeroclaw" "$stub_log" >/dev/null || fail "user setup creates the ZeroClaw lazy stub" pass "user setup creates the custom agent lazy stubs" : >"$stub_log" @@ -208,6 +211,7 @@ grep -Fx "$goose_package goose" "$stub_log" >/dev/null || fail "user setup recre grep -Fx "$kiro_cli_package kiro-cli" "$stub_log" >/dev/null || fail "user setup recreates the Kiro CLI lazy stub" grep -Fx "$vibe_package vibe" "$stub_log" >/dev/null || fail "user setup recreates the Mistral Vibe lazy stub" grep -Fx "$freebuff_package freebuff" "$stub_log" >/dev/null || fail "user setup recreates the Freebuff lazy stub" +grep -Fx "$zeroclaw_package zeroclaw" "$stub_log" >/dev/null || fail "user setup recreates the ZeroClaw lazy stub" pass "user setup recreates wrappers even when the commands are present" : >"$stub_log" @@ -256,6 +260,10 @@ grep -Fx "$vibe_package vibe" "$stub_log" >/dev/null || fail "Mistral Vibe migra OMARCHY_TEST_MISSING_COMMAND=freebuff source "$ROOT/migrations/1792033907.sh" >/dev/null grep -Fx "$freebuff_package freebuff" "$stub_log" >/dev/null || fail "Freebuff migration creates its lazy stub" +: >"$stub_log" +OMARCHY_TEST_MISSING_COMMAND=zeroclaw source "$ROOT/migrations/1792034058.sh" >/dev/null +grep -Fx "$zeroclaw_package zeroclaw" "$stub_log" >/dev/null || fail "ZeroClaw migration creates its lazy stub" + : >"$stub_log" source "$ROOT/migrations/1785617047.sh" >/dev/null grep -Fx "$omp_package omp" "$stub_log" >/dev/null || fail "Oh My Pi migration creates a working lazy stub" @@ -366,6 +374,7 @@ OMARCHY_TEST_MISSING_COMMAND=goose source "$ROOT/migrations/1792033725.sh" >/dev OMARCHY_TEST_MISSING_COMMAND=kiro-cli source "$ROOT/migrations/1792033795.sh" >/dev/null OMARCHY_TEST_MISSING_COMMAND=vibe source "$ROOT/migrations/1792033880.sh" >/dev/null OMARCHY_TEST_MISSING_COMMAND=freebuff source "$ROOT/migrations/1792033907.sh" >/dev/null +OMARCHY_TEST_MISSING_COMMAND=zeroclaw source "$ROOT/migrations/1792034058.sh" >/dev/null [[ ! -s $stub_log ]] || fail "agent migrations respect the preinstall opt-out" [[ ! -e $test_home/.local/bin/omp ]] || fail "agent migration removes the obsolete Oh My Pi wrapper after opt-out" @@ -506,9 +515,10 @@ pass "agent migrations install working wrappers without overriding the preinstal "$ROOT/bin/omarchy-mise-install" "$kiro_cli_package" kiro-cli "$ROOT/bin/omarchy-mise-install" "$vibe_package" vibe "$ROOT/bin/omarchy-mise-install" "$freebuff_package" freebuff +"$ROOT/bin/omarchy-mise-install" "$zeroclaw_package" zeroclaw touch "$test_home/.local/bin/agy" "$test_home/.local/bin/ori" omarchy-remove-preinstalls >/dev/null -for command in agy omp ori grok crush cursor-agent muse kimi qwen dsh mcode goose kiro-cli vibe freebuff; do +for command in agy omp ori grok crush cursor-agent muse kimi qwen dsh mcode goose kiro-cli vibe freebuff zeroclaw; do [[ ! -e $test_home/.local/bin/$command ]] || fail "Remove Preinstalls deletes the $command lazy stub" done pass "Remove Preinstalls deletes every optional agent lazy stub" @@ -606,6 +616,7 @@ declare -A expected_agents=( [mistral]="vibe" [mistral-vibe]="vibe" [freebuff]="freebuff" + [zeroclaw]="zeroclaw" [cursor]="cursor-agent" [cursor-agent]="cursor-agent" [muse]="muse" @@ -632,6 +643,7 @@ declare -A expected_packages=( [kiro-cli]="$kiro_cli_package" [vibe]="$vibe_package" [freebuff]="$freebuff_package" + [zeroclaw]="$zeroclaw_package" [cursor-agent]="$cursor_agent_package" [muse]="$muse_package" ) @@ -916,6 +928,7 @@ assert_launch kiro-cli kiro-cli chat --trust-all-tools -- "Review this project" assert_launch vibe vibe --auto-approve -- "Review this project" : >"$notification_history" assert_launch freebuff freebuff +assert_launch zeroclaw zerocode [[ $(<"$clipboard_log") == "Review this project" ]] || fail "Freebuff launcher puts the prompt on the clipboard" mapfile -d '' -t notification_args <"$notification_history" [[ ${notification_args[*]} == *"Freebuff takes no starting prompt"* ]] || fail "Freebuff launcher says where the prompt went" @@ -957,6 +970,7 @@ assert_bypass vibe vibe --auto-approve : >"$clipboard_log" : >"$notification_history" assert_bypass freebuff freebuff +assert_bypass zeroclaw zerocode [[ ! -s $clipboard_log && ! -s $notification_history ]] || fail "Freebuff launcher leaves the clipboard alone without a prompt" pass "agent launcher skips permission prompts for every supported agent" diff --git a/test/shell.d/menu-test.sh b/test/shell.d/menu-test.sh index 50ba69abae8..29b27839f76 100644 --- a/test/shell.d/menu-test.sh +++ b/test/shell.d/menu-test.sh @@ -260,6 +260,7 @@ const expectedAgents = { 'kiro-cli': { icon: '\ue913', iconFont: 'omarchy', label: 'Kiro CLI' }, vibe: { icon: '\ue914', iconFont: 'omarchy', label: 'Mistral Vibe' }, freebuff: { icon: '\ue915', iconFont: 'omarchy', label: 'Freebuff' }, + zeroclaw: { icon: '\u{f0b21}', label: 'ZeroClaw' }, hermes: { icon: '\ue90a', iconFont: 'omarchy', label: 'Hermes' }, openclaw: { icon: '\ue90c', iconFont: 'omarchy', label: 'OpenClaw' }, copilot: { icon: '', label: 'Copilot' }, @@ -285,7 +286,7 @@ assertDeepEqual( defaultItems .filter(item => item.parent === 'setup.default.agent') .map(item => item.label), - ['Antigravity', 'Claude', 'Codex', 'Copilot', 'Crush', 'Cursor CLI', 'DeepSeek Harness', 'Freebuff', 'Goose', 'Grok', 'Hermes', 'Kimi', 'Kiro CLI', 'MiniMax Code', 'Mistral Vibe', 'Muse Code', 'omp', 'OpenClaw', 'OpenCode', 'Ori', 'Pi', 'Qwen Code'], + ['Antigravity', 'Claude', 'Codex', 'Copilot', 'Crush', 'Cursor CLI', 'DeepSeek Harness', 'Freebuff', 'Goose', 'Grok', 'Hermes', 'Kimi', 'Kiro CLI', 'MiniMax Code', 'Mistral Vibe', 'Muse Code', 'omp', 'OpenClaw', 'OpenCode', 'Ori', 'Pi', 'Qwen Code', 'ZeroClaw'], 'menu sorts coding agents alphabetically' ) const expectedDefaults = { diff --git a/test/shell.d/mise-provision-test.sh b/test/shell.d/mise-provision-test.sh index 4c641df3575..8d1b986000f 100644 --- a/test/shell.d/mise-provision-test.sh +++ b/test/shell.d/mise-provision-test.sh @@ -17,7 +17,7 @@ SH chmod +x "$mock_bin/mise" source "$ROOT/install/user/mise.sh" -for agent in kimi cursor-agent muse qwen dsh mcode goose kiro-cli vibe freebuff; do +for agent in kimi cursor-agent muse qwen dsh mcode goose kiro-cli vibe freebuff zeroclaw; do [[ -x $HOME/.local/bin/$agent && ! -L $HOME/.local/bin/$agent ]] || fail "provisioning creates the missing $agent wrapper" cp "$HOME/.local/bin/$agent" "$test_tmp/$agent-wrapper" @@ -27,7 +27,7 @@ mapfile -d '' -t args <"$OMARCHY_TEST_MISE_LOG" fail "provisioning creates lazy wrappers without installing packages" pass "provisioning creates lazy Kimi, Cursor, and Muse wrappers" -for agent in kimi cursor-agent muse qwen dsh mcode goose kiro-cli vibe freebuff; do +for agent in kimi cursor-agent muse qwen dsh mcode goose kiro-cli vibe freebuff zeroclaw; do launcher="$HOME/.local/bin/$agent" for kind in script binary symlink; do rm -f "$launcher" diff --git a/test/shell.d/zeroclaw-agent-test.sh b/test/shell.d/zeroclaw-agent-test.sh new file mode 100644 index 00000000000..5653a7d347d --- /dev/null +++ b/test/shell.d/zeroclaw-agent-test.sh @@ -0,0 +1,85 @@ +#!/bin/bash + +set -euo pipefail +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" + +test_tmp=$(mktemp -d) +trap 'rm -rf "$test_tmp"' EXIT +mock_bin="$test_tmp/bin" +export HOME="$test_tmp/home with spaces" +export PATH="$mock_bin:$HOME/.local/bin:$ROOT/bin:/usr/local/bin:/usr/bin" +export OMARCHY_PATH="$ROOT" +export OMARCHY_TEST_MISE_LOG="$test_tmp/mise-log" +mkdir -p "$mock_bin" "$HOME/.local/bin" +zeroclaw="$HOME/.local/bin/zeroclaw" +package="github:zeroclaw-labs/zeroclaw" + +cat >"$mock_bin/mise" <<'SH' +#!/bin/bash +printf '%s\0' "$@" >>"$OMARCHY_TEST_MISE_LOG" +SH +cat >"$mock_bin/noop" <<'SH' +#!/bin/bash +exit 0 +SH +chmod +x "$mock_bin/mise" "$mock_bin/noop" +for command in gum hyprctl omarchy-webapp-remove-all omarchy-tui-remove-all omarchy-pkg-drop; do + ln -s noop "$mock_bin/$command" +done + +apply_zeroclaw() { + case $1 in + migration) source "$ROOT/migrations/1792034058.sh" >/dev/null ;; + removal) omarchy-remove-preinstalls >/dev/null ;; + esac +} + +apply_zeroclaw migration +[[ -x $zeroclaw && ! -L $zeroclaw ]] || fail "ZeroClaw migration creates a missing wrapper" +[[ ! -s $OMARCHY_TEST_MISE_LOG ]] || fail "ZeroClaw migration installs the package before first launch" +cp "$zeroclaw" "$test_tmp/wrapper" +apply_zeroclaw migration +cmp -s "$zeroclaw" "$test_tmp/wrapper" || fail "ZeroClaw migration is idempotent" +"$zeroclaw" --prompt 'Review this project' +mapfile -d '' -t args <"$OMARCHY_TEST_MISE_LOG" +[[ ${args[*]} == "use -g --quiet $package x $package -- zeroclaw --prompt Review this project" ]] || + fail "ZeroClaw migration wrapper installs the official package and forwards arguments" +pass "ZeroClaw migration creates a lazy, repeatable wrapper" + +rm -f "$zeroclaw" +mkdir -p "$HOME/.local/state/omarchy" +touch "$HOME/.local/state/omarchy/preinstalls-removed" +apply_zeroclaw migration +[[ ! -e $zeroclaw ]] || fail "ZeroClaw migration respects removed preinstalls with a missing command" +rm "$HOME/.local/state/omarchy/preinstalls-removed" +pass "ZeroClaw migration respects the preinstall opt-out" + +for operation in migration removal; do + for kind in script binary symlink; do + rm -f "$zeroclaw" + case $kind in + script) printf '#!/bin/bash\necho user-managed zeroclaw\n' >"$zeroclaw" ;; + binary) cp /usr/bin/true "$zeroclaw" ;; + symlink) ln -s /usr/bin/true "$zeroclaw" ;; + esac + [[ $kind == "symlink" ]] || chmod +x "$zeroclaw" + rm -f "$test_tmp/saved-$kind" + cp -P "$zeroclaw" "$test_tmp/saved-$kind" + for attempt in 1 2; do + apply_zeroclaw "$operation" + cmp -s "$zeroclaw" "$test_tmp/saved-$kind" || fail "ZeroClaw $operation preserves a user $kind" + if [[ $kind == "symlink" ]]; then + [[ -L $zeroclaw && $(readlink "$zeroclaw") == "/usr/bin/true" ]] || + fail "ZeroClaw $operation preserves the user's symlink" + else + [[ ! -L $zeroclaw ]] || fail "ZeroClaw $operation keeps the user's $kind as a regular file" + fi + done + pass "ZeroClaw $operation preserves a user $kind on repeated runs" + done +done + +omarchy-mise-install "$package" zeroclaw +apply_zeroclaw removal +[[ ! -e $zeroclaw ]] || fail "Remove Preinstalls removes Omarchy's ZeroClaw wrapper" +pass "Remove Preinstalls removes Omarchy's ZeroClaw wrapper" From 797f24a0dab75355af93fc16e32983e7d3d3a418 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:16:38 -0400 Subject: [PATCH 104/546] Add Kilo as a selectable default coding agent Kilo (kilo) joins the default agent picker, installed through mise as npm:@kilocode/cli with a lazy stub and migration like the other agents. It is built on OpenCode, so the launcher uses the same --auto, and seeds a prompt with --prompt= bound so a leading dash stays literal. From #11129. Co-authored-by: webreflection Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent | 6 ++ bin/omarchy-default-agent | 5 +- bin/omarchy-remove-preinstalls | 6 ++ default/fonts/omarchy/README.md | 1 + default/fonts/omarchy/omarchy.ttf | Bin 12496 -> 12680 bytes default/i18n/en/menu.json | 3 + default/omarchy/omarchy-menu.jsonc | 1 + install/user/mise.sh | 1 + manual/17-ai.md | 1 + migrations/1792034151.sh | 5 ++ test/shell.d/default-agent-test.sh | 18 +++++- test/shell.d/kilo-agent-test.sh | 85 ++++++++++++++++++++++++++++ test/shell.d/menu-test.sh | 5 +- test/shell.d/mise-provision-test.sh | 4 +- 14 files changed, 134 insertions(+), 7 deletions(-) create mode 100644 migrations/1792034151.sh create mode 100644 test/shell.d/kilo-agent-test.sh diff --git a/bin/omarchy-agent b/bin/omarchy-agent index 898d5d229e6..70a858f10ac 100755 --- a/bin/omarchy-agent +++ b/bin/omarchy-agent @@ -191,6 +191,12 @@ zeroclaw) command=(zerocode) prompt_via_clipboard="ZeroClaw" ;; +kilo) + # Kilo is built on OpenCode and shares its flags. Bind the prompt so one + # starting with a dash stays literal. + command=(kilo --auto) + [[ -n ${prompt:-} ]] && command+=("--prompt=$prompt") + ;; *) echo "Unsupported default agent: $agent" >&2 exit 1 diff --git a/bin/omarchy-default-agent b/bin/omarchy-default-agent index e2df25347dc..4270d2ec02b 100755 --- a/bin/omarchy-default-agent +++ b/bin/omarchy-default-agent @@ -1,7 +1,7 @@ #!/bin/bash # omarchy:summary=Set and launch the default coding agent -# omarchy:args=[pi|omp|opencode|ori|claude|codex|grok|openclaw|agy|hermes|copilot|crush|cursor-agent|muse|kimi|qwen|dsh|mcode|goose|kiro-cli|vibe|freebuff|zeroclaw] +# omarchy:args=[pi|omp|opencode|ori|claude|codex|grok|openclaw|agy|hermes|copilot|crush|cursor-agent|muse|kimi|qwen|dsh|mcode|goose|kiro-cli|vibe|freebuff|zeroclaw|kilo] # omarchy:examples=omarchy default agent | omarchy default agent codex | omarchy default agent claude installing=false @@ -61,8 +61,9 @@ kiro-cli | kiro) vibe | mistral | mistral-vibe) agent="vibe"; name="Mistral Vibe"; agent_package="github:mistralai/mistral-vibe[matching=vibe-linux]" ;; freebuff) agent="freebuff"; name="Freebuff"; agent_package="npm:freebuff" ;; zeroclaw) agent="zeroclaw"; name="ZeroClaw"; agent_package="github:zeroclaw-labs/zeroclaw" ;; +kilo | kilocode | kilo-code) agent="kilo"; name="Kilo"; agent_package="npm:@kilocode/cli" ;; *) - echo "Usage: omarchy-default-agent " + echo "Usage: omarchy-default-agent " exit 1 ;; esac diff --git a/bin/omarchy-remove-preinstalls b/bin/omarchy-remove-preinstalls index aeab51ab8ce..13f6f5f60f7 100755 --- a/bin/omarchy-remove-preinstalls +++ b/bin/omarchy-remove-preinstalls @@ -84,6 +84,12 @@ if gum confirm --affirmative "$(omarchy-i18n common.yes)" --negative "$(omarchy- rm -f ~/.local/bin/zeroclaw fi + # Preserve a user-managed Kilo launcher at the same path. + if [[ -f ~/.local/bin/kilo && ! -L ~/.local/bin/kilo ]] && + grep -Eq '^mise use -g .*"npm:@kilocode/cli"' ~/.local/bin/kilo; then + rm -f ~/.local/bin/kilo + fi + omarchy-pkg-drop \ aether \ cliamp \ diff --git a/default/fonts/omarchy/README.md b/default/fonts/omarchy/README.md index c342fa4474b..3144c927aee 100644 --- a/default/fonts/omarchy/README.md +++ b/default/fonts/omarchy/README.md @@ -24,6 +24,7 @@ The private-use glyphs in `omarchy.ttf` are: - `U+E913` — Kiro, from the ghost in , with its tile dropped - `U+E914` — Mistral AI, from - `U+E915` — Freebuff, from , its rects joined into one path +- `U+E916` — Kilo, from The agent marks are monochrome so the menu can render them using the active theme's foreground and selection colors. diff --git a/default/fonts/omarchy/omarchy.ttf b/default/fonts/omarchy/omarchy.ttf index cf881e5a8362a4386b5eb9590a4096beac30504b..6d9c49ee8fbb016fd74e39b48a9d84b2b2e3a93b 100644 GIT binary patch delta 413 zcmW-cJxBs!7{{OY?o9i};j9MnIB>Cxrk|NH%)2j2Hxucj`S%0>4Y0I32n$-V5+?)+x+ z8-Q&C_}+}PzO(Qcy8sYw03c$Za1qXrY$4AYxngMxocG8!@^Lnw(l?uW7poF-BBz%M zWSI#gKO_4O^jumx`#M3!`_XzKe^?xQ+JD6|e2xot{l!a&$NVHe!yaTFA2!tZU^vof zvtD0$7!r_RTWkwhaKZ%kGvJ4Lh(HWC-2t~F;0_RhpOR!DIw$ZFFNJVA`z#9~d7x0p zeJf*&CW?`5*%fx>k7Mj!hl&~PN;EVDBelebnnLcx?PSvGZNwEtqg1O=Iyk5(e;gUl zh|~}j?Y2K_r~OFjm8GhxqN+gf6GW1Nan3O9;5zY{t3t_a35sj7GbnwHk=gDj`rpMk PL#EA0_}ankXR`T!i!N@@ delta 231 zcmeB3zK}S(k22KG!K-v`LA$jMJmjMy8*3FIFE^6hdHD+(B`m<)hQ zS%7?nyu{p8/dev/null || fail "user setup grep -Fx "$vibe_package vibe" "$stub_log" >/dev/null || fail "user setup creates the Mistral Vibe lazy stub" grep -Fx "$freebuff_package freebuff" "$stub_log" >/dev/null || fail "user setup creates the Freebuff lazy stub" grep -Fx "$zeroclaw_package zeroclaw" "$stub_log" >/dev/null || fail "user setup creates the ZeroClaw lazy stub" +grep -Fx "$kilo_package kilo" "$stub_log" >/dev/null || fail "user setup creates the Kilo lazy stub" pass "user setup creates the custom agent lazy stubs" : >"$stub_log" @@ -212,6 +215,7 @@ grep -Fx "$kiro_cli_package kiro-cli" "$stub_log" >/dev/null || fail "user setup grep -Fx "$vibe_package vibe" "$stub_log" >/dev/null || fail "user setup recreates the Mistral Vibe lazy stub" grep -Fx "$freebuff_package freebuff" "$stub_log" >/dev/null || fail "user setup recreates the Freebuff lazy stub" grep -Fx "$zeroclaw_package zeroclaw" "$stub_log" >/dev/null || fail "user setup recreates the ZeroClaw lazy stub" +grep -Fx "$kilo_package kilo" "$stub_log" >/dev/null || fail "user setup recreates the Kilo lazy stub" pass "user setup recreates wrappers even when the commands are present" : >"$stub_log" @@ -264,6 +268,10 @@ grep -Fx "$freebuff_package freebuff" "$stub_log" >/dev/null || fail "Freebuff m OMARCHY_TEST_MISSING_COMMAND=zeroclaw source "$ROOT/migrations/1792034058.sh" >/dev/null grep -Fx "$zeroclaw_package zeroclaw" "$stub_log" >/dev/null || fail "ZeroClaw migration creates its lazy stub" +: >"$stub_log" +OMARCHY_TEST_MISSING_COMMAND=kilo source "$ROOT/migrations/1792034151.sh" >/dev/null +grep -Fx "$kilo_package kilo" "$stub_log" >/dev/null || fail "Kilo migration creates its lazy stub" + : >"$stub_log" source "$ROOT/migrations/1785617047.sh" >/dev/null grep -Fx "$omp_package omp" "$stub_log" >/dev/null || fail "Oh My Pi migration creates a working lazy stub" @@ -375,6 +383,7 @@ OMARCHY_TEST_MISSING_COMMAND=kiro-cli source "$ROOT/migrations/1792033795.sh" >/ OMARCHY_TEST_MISSING_COMMAND=vibe source "$ROOT/migrations/1792033880.sh" >/dev/null OMARCHY_TEST_MISSING_COMMAND=freebuff source "$ROOT/migrations/1792033907.sh" >/dev/null OMARCHY_TEST_MISSING_COMMAND=zeroclaw source "$ROOT/migrations/1792034058.sh" >/dev/null +OMARCHY_TEST_MISSING_COMMAND=kilo source "$ROOT/migrations/1792034151.sh" >/dev/null [[ ! -s $stub_log ]] || fail "agent migrations respect the preinstall opt-out" [[ ! -e $test_home/.local/bin/omp ]] || fail "agent migration removes the obsolete Oh My Pi wrapper after opt-out" @@ -516,9 +525,10 @@ pass "agent migrations install working wrappers without overriding the preinstal "$ROOT/bin/omarchy-mise-install" "$vibe_package" vibe "$ROOT/bin/omarchy-mise-install" "$freebuff_package" freebuff "$ROOT/bin/omarchy-mise-install" "$zeroclaw_package" zeroclaw +"$ROOT/bin/omarchy-mise-install" "$kilo_package" kilo touch "$test_home/.local/bin/agy" "$test_home/.local/bin/ori" omarchy-remove-preinstalls >/dev/null -for command in agy omp ori grok crush cursor-agent muse kimi qwen dsh mcode goose kiro-cli vibe freebuff zeroclaw; do +for command in agy omp ori grok crush cursor-agent muse kimi qwen dsh mcode goose kiro-cli vibe freebuff zeroclaw kilo; do [[ ! -e $test_home/.local/bin/$command ]] || fail "Remove Preinstalls deletes the $command lazy stub" done pass "Remove Preinstalls deletes every optional agent lazy stub" @@ -617,6 +627,9 @@ declare -A expected_agents=( [mistral-vibe]="vibe" [freebuff]="freebuff" [zeroclaw]="zeroclaw" + [kilo]="kilo" + [kilocode]="kilo" + [kilo-code]="kilo" [cursor]="cursor-agent" [cursor-agent]="cursor-agent" [muse]="muse" @@ -644,6 +657,7 @@ declare -A expected_packages=( [vibe]="$vibe_package" [freebuff]="$freebuff_package" [zeroclaw]="$zeroclaw_package" + [kilo]="$kilo_package" [cursor-agent]="$cursor_agent_package" [muse]="$muse_package" ) @@ -929,6 +943,7 @@ assert_launch vibe vibe --auto-approve -- "Review this project" : >"$notification_history" assert_launch freebuff freebuff assert_launch zeroclaw zerocode +assert_launch kilo kilo --auto "--prompt=Review this project" [[ $(<"$clipboard_log") == "Review this project" ]] || fail "Freebuff launcher puts the prompt on the clipboard" mapfile -d '' -t notification_args <"$notification_history" [[ ${notification_args[*]} == *"Freebuff takes no starting prompt"* ]] || fail "Freebuff launcher says where the prompt went" @@ -971,6 +986,7 @@ assert_bypass vibe vibe --auto-approve : >"$notification_history" assert_bypass freebuff freebuff assert_bypass zeroclaw zerocode +assert_bypass kilo kilo --auto [[ ! -s $clipboard_log && ! -s $notification_history ]] || fail "Freebuff launcher leaves the clipboard alone without a prompt" pass "agent launcher skips permission prompts for every supported agent" diff --git a/test/shell.d/kilo-agent-test.sh b/test/shell.d/kilo-agent-test.sh new file mode 100644 index 00000000000..29e520f2620 --- /dev/null +++ b/test/shell.d/kilo-agent-test.sh @@ -0,0 +1,85 @@ +#!/bin/bash + +set -euo pipefail +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" + +test_tmp=$(mktemp -d) +trap 'rm -rf "$test_tmp"' EXIT +mock_bin="$test_tmp/bin" +export HOME="$test_tmp/home with spaces" +export PATH="$mock_bin:$HOME/.local/bin:$ROOT/bin:/usr/local/bin:/usr/bin" +export OMARCHY_PATH="$ROOT" +export OMARCHY_TEST_MISE_LOG="$test_tmp/mise-log" +mkdir -p "$mock_bin" "$HOME/.local/bin" +kilo="$HOME/.local/bin/kilo" +package="npm:@kilocode/cli" + +cat >"$mock_bin/mise" <<'SH' +#!/bin/bash +printf '%s\0' "$@" >>"$OMARCHY_TEST_MISE_LOG" +SH +cat >"$mock_bin/noop" <<'SH' +#!/bin/bash +exit 0 +SH +chmod +x "$mock_bin/mise" "$mock_bin/noop" +for command in gum hyprctl omarchy-webapp-remove-all omarchy-tui-remove-all omarchy-pkg-drop; do + ln -s noop "$mock_bin/$command" +done + +apply_kilo() { + case $1 in + migration) source "$ROOT/migrations/1792034151.sh" >/dev/null ;; + removal) omarchy-remove-preinstalls >/dev/null ;; + esac +} + +apply_kilo migration +[[ -x $kilo && ! -L $kilo ]] || fail "Kilo migration creates a missing wrapper" +[[ ! -s $OMARCHY_TEST_MISE_LOG ]] || fail "Kilo migration installs the package before first launch" +cp "$kilo" "$test_tmp/wrapper" +apply_kilo migration +cmp -s "$kilo" "$test_tmp/wrapper" || fail "Kilo migration is idempotent" +"$kilo" --prompt 'Review this project' +mapfile -d '' -t args <"$OMARCHY_TEST_MISE_LOG" +[[ ${args[*]} == "use -g --quiet $package x $package -- kilo --prompt Review this project" ]] || + fail "Kilo migration wrapper installs the official package and forwards arguments" +pass "Kilo migration creates a lazy, repeatable wrapper" + +rm -f "$kilo" +mkdir -p "$HOME/.local/state/omarchy" +touch "$HOME/.local/state/omarchy/preinstalls-removed" +apply_kilo migration +[[ ! -e $kilo ]] || fail "Kilo migration respects removed preinstalls with a missing command" +rm "$HOME/.local/state/omarchy/preinstalls-removed" +pass "Kilo migration respects the preinstall opt-out" + +for operation in migration removal; do + for kind in script binary symlink; do + rm -f "$kilo" + case $kind in + script) printf '#!/bin/bash\necho user-managed kilo\n' >"$kilo" ;; + binary) cp /usr/bin/true "$kilo" ;; + symlink) ln -s /usr/bin/true "$kilo" ;; + esac + [[ $kind == "symlink" ]] || chmod +x "$kilo" + rm -f "$test_tmp/saved-$kind" + cp -P "$kilo" "$test_tmp/saved-$kind" + for attempt in 1 2; do + apply_kilo "$operation" + cmp -s "$kilo" "$test_tmp/saved-$kind" || fail "Kilo $operation preserves a user $kind" + if [[ $kind == "symlink" ]]; then + [[ -L $kilo && $(readlink "$kilo") == "/usr/bin/true" ]] || + fail "Kilo $operation preserves the user's symlink" + else + [[ ! -L $kilo ]] || fail "Kilo $operation keeps the user's $kind as a regular file" + fi + done + pass "Kilo $operation preserves a user $kind on repeated runs" + done +done + +omarchy-mise-install "$package" kilo +apply_kilo removal +[[ ! -e $kilo ]] || fail "Remove Preinstalls removes Omarchy's Kilo wrapper" +pass "Remove Preinstalls removes Omarchy's Kilo wrapper" diff --git a/test/shell.d/menu-test.sh b/test/shell.d/menu-test.sh index 29b27839f76..961a1753d91 100644 --- a/test/shell.d/menu-test.sh +++ b/test/shell.d/menu-test.sh @@ -261,6 +261,7 @@ const expectedAgents = { vibe: { icon: '\ue914', iconFont: 'omarchy', label: 'Mistral Vibe' }, freebuff: { icon: '\ue915', iconFont: 'omarchy', label: 'Freebuff' }, zeroclaw: { icon: '\u{f0b21}', label: 'ZeroClaw' }, + kilo: { icon: '\ue916', iconFont: 'omarchy', label: 'Kilo' }, hermes: { icon: '\ue90a', iconFont: 'omarchy', label: 'Hermes' }, openclaw: { icon: '\ue90c', iconFont: 'omarchy', label: 'OpenClaw' }, copilot: { icon: '', label: 'Copilot' }, @@ -286,7 +287,7 @@ assertDeepEqual( defaultItems .filter(item => item.parent === 'setup.default.agent') .map(item => item.label), - ['Antigravity', 'Claude', 'Codex', 'Copilot', 'Crush', 'Cursor CLI', 'DeepSeek Harness', 'Freebuff', 'Goose', 'Grok', 'Hermes', 'Kimi', 'Kiro CLI', 'MiniMax Code', 'Mistral Vibe', 'Muse Code', 'omp', 'OpenClaw', 'OpenCode', 'Ori', 'Pi', 'Qwen Code', 'ZeroClaw'], + ['Antigravity', 'Claude', 'Codex', 'Copilot', 'Crush', 'Cursor CLI', 'DeepSeek Harness', 'Freebuff', 'Goose', 'Grok', 'Hermes', 'Kilo', 'Kimi', 'Kiro CLI', 'MiniMax Code', 'Mistral Vibe', 'Muse Code', 'omp', 'OpenClaw', 'OpenCode', 'Ori', 'Pi', 'Qwen Code', 'ZeroClaw'], 'menu sorts coding agents alphabetically' ) const expectedDefaults = { @@ -883,5 +884,5 @@ assert( JS font_charset=$(fc-query --format='%{charset}' "$ROOT/default/fonts/omarchy/omarchy.ttf") -[[ $font_charset == *"e900-e915"* ]] || fail "Omarchy icon font includes every custom menu glyph" +[[ $font_charset == *"e900-e916"* ]] || fail "Omarchy icon font includes every custom menu glyph" pass "Omarchy icon font includes the official agent marks" diff --git a/test/shell.d/mise-provision-test.sh b/test/shell.d/mise-provision-test.sh index 8d1b986000f..c866790a5e8 100644 --- a/test/shell.d/mise-provision-test.sh +++ b/test/shell.d/mise-provision-test.sh @@ -17,7 +17,7 @@ SH chmod +x "$mock_bin/mise" source "$ROOT/install/user/mise.sh" -for agent in kimi cursor-agent muse qwen dsh mcode goose kiro-cli vibe freebuff zeroclaw; do +for agent in kimi cursor-agent muse qwen dsh mcode goose kiro-cli vibe freebuff zeroclaw kilo; do [[ -x $HOME/.local/bin/$agent && ! -L $HOME/.local/bin/$agent ]] || fail "provisioning creates the missing $agent wrapper" cp "$HOME/.local/bin/$agent" "$test_tmp/$agent-wrapper" @@ -27,7 +27,7 @@ mapfile -d '' -t args <"$OMARCHY_TEST_MISE_LOG" fail "provisioning creates lazy wrappers without installing packages" pass "provisioning creates lazy Kimi, Cursor, and Muse wrappers" -for agent in kimi cursor-agent muse qwen dsh mcode goose kiro-cli vibe freebuff zeroclaw; do +for agent in kimi cursor-agent muse qwen dsh mcode goose kiro-cli vibe freebuff zeroclaw kilo; do launcher="$HOME/.local/bin/$agent" for kind in script binary symlink; do rm -f "$launcher" From c14de24baffed7226d9bfc79e4e4530e4ad86ab7 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:17:23 -0400 Subject: [PATCH 105/546] Add Devin CLI as a selectable default coding agent Devin CLI (devin) joins the default agent picker. mise's HTTP backend installs Cognition's per-platform release archive for x86_64 or aarch64, resolving the version from its current manifest, in place of a wrapper around the curl | bash installer. The launcher runs it with the dangerous permission mode and without the workspace trust check, and passes a prompt after --, since a bare positional opens Devin Desktop. The launcher part of #8124; its usage collector is left to the agents panel work. Co-authored-by: Joao Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent | 6 ++ bin/omarchy-default-agent | 9 ++- bin/omarchy-remove-preinstalls | 6 ++ default/fonts/omarchy/README.md | 1 + default/fonts/omarchy/omarchy.ttf | Bin 12680 -> 13108 bytes default/i18n/en/menu.json | 3 + default/omarchy/omarchy-menu.jsonc | 1 + install/user/mise.sh | 1 + manual/17-ai.md | 1 + migrations/1792034198.sh | 5 ++ test/shell.d/default-agent-test.sh | 17 +++++- test/shell.d/devin-agent-test.sh | 85 ++++++++++++++++++++++++++++ test/shell.d/menu-test.sh | 5 +- test/shell.d/mise-provision-test.sh | 4 +- 14 files changed, 137 insertions(+), 7 deletions(-) create mode 100644 migrations/1792034198.sh create mode 100644 test/shell.d/devin-agent-test.sh diff --git a/bin/omarchy-agent b/bin/omarchy-agent index 70a858f10ac..56885400ff9 100755 --- a/bin/omarchy-agent +++ b/bin/omarchy-agent @@ -197,6 +197,12 @@ kilo) command=(kilo --auto) [[ -n ${prompt:-} ]] && command+=("--prompt=$prompt") ;; +devin) + # A bare positional is a path that opens Devin Desktop, so the prompt + # always follows --, which starts an interactive session with it. + command=(devin --permission-mode dangerous --respect-workspace-trust false) + [[ -n ${prompt:-} ]] && command+=(-- "$prompt") + ;; *) echo "Unsupported default agent: $agent" >&2 exit 1 diff --git a/bin/omarchy-default-agent b/bin/omarchy-default-agent index 4270d2ec02b..9773caeeb97 100755 --- a/bin/omarchy-default-agent +++ b/bin/omarchy-default-agent @@ -1,7 +1,7 @@ #!/bin/bash # omarchy:summary=Set and launch the default coding agent -# omarchy:args=[pi|omp|opencode|ori|claude|codex|grok|openclaw|agy|hermes|copilot|crush|cursor-agent|muse|kimi|qwen|dsh|mcode|goose|kiro-cli|vibe|freebuff|zeroclaw|kilo] +# omarchy:args=[pi|omp|opencode|ori|claude|codex|grok|openclaw|agy|hermes|copilot|crush|cursor-agent|muse|kimi|qwen|dsh|mcode|goose|kiro-cli|vibe|freebuff|zeroclaw|kilo|devin] # omarchy:examples=omarchy default agent | omarchy default agent codex | omarchy default agent claude installing=false @@ -62,8 +62,13 @@ vibe | mistral | mistral-vibe) agent="vibe"; name="Mistral Vibe"; agent_package= freebuff) agent="freebuff"; name="Freebuff"; agent_package="npm:freebuff" ;; zeroclaw) agent="zeroclaw"; name="ZeroClaw"; agent_package="github:zeroclaw-labs/zeroclaw" ;; kilo | kilocode | kilo-code) agent="kilo"; name="Kilo"; agent_package="npm:@kilocode/cli" ;; +devin | devin-cli) + agent="devin"; name="Devin CLI" + # Cognition publishes per-platform archives behind a manifest of the current release. + agent_package="http:devin[platforms_linux_x64_url=https://static.devin.ai/cli/{{ version }}/devin-{{ version }}-x86_64-unknown-linux.tar.gz,platforms_linux_arm64_url=https://static.devin.ai/cli/{{ version }}/devin-{{ version }}-aarch64-unknown-linux.tar.gz,bin_path=bin,version_list_url=https://static.devin.ai/cli/current/manifest.json,version_json_path=.version]" + ;; *) - echo "Usage: omarchy-default-agent " + echo "Usage: omarchy-default-agent " exit 1 ;; esac diff --git a/bin/omarchy-remove-preinstalls b/bin/omarchy-remove-preinstalls index 13f6f5f60f7..03f1f4a8403 100755 --- a/bin/omarchy-remove-preinstalls +++ b/bin/omarchy-remove-preinstalls @@ -90,6 +90,12 @@ if gum confirm --affirmative "$(omarchy-i18n common.yes)" --negative "$(omarchy- rm -f ~/.local/bin/kilo fi + # Preserve a user-managed Devin CLI launcher at the same path. + if [[ -f ~/.local/bin/devin && ! -L ~/.local/bin/devin ]] && + grep -Eq '^mise use -g .*"http:devin\[' ~/.local/bin/devin; then + rm -f ~/.local/bin/devin + fi + omarchy-pkg-drop \ aether \ cliamp \ diff --git a/default/fonts/omarchy/README.md b/default/fonts/omarchy/README.md index 3144c927aee..1ee002715b3 100644 --- a/default/fonts/omarchy/README.md +++ b/default/fonts/omarchy/README.md @@ -25,6 +25,7 @@ The private-use glyphs in `omarchy.ttf` are: - `U+E914` — Mistral AI, from - `U+E915` — Freebuff, from , its rects joined into one path - `U+E916` — Kilo, from +- `U+E917` — Devin, from the official mark published by Cognition at The agent marks are monochrome so the menu can render them using the active theme's foreground and selection colors. diff --git a/default/fonts/omarchy/omarchy.ttf b/default/fonts/omarchy/omarchy.ttf index 6d9c49ee8fbb016fd74e39b48a9d84b2b2e3a93b..778dd272b8279d18ce3b064c6960eeeaf35a0318 100644 GIT binary patch delta 657 zcmZ8fO=}ZT6uo!eWG0zpprJGGWfs~r6EjLd=A)SbDNN;ri+3>tOV_xCVfS z0LtC1ohO_1^N9+;&Hxs>-A;Ss=JeJpfNLuN5X(Xq)kVU6!qslCKYRe$1;TF#Kik<` zZ*MDoy5 zZMp=yMjDy4oFrioCw^)t!ibABAx`4@vgGk6cO2eORlNKdPKt=dNwIjOGNveqSTH$v zN}|D3lqx!CpUjacZDO<5hBO-Sf{E0lFs2@5bGh7qGE`YZlpM~@0=X*pOmiuXz?NZ6 zEu(1|O@$#NOD3~g4!6yjMoA^}`KFwR0+y(5bEh>Mn>rY_GMVe-^T9G&a)eY0cU*ht*YM?(|YqKKWF=$P+#FZ!74GfBA2<%tmLh Gz55TLziZF{ delta 232 zcmdmz){#8Hu)ajlkAZ=)gMop`D?O(&O*>-mV+IE1D+~`5_40HyYJKj`9S^L8w>J_OL(?s9|lSq^Z*t7*|/dev/null || fail "user setup creates grep -Fx "$freebuff_package freebuff" "$stub_log" >/dev/null || fail "user setup creates the Freebuff lazy stub" grep -Fx "$zeroclaw_package zeroclaw" "$stub_log" >/dev/null || fail "user setup creates the ZeroClaw lazy stub" grep -Fx "$kilo_package kilo" "$stub_log" >/dev/null || fail "user setup creates the Kilo lazy stub" +grep -Fx "$devin_package devin" "$stub_log" >/dev/null || fail "user setup creates the Devin CLI lazy stub" pass "user setup creates the custom agent lazy stubs" : >"$stub_log" @@ -216,6 +219,7 @@ grep -Fx "$vibe_package vibe" "$stub_log" >/dev/null || fail "user setup recreat grep -Fx "$freebuff_package freebuff" "$stub_log" >/dev/null || fail "user setup recreates the Freebuff lazy stub" grep -Fx "$zeroclaw_package zeroclaw" "$stub_log" >/dev/null || fail "user setup recreates the ZeroClaw lazy stub" grep -Fx "$kilo_package kilo" "$stub_log" >/dev/null || fail "user setup recreates the Kilo lazy stub" +grep -Fx "$devin_package devin" "$stub_log" >/dev/null || fail "user setup recreates the Devin CLI lazy stub" pass "user setup recreates wrappers even when the commands are present" : >"$stub_log" @@ -272,6 +276,10 @@ grep -Fx "$zeroclaw_package zeroclaw" "$stub_log" >/dev/null || fail "ZeroClaw m OMARCHY_TEST_MISSING_COMMAND=kilo source "$ROOT/migrations/1792034151.sh" >/dev/null grep -Fx "$kilo_package kilo" "$stub_log" >/dev/null || fail "Kilo migration creates its lazy stub" +: >"$stub_log" +OMARCHY_TEST_MISSING_COMMAND=devin source "$ROOT/migrations/1792034198.sh" >/dev/null +grep -Fx "$devin_package devin" "$stub_log" >/dev/null || fail "Devin CLI migration creates its lazy stub" + : >"$stub_log" source "$ROOT/migrations/1785617047.sh" >/dev/null grep -Fx "$omp_package omp" "$stub_log" >/dev/null || fail "Oh My Pi migration creates a working lazy stub" @@ -384,6 +392,7 @@ OMARCHY_TEST_MISSING_COMMAND=vibe source "$ROOT/migrations/1792033880.sh" >/dev/ OMARCHY_TEST_MISSING_COMMAND=freebuff source "$ROOT/migrations/1792033907.sh" >/dev/null OMARCHY_TEST_MISSING_COMMAND=zeroclaw source "$ROOT/migrations/1792034058.sh" >/dev/null OMARCHY_TEST_MISSING_COMMAND=kilo source "$ROOT/migrations/1792034151.sh" >/dev/null +OMARCHY_TEST_MISSING_COMMAND=devin source "$ROOT/migrations/1792034198.sh" >/dev/null [[ ! -s $stub_log ]] || fail "agent migrations respect the preinstall opt-out" [[ ! -e $test_home/.local/bin/omp ]] || fail "agent migration removes the obsolete Oh My Pi wrapper after opt-out" @@ -526,9 +535,10 @@ pass "agent migrations install working wrappers without overriding the preinstal "$ROOT/bin/omarchy-mise-install" "$freebuff_package" freebuff "$ROOT/bin/omarchy-mise-install" "$zeroclaw_package" zeroclaw "$ROOT/bin/omarchy-mise-install" "$kilo_package" kilo +"$ROOT/bin/omarchy-mise-install" "$devin_package" devin touch "$test_home/.local/bin/agy" "$test_home/.local/bin/ori" omarchy-remove-preinstalls >/dev/null -for command in agy omp ori grok crush cursor-agent muse kimi qwen dsh mcode goose kiro-cli vibe freebuff zeroclaw kilo; do +for command in agy omp ori grok crush cursor-agent muse kimi qwen dsh mcode goose kiro-cli vibe freebuff zeroclaw kilo devin; do [[ ! -e $test_home/.local/bin/$command ]] || fail "Remove Preinstalls deletes the $command lazy stub" done pass "Remove Preinstalls deletes every optional agent lazy stub" @@ -630,6 +640,8 @@ declare -A expected_agents=( [kilo]="kilo" [kilocode]="kilo" [kilo-code]="kilo" + [devin]="devin" + [devin-cli]="devin" [cursor]="cursor-agent" [cursor-agent]="cursor-agent" [muse]="muse" @@ -658,6 +670,7 @@ declare -A expected_packages=( [freebuff]="$freebuff_package" [zeroclaw]="$zeroclaw_package" [kilo]="$kilo_package" + [devin]="$devin_package" [cursor-agent]="$cursor_agent_package" [muse]="$muse_package" ) @@ -944,6 +957,7 @@ assert_launch vibe vibe --auto-approve -- "Review this project" assert_launch freebuff freebuff assert_launch zeroclaw zerocode assert_launch kilo kilo --auto "--prompt=Review this project" +assert_launch devin devin --permission-mode dangerous --respect-workspace-trust false -- "Review this project" [[ $(<"$clipboard_log") == "Review this project" ]] || fail "Freebuff launcher puts the prompt on the clipboard" mapfile -d '' -t notification_args <"$notification_history" [[ ${notification_args[*]} == *"Freebuff takes no starting prompt"* ]] || fail "Freebuff launcher says where the prompt went" @@ -987,6 +1001,7 @@ assert_bypass vibe vibe --auto-approve assert_bypass freebuff freebuff assert_bypass zeroclaw zerocode assert_bypass kilo kilo --auto +assert_bypass devin devin --permission-mode dangerous --respect-workspace-trust false [[ ! -s $clipboard_log && ! -s $notification_history ]] || fail "Freebuff launcher leaves the clipboard alone without a prompt" pass "agent launcher skips permission prompts for every supported agent" diff --git a/test/shell.d/devin-agent-test.sh b/test/shell.d/devin-agent-test.sh new file mode 100644 index 00000000000..b9ad84f80f7 --- /dev/null +++ b/test/shell.d/devin-agent-test.sh @@ -0,0 +1,85 @@ +#!/bin/bash + +set -euo pipefail +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" + +test_tmp=$(mktemp -d) +trap 'rm -rf "$test_tmp"' EXIT +mock_bin="$test_tmp/bin" +export HOME="$test_tmp/home with spaces" +export PATH="$mock_bin:$HOME/.local/bin:$ROOT/bin:/usr/local/bin:/usr/bin" +export OMARCHY_PATH="$ROOT" +export OMARCHY_TEST_MISE_LOG="$test_tmp/mise-log" +mkdir -p "$mock_bin" "$HOME/.local/bin" +devin="$HOME/.local/bin/devin" +package="http:devin[platforms_linux_x64_url=https://static.devin.ai/cli/{{ version }}/devin-{{ version }}-x86_64-unknown-linux.tar.gz,platforms_linux_arm64_url=https://static.devin.ai/cli/{{ version }}/devin-{{ version }}-aarch64-unknown-linux.tar.gz,bin_path=bin,version_list_url=https://static.devin.ai/cli/current/manifest.json,version_json_path=.version]" + +cat >"$mock_bin/mise" <<'SH' +#!/bin/bash +printf '%s\0' "$@" >>"$OMARCHY_TEST_MISE_LOG" +SH +cat >"$mock_bin/noop" <<'SH' +#!/bin/bash +exit 0 +SH +chmod +x "$mock_bin/mise" "$mock_bin/noop" +for command in gum hyprctl omarchy-webapp-remove-all omarchy-tui-remove-all omarchy-pkg-drop; do + ln -s noop "$mock_bin/$command" +done + +apply_devin() { + case $1 in + migration) source "$ROOT/migrations/1792034198.sh" >/dev/null ;; + removal) omarchy-remove-preinstalls >/dev/null ;; + esac +} + +apply_devin migration +[[ -x $devin && ! -L $devin ]] || fail "Devin CLI migration creates a missing wrapper" +[[ ! -s $OMARCHY_TEST_MISE_LOG ]] || fail "Devin CLI migration installs the package before first launch" +cp "$devin" "$test_tmp/wrapper" +apply_devin migration +cmp -s "$devin" "$test_tmp/wrapper" || fail "Devin CLI migration is idempotent" +"$devin" --prompt 'Review this project' +mapfile -d '' -t args <"$OMARCHY_TEST_MISE_LOG" +[[ ${args[*]} == "use -g --quiet $package x $package -- devin --prompt Review this project" ]] || + fail "Devin CLI migration wrapper installs the official package and forwards arguments" +pass "Devin CLI migration creates a lazy, repeatable wrapper" + +rm -f "$devin" +mkdir -p "$HOME/.local/state/omarchy" +touch "$HOME/.local/state/omarchy/preinstalls-removed" +apply_devin migration +[[ ! -e $devin ]] || fail "Devin CLI migration respects removed preinstalls with a missing command" +rm "$HOME/.local/state/omarchy/preinstalls-removed" +pass "Devin CLI migration respects the preinstall opt-out" + +for operation in migration removal; do + for kind in script binary symlink; do + rm -f "$devin" + case $kind in + script) printf '#!/bin/bash\necho user-managed devin\n' >"$devin" ;; + binary) cp /usr/bin/true "$devin" ;; + symlink) ln -s /usr/bin/true "$devin" ;; + esac + [[ $kind == "symlink" ]] || chmod +x "$devin" + rm -f "$test_tmp/saved-$kind" + cp -P "$devin" "$test_tmp/saved-$kind" + for attempt in 1 2; do + apply_devin "$operation" + cmp -s "$devin" "$test_tmp/saved-$kind" || fail "Devin CLI $operation preserves a user $kind" + if [[ $kind == "symlink" ]]; then + [[ -L $devin && $(readlink "$devin") == "/usr/bin/true" ]] || + fail "Devin CLI $operation preserves the user's symlink" + else + [[ ! -L $devin ]] || fail "Devin CLI $operation keeps the user's $kind as a regular file" + fi + done + pass "Devin CLI $operation preserves a user $kind on repeated runs" + done +done + +omarchy-mise-install "$package" devin +apply_devin removal +[[ ! -e $devin ]] || fail "Remove Preinstalls removes Omarchy's Devin CLI wrapper" +pass "Remove Preinstalls removes Omarchy's Devin CLI wrapper" diff --git a/test/shell.d/menu-test.sh b/test/shell.d/menu-test.sh index 961a1753d91..839bc23af42 100644 --- a/test/shell.d/menu-test.sh +++ b/test/shell.d/menu-test.sh @@ -262,6 +262,7 @@ const expectedAgents = { freebuff: { icon: '\ue915', iconFont: 'omarchy', label: 'Freebuff' }, zeroclaw: { icon: '\u{f0b21}', label: 'ZeroClaw' }, kilo: { icon: '\ue916', iconFont: 'omarchy', label: 'Kilo' }, + devin: { icon: '\ue917', iconFont: 'omarchy', label: 'Devin CLI' }, hermes: { icon: '\ue90a', iconFont: 'omarchy', label: 'Hermes' }, openclaw: { icon: '\ue90c', iconFont: 'omarchy', label: 'OpenClaw' }, copilot: { icon: '', label: 'Copilot' }, @@ -287,7 +288,7 @@ assertDeepEqual( defaultItems .filter(item => item.parent === 'setup.default.agent') .map(item => item.label), - ['Antigravity', 'Claude', 'Codex', 'Copilot', 'Crush', 'Cursor CLI', 'DeepSeek Harness', 'Freebuff', 'Goose', 'Grok', 'Hermes', 'Kilo', 'Kimi', 'Kiro CLI', 'MiniMax Code', 'Mistral Vibe', 'Muse Code', 'omp', 'OpenClaw', 'OpenCode', 'Ori', 'Pi', 'Qwen Code', 'ZeroClaw'], + ['Antigravity', 'Claude', 'Codex', 'Copilot', 'Crush', 'Cursor CLI', 'DeepSeek Harness', 'Devin CLI', 'Freebuff', 'Goose', 'Grok', 'Hermes', 'Kilo', 'Kimi', 'Kiro CLI', 'MiniMax Code', 'Mistral Vibe', 'Muse Code', 'omp', 'OpenClaw', 'OpenCode', 'Ori', 'Pi', 'Qwen Code', 'ZeroClaw'], 'menu sorts coding agents alphabetically' ) const expectedDefaults = { @@ -884,5 +885,5 @@ assert( JS font_charset=$(fc-query --format='%{charset}' "$ROOT/default/fonts/omarchy/omarchy.ttf") -[[ $font_charset == *"e900-e916"* ]] || fail "Omarchy icon font includes every custom menu glyph" +[[ $font_charset == *"e900-e917"* ]] || fail "Omarchy icon font includes every custom menu glyph" pass "Omarchy icon font includes the official agent marks" diff --git a/test/shell.d/mise-provision-test.sh b/test/shell.d/mise-provision-test.sh index c866790a5e8..bbba06a6b92 100644 --- a/test/shell.d/mise-provision-test.sh +++ b/test/shell.d/mise-provision-test.sh @@ -17,7 +17,7 @@ SH chmod +x "$mock_bin/mise" source "$ROOT/install/user/mise.sh" -for agent in kimi cursor-agent muse qwen dsh mcode goose kiro-cli vibe freebuff zeroclaw kilo; do +for agent in kimi cursor-agent muse qwen dsh mcode goose kiro-cli vibe freebuff zeroclaw kilo devin; do [[ -x $HOME/.local/bin/$agent && ! -L $HOME/.local/bin/$agent ]] || fail "provisioning creates the missing $agent wrapper" cp "$HOME/.local/bin/$agent" "$test_tmp/$agent-wrapper" @@ -27,7 +27,7 @@ mapfile -d '' -t args <"$OMARCHY_TEST_MISE_LOG" fail "provisioning creates lazy wrappers without installing packages" pass "provisioning creates lazy Kimi, Cursor, and Muse wrappers" -for agent in kimi cursor-agent muse qwen dsh mcode goose kiro-cli vibe freebuff zeroclaw kilo; do +for agent in kimi cursor-agent muse qwen dsh mcode goose kiro-cli vibe freebuff zeroclaw kilo devin; do launcher="$HOME/.local/bin/$agent" for kind in script binary symlink; do rm -f "$launcher" From ddabe592e6d99066e178386a09ed07c27ac171c4 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:18:26 -0400 Subject: [PATCH 106/546] Add Amp as a selectable default coding agent Amp (amp) joins the default agent picker, installed through mise's registry amp package (npm:@ampcode/cli) instead of the curl | bash installer, with a lazy stub and migration like the other agents. The launcher runs it with --dangerously-allow-all, and seeds a prompt by piping it in, which Amp takes as the first message of an interactive session; a positional message only works with --execute, which exits after one turn. From #6844. Co-authored-by: Dave Ashworth Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent | 9 +++ bin/omarchy-default-agent | 5 +- bin/omarchy-remove-preinstalls | 6 ++ default/fonts/omarchy/README.md | 1 + default/fonts/omarchy/omarchy.ttf | Bin 13108 -> 13216 bytes default/i18n/en/menu.json | 3 + default/omarchy/omarchy-menu.jsonc | 1 + install/user/mise.sh | 1 + manual/17-ai.md | 1 + migrations/1792034243.sh | 5 ++ test/shell.d/amp-agent-test.sh | 85 ++++++++++++++++++++++++++++ test/shell.d/default-agent-test.sh | 17 +++++- test/shell.d/menu-test.sh | 5 +- test/shell.d/mise-provision-test.sh | 4 +- 14 files changed, 136 insertions(+), 7 deletions(-) create mode 100644 migrations/1792034243.sh create mode 100644 test/shell.d/amp-agent-test.sh diff --git a/bin/omarchy-agent b/bin/omarchy-agent index 56885400ff9..95c26a0425e 100755 --- a/bin/omarchy-agent +++ b/bin/omarchy-agent @@ -203,6 +203,15 @@ devin) command=(devin --permission-mode dangerous --respect-workspace-trust false) [[ -n ${prompt:-} ]] && command+=(-- "$prompt") ;; +amp) + # Amp takes its first interactive message from piped input; a positional + # message is for --execute, which exits after one turn. + if [[ -n ${prompt:-} ]]; then + command=(bash -c 'printf "%s\n" "$1" | exec amp --dangerously-allow-all' amp "$prompt") + else + command=(amp --dangerously-allow-all) + fi + ;; *) echo "Unsupported default agent: $agent" >&2 exit 1 diff --git a/bin/omarchy-default-agent b/bin/omarchy-default-agent index 9773caeeb97..c29da0c8ebe 100755 --- a/bin/omarchy-default-agent +++ b/bin/omarchy-default-agent @@ -1,7 +1,7 @@ #!/bin/bash # omarchy:summary=Set and launch the default coding agent -# omarchy:args=[pi|omp|opencode|ori|claude|codex|grok|openclaw|agy|hermes|copilot|crush|cursor-agent|muse|kimi|qwen|dsh|mcode|goose|kiro-cli|vibe|freebuff|zeroclaw|kilo|devin] +# omarchy:args=[pi|omp|opencode|ori|claude|codex|grok|openclaw|agy|hermes|copilot|crush|cursor-agent|muse|kimi|qwen|dsh|mcode|goose|kiro-cli|vibe|freebuff|zeroclaw|kilo|devin|amp] # omarchy:examples=omarchy default agent | omarchy default agent codex | omarchy default agent claude installing=false @@ -67,8 +67,9 @@ devin | devin-cli) # Cognition publishes per-platform archives behind a manifest of the current release. agent_package="http:devin[platforms_linux_x64_url=https://static.devin.ai/cli/{{ version }}/devin-{{ version }}-x86_64-unknown-linux.tar.gz,platforms_linux_arm64_url=https://static.devin.ai/cli/{{ version }}/devin-{{ version }}-aarch64-unknown-linux.tar.gz,bin_path=bin,version_list_url=https://static.devin.ai/cli/current/manifest.json,version_json_path=.version]" ;; +amp | ampcode) agent="amp"; name="Amp" ;; *) - echo "Usage: omarchy-default-agent " + echo "Usage: omarchy-default-agent " exit 1 ;; esac diff --git a/bin/omarchy-remove-preinstalls b/bin/omarchy-remove-preinstalls index 03f1f4a8403..8841c1ca398 100755 --- a/bin/omarchy-remove-preinstalls +++ b/bin/omarchy-remove-preinstalls @@ -96,6 +96,12 @@ if gum confirm --affirmative "$(omarchy-i18n common.yes)" --negative "$(omarchy- rm -f ~/.local/bin/devin fi + # Preserve a user-managed Amp launcher at the same path. + if [[ -f ~/.local/bin/amp && ! -L ~/.local/bin/amp ]] && + grep -Eq '^mise use -g .*"amp"' ~/.local/bin/amp; then + rm -f ~/.local/bin/amp + fi + omarchy-pkg-drop \ aether \ cliamp \ diff --git a/default/fonts/omarchy/README.md b/default/fonts/omarchy/README.md index 1ee002715b3..e5a6758bb51 100644 --- a/default/fonts/omarchy/README.md +++ b/default/fonts/omarchy/README.md @@ -26,6 +26,7 @@ The private-use glyphs in `omarchy.ttf` are: - `U+E915` — Freebuff, from , its rects joined into one path - `U+E916` — Kilo, from - `U+E917` — Devin, from the official mark published by Cognition at +- `U+E918` — Amp, from , its three paths joined into one The agent marks are monochrome so the menu can render them using the active theme's foreground and selection colors. diff --git a/default/fonts/omarchy/omarchy.ttf b/default/fonts/omarchy/omarchy.ttf index 778dd272b8279d18ce3b064c6960eeeaf35a0318..71fc91429f9c1a505f437cf2132f241db9fd8fe8 100644 GIT binary patch delta 367 zcmdmzwjh0iVSSlk00RSK2Ll6>S9(rm8rz=BFBuq^uP`tup2$c|Own6s)yKe~I){OQ z!7KwPz`@9(1mpwd8B{WIODa+rIHH04A3%OfPJVJ?W?|7~1_renAm1f7v7&&{mdOCf zKLF$_gN!MJ5);%~zN+73-N9Y#1Dv_b?w|;9%fm5CK}t&djcEu4>N8$j%HWnTm`2u7zCv zcTb*iDWgPD!oS%dhB$~J$G9wBo>2uT>sM6#pFy56HKC~Z-#90FfafB(Rg5e delta 236 zcmZ3Gz9nsfVSTBfKLZ0}2Ll6>S9(rmn(3`+EDQ|HR~Q%+Gcr;WQ}otZ^)WE0+AuIM zn1KX17+Dm7{1PBvB_p?_B7uQJ9>|{omy0=1fIyMxYp@vGH%=7P-kQRQ=f$fP4m)$rqXAC(EeCOs-LL;sB`tVuj7e S)ZTM5DooyH"$mock_bin/mise" <<'SH' +#!/bin/bash +printf '%s\0' "$@" >>"$OMARCHY_TEST_MISE_LOG" +SH +cat >"$mock_bin/noop" <<'SH' +#!/bin/bash +exit 0 +SH +chmod +x "$mock_bin/mise" "$mock_bin/noop" +for command in gum hyprctl omarchy-webapp-remove-all omarchy-tui-remove-all omarchy-pkg-drop; do + ln -s noop "$mock_bin/$command" +done + +apply_amp() { + case $1 in + migration) source "$ROOT/migrations/1792034243.sh" >/dev/null ;; + removal) omarchy-remove-preinstalls >/dev/null ;; + esac +} + +apply_amp migration +[[ -x $amp && ! -L $amp ]] || fail "Amp migration creates a missing wrapper" +[[ ! -s $OMARCHY_TEST_MISE_LOG ]] || fail "Amp migration installs the package before first launch" +cp "$amp" "$test_tmp/wrapper" +apply_amp migration +cmp -s "$amp" "$test_tmp/wrapper" || fail "Amp migration is idempotent" +"$amp" --prompt 'Review this project' +mapfile -d '' -t args <"$OMARCHY_TEST_MISE_LOG" +[[ ${args[*]} == "use -g --quiet $package x $package -- amp --prompt Review this project" ]] || + fail "Amp migration wrapper installs the official package and forwards arguments" +pass "Amp migration creates a lazy, repeatable wrapper" + +rm -f "$amp" +mkdir -p "$HOME/.local/state/omarchy" +touch "$HOME/.local/state/omarchy/preinstalls-removed" +apply_amp migration +[[ ! -e $amp ]] || fail "Amp migration respects removed preinstalls with a missing command" +rm "$HOME/.local/state/omarchy/preinstalls-removed" +pass "Amp migration respects the preinstall opt-out" + +for operation in migration removal; do + for kind in script binary symlink; do + rm -f "$amp" + case $kind in + script) printf '#!/bin/bash\necho user-managed amp\n' >"$amp" ;; + binary) cp /usr/bin/true "$amp" ;; + symlink) ln -s /usr/bin/true "$amp" ;; + esac + [[ $kind == "symlink" ]] || chmod +x "$amp" + rm -f "$test_tmp/saved-$kind" + cp -P "$amp" "$test_tmp/saved-$kind" + for attempt in 1 2; do + apply_amp "$operation" + cmp -s "$amp" "$test_tmp/saved-$kind" || fail "Amp $operation preserves a user $kind" + if [[ $kind == "symlink" ]]; then + [[ -L $amp && $(readlink "$amp") == "/usr/bin/true" ]] || + fail "Amp $operation preserves the user's symlink" + else + [[ ! -L $amp ]] || fail "Amp $operation keeps the user's $kind as a regular file" + fi + done + pass "Amp $operation preserves a user $kind on repeated runs" + done +done + +omarchy-mise-install "$package" amp +apply_amp removal +[[ ! -e $amp ]] || fail "Remove Preinstalls removes Omarchy's Amp wrapper" +pass "Remove Preinstalls removes Omarchy's Amp wrapper" diff --git a/test/shell.d/default-agent-test.sh b/test/shell.d/default-agent-test.sh index c3cad648429..c58f9c160f3 100644 --- a/test/shell.d/default-agent-test.sh +++ b/test/shell.d/default-agent-test.sh @@ -149,6 +149,7 @@ freebuff_package="npm:freebuff" zeroclaw_package="github:zeroclaw-labs/zeroclaw" kilo_package="npm:@kilocode/cli" devin_package="http:devin[platforms_linux_x64_url=https://static.devin.ai/cli/{{ version }}/devin-{{ version }}-x86_64-unknown-linux.tar.gz,platforms_linux_arm64_url=https://static.devin.ai/cli/{{ version }}/devin-{{ version }}-aarch64-unknown-linux.tar.gz,bin_path=bin,version_list_url=https://static.devin.ai/cli/current/manifest.json,version_json_path=.version]" +amp_package="amp" cursor_agent_package="cursor-agent" muse_package="http:muse[url=https://api.meta.ai/muse-launcher.sh,bin=muse,version_list_url=https://api.meta.ai/muse-code/channels/muse-stable,version_json_path=.version]" @@ -180,6 +181,7 @@ assert_lazy_stub "$freebuff_package" freebuff assert_lazy_stub "$zeroclaw_package" zeroclaw assert_lazy_stub "$kilo_package" kilo assert_lazy_stub "$devin_package" devin +assert_lazy_stub "$amp_package" amp assert_lazy_stub "$cursor_agent_package" cursor-agent assert_lazy_stub "$muse_package" muse pass "custom agent lazy stubs preserve their mise packages" @@ -203,6 +205,7 @@ grep -Fx "$freebuff_package freebuff" "$stub_log" >/dev/null || fail "user setup grep -Fx "$zeroclaw_package zeroclaw" "$stub_log" >/dev/null || fail "user setup creates the ZeroClaw lazy stub" grep -Fx "$kilo_package kilo" "$stub_log" >/dev/null || fail "user setup creates the Kilo lazy stub" grep -Fx "$devin_package devin" "$stub_log" >/dev/null || fail "user setup creates the Devin CLI lazy stub" +grep -Fx "$amp_package" "$stub_log" >/dev/null || fail "user setup creates the Amp lazy stub" pass "user setup creates the custom agent lazy stubs" : >"$stub_log" @@ -220,6 +223,7 @@ grep -Fx "$freebuff_package freebuff" "$stub_log" >/dev/null || fail "user setup grep -Fx "$zeroclaw_package zeroclaw" "$stub_log" >/dev/null || fail "user setup recreates the ZeroClaw lazy stub" grep -Fx "$kilo_package kilo" "$stub_log" >/dev/null || fail "user setup recreates the Kilo lazy stub" grep -Fx "$devin_package devin" "$stub_log" >/dev/null || fail "user setup recreates the Devin CLI lazy stub" +grep -Fx "$amp_package" "$stub_log" >/dev/null || fail "user setup recreates the Amp lazy stub" pass "user setup recreates wrappers even when the commands are present" : >"$stub_log" @@ -280,6 +284,10 @@ grep -Fx "$kilo_package kilo" "$stub_log" >/dev/null || fail "Kilo migration cre OMARCHY_TEST_MISSING_COMMAND=devin source "$ROOT/migrations/1792034198.sh" >/dev/null grep -Fx "$devin_package devin" "$stub_log" >/dev/null || fail "Devin CLI migration creates its lazy stub" +: >"$stub_log" +OMARCHY_TEST_MISSING_COMMAND=amp source "$ROOT/migrations/1792034243.sh" >/dev/null +grep -Fx "$amp_package" "$stub_log" >/dev/null || fail "Amp migration creates its lazy stub" + : >"$stub_log" source "$ROOT/migrations/1785617047.sh" >/dev/null grep -Fx "$omp_package omp" "$stub_log" >/dev/null || fail "Oh My Pi migration creates a working lazy stub" @@ -393,6 +401,7 @@ OMARCHY_TEST_MISSING_COMMAND=freebuff source "$ROOT/migrations/1792033907.sh" >/ OMARCHY_TEST_MISSING_COMMAND=zeroclaw source "$ROOT/migrations/1792034058.sh" >/dev/null OMARCHY_TEST_MISSING_COMMAND=kilo source "$ROOT/migrations/1792034151.sh" >/dev/null OMARCHY_TEST_MISSING_COMMAND=devin source "$ROOT/migrations/1792034198.sh" >/dev/null +OMARCHY_TEST_MISSING_COMMAND=amp source "$ROOT/migrations/1792034243.sh" >/dev/null [[ ! -s $stub_log ]] || fail "agent migrations respect the preinstall opt-out" [[ ! -e $test_home/.local/bin/omp ]] || fail "agent migration removes the obsolete Oh My Pi wrapper after opt-out" @@ -536,9 +545,10 @@ pass "agent migrations install working wrappers without overriding the preinstal "$ROOT/bin/omarchy-mise-install" "$zeroclaw_package" zeroclaw "$ROOT/bin/omarchy-mise-install" "$kilo_package" kilo "$ROOT/bin/omarchy-mise-install" "$devin_package" devin +"$ROOT/bin/omarchy-mise-install" "$amp_package" amp touch "$test_home/.local/bin/agy" "$test_home/.local/bin/ori" omarchy-remove-preinstalls >/dev/null -for command in agy omp ori grok crush cursor-agent muse kimi qwen dsh mcode goose kiro-cli vibe freebuff zeroclaw kilo devin; do +for command in agy omp ori grok crush cursor-agent muse kimi qwen dsh mcode goose kiro-cli vibe freebuff zeroclaw kilo devin amp; do [[ ! -e $test_home/.local/bin/$command ]] || fail "Remove Preinstalls deletes the $command lazy stub" done pass "Remove Preinstalls deletes every optional agent lazy stub" @@ -642,6 +652,8 @@ declare -A expected_agents=( [kilo-code]="kilo" [devin]="devin" [devin-cli]="devin" + [amp]="amp" + [ampcode]="amp" [cursor]="cursor-agent" [cursor-agent]="cursor-agent" [muse]="muse" @@ -671,6 +683,7 @@ declare -A expected_packages=( [zeroclaw]="$zeroclaw_package" [kilo]="$kilo_package" [devin]="$devin_package" + [amp]="$amp_package" [cursor-agent]="$cursor_agent_package" [muse]="$muse_package" ) @@ -958,6 +971,7 @@ assert_launch freebuff freebuff assert_launch zeroclaw zerocode assert_launch kilo kilo --auto "--prompt=Review this project" assert_launch devin devin --permission-mode dangerous --respect-workspace-trust false -- "Review this project" +assert_launch amp bash -c 'printf "%s\n" "$1" | exec amp --dangerously-allow-all' amp "Review this project" [[ $(<"$clipboard_log") == "Review this project" ]] || fail "Freebuff launcher puts the prompt on the clipboard" mapfile -d '' -t notification_args <"$notification_history" [[ ${notification_args[*]} == *"Freebuff takes no starting prompt"* ]] || fail "Freebuff launcher says where the prompt went" @@ -1002,6 +1016,7 @@ assert_bypass freebuff freebuff assert_bypass zeroclaw zerocode assert_bypass kilo kilo --auto assert_bypass devin devin --permission-mode dangerous --respect-workspace-trust false +assert_bypass amp amp --dangerously-allow-all [[ ! -s $clipboard_log && ! -s $notification_history ]] || fail "Freebuff launcher leaves the clipboard alone without a prompt" pass "agent launcher skips permission prompts for every supported agent" diff --git a/test/shell.d/menu-test.sh b/test/shell.d/menu-test.sh index 839bc23af42..d1f7ce5b74e 100644 --- a/test/shell.d/menu-test.sh +++ b/test/shell.d/menu-test.sh @@ -263,6 +263,7 @@ const expectedAgents = { zeroclaw: { icon: '\u{f0b21}', label: 'ZeroClaw' }, kilo: { icon: '\ue916', iconFont: 'omarchy', label: 'Kilo' }, devin: { icon: '\ue917', iconFont: 'omarchy', label: 'Devin CLI' }, + amp: { icon: '\ue918', iconFont: 'omarchy', label: 'Amp' }, hermes: { icon: '\ue90a', iconFont: 'omarchy', label: 'Hermes' }, openclaw: { icon: '\ue90c', iconFont: 'omarchy', label: 'OpenClaw' }, copilot: { icon: '', label: 'Copilot' }, @@ -288,7 +289,7 @@ assertDeepEqual( defaultItems .filter(item => item.parent === 'setup.default.agent') .map(item => item.label), - ['Antigravity', 'Claude', 'Codex', 'Copilot', 'Crush', 'Cursor CLI', 'DeepSeek Harness', 'Devin CLI', 'Freebuff', 'Goose', 'Grok', 'Hermes', 'Kilo', 'Kimi', 'Kiro CLI', 'MiniMax Code', 'Mistral Vibe', 'Muse Code', 'omp', 'OpenClaw', 'OpenCode', 'Ori', 'Pi', 'Qwen Code', 'ZeroClaw'], + ['Amp', 'Antigravity', 'Claude', 'Codex', 'Copilot', 'Crush', 'Cursor CLI', 'DeepSeek Harness', 'Devin CLI', 'Freebuff', 'Goose', 'Grok', 'Hermes', 'Kilo', 'Kimi', 'Kiro CLI', 'MiniMax Code', 'Mistral Vibe', 'Muse Code', 'omp', 'OpenClaw', 'OpenCode', 'Ori', 'Pi', 'Qwen Code', 'ZeroClaw'], 'menu sorts coding agents alphabetically' ) const expectedDefaults = { @@ -885,5 +886,5 @@ assert( JS font_charset=$(fc-query --format='%{charset}' "$ROOT/default/fonts/omarchy/omarchy.ttf") -[[ $font_charset == *"e900-e917"* ]] || fail "Omarchy icon font includes every custom menu glyph" +[[ $font_charset == *"e900-e918"* ]] || fail "Omarchy icon font includes every custom menu glyph" pass "Omarchy icon font includes the official agent marks" diff --git a/test/shell.d/mise-provision-test.sh b/test/shell.d/mise-provision-test.sh index bbba06a6b92..d3099f37f65 100644 --- a/test/shell.d/mise-provision-test.sh +++ b/test/shell.d/mise-provision-test.sh @@ -17,7 +17,7 @@ SH chmod +x "$mock_bin/mise" source "$ROOT/install/user/mise.sh" -for agent in kimi cursor-agent muse qwen dsh mcode goose kiro-cli vibe freebuff zeroclaw kilo devin; do +for agent in kimi cursor-agent muse qwen dsh mcode goose kiro-cli vibe freebuff zeroclaw kilo devin amp; do [[ -x $HOME/.local/bin/$agent && ! -L $HOME/.local/bin/$agent ]] || fail "provisioning creates the missing $agent wrapper" cp "$HOME/.local/bin/$agent" "$test_tmp/$agent-wrapper" @@ -27,7 +27,7 @@ mapfile -d '' -t args <"$OMARCHY_TEST_MISE_LOG" fail "provisioning creates lazy wrappers without installing packages" pass "provisioning creates lazy Kimi, Cursor, and Muse wrappers" -for agent in kimi cursor-agent muse qwen dsh mcode goose kiro-cli vibe freebuff zeroclaw kilo devin; do +for agent in kimi cursor-agent muse qwen dsh mcode goose kiro-cli vibe freebuff zeroclaw kilo devin amp; do launcher="$HOME/.local/bin/$agent" for kind in script binary symlink; do rm -f "$launcher" From aa8ccdb5534ffec9db8a6e39e30b6cc78e03977d Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:19:13 -0400 Subject: [PATCH 107/546] Add Command Code as a selectable default coding agent Command Code (commandcode) joins the default agent picker, installed through mise as npm:command-code with a lazy stub and migration like the other agents. The launcher runs it with --yolo, which also implies --trust, and seeds the interactive session with a prompt after --. From #10730, which also covers the Command Code part of #11839. Co-authored-by: Ahmad Awais Co-authored-by: HIMANSHU11827 Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent | 6 ++ bin/omarchy-default-agent | 5 +- bin/omarchy-remove-preinstalls | 6 ++ default/fonts/omarchy/README.md | 1 + default/fonts/omarchy/omarchy.ttf | Bin 13216 -> 13996 bytes default/i18n/en/menu.json | 3 + default/omarchy/omarchy-menu.jsonc | 1 + install/user/mise.sh | 1 + manual/17-ai.md | 1 + migrations/1792034306.sh | 5 ++ test/shell.d/commandcode-agent-test.sh | 85 +++++++++++++++++++++++++ test/shell.d/default-agent-test.sh | 18 +++++- test/shell.d/menu-test.sh | 5 +- test/shell.d/mise-provision-test.sh | 4 +- 14 files changed, 134 insertions(+), 7 deletions(-) create mode 100644 migrations/1792034306.sh create mode 100644 test/shell.d/commandcode-agent-test.sh diff --git a/bin/omarchy-agent b/bin/omarchy-agent index 95c26a0425e..fb0150602dc 100755 --- a/bin/omarchy-agent +++ b/bin/omarchy-agent @@ -212,6 +212,12 @@ amp) command=(amp --dangerously-allow-all) fi ;; +commandcode) + # --yolo skips permission prompts and implies --trust. A positional prompt + # seeds the interactive session; -- keeps one starting with a dash literal. + command=(commandcode --yolo) + [[ -n ${prompt:-} ]] && command+=(-- "$prompt") + ;; *) echo "Unsupported default agent: $agent" >&2 exit 1 diff --git a/bin/omarchy-default-agent b/bin/omarchy-default-agent index c29da0c8ebe..eebaee131e7 100755 --- a/bin/omarchy-default-agent +++ b/bin/omarchy-default-agent @@ -1,7 +1,7 @@ #!/bin/bash # omarchy:summary=Set and launch the default coding agent -# omarchy:args=[pi|omp|opencode|ori|claude|codex|grok|openclaw|agy|hermes|copilot|crush|cursor-agent|muse|kimi|qwen|dsh|mcode|goose|kiro-cli|vibe|freebuff|zeroclaw|kilo|devin|amp] +# omarchy:args=[pi|omp|opencode|ori|claude|codex|grok|openclaw|agy|hermes|copilot|crush|cursor-agent|muse|kimi|qwen|dsh|mcode|goose|kiro-cli|vibe|freebuff|zeroclaw|kilo|devin|amp|commandcode] # omarchy:examples=omarchy default agent | omarchy default agent codex | omarchy default agent claude installing=false @@ -68,8 +68,9 @@ devin | devin-cli) agent_package="http:devin[platforms_linux_x64_url=https://static.devin.ai/cli/{{ version }}/devin-{{ version }}-x86_64-unknown-linux.tar.gz,platforms_linux_arm64_url=https://static.devin.ai/cli/{{ version }}/devin-{{ version }}-aarch64-unknown-linux.tar.gz,bin_path=bin,version_list_url=https://static.devin.ai/cli/current/manifest.json,version_json_path=.version]" ;; amp | ampcode) agent="amp"; name="Amp" ;; +commandcode | command-code | cmdc) agent="commandcode"; name="Command Code"; agent_package="npm:command-code" ;; *) - echo "Usage: omarchy-default-agent " + echo "Usage: omarchy-default-agent " exit 1 ;; esac diff --git a/bin/omarchy-remove-preinstalls b/bin/omarchy-remove-preinstalls index 8841c1ca398..c12c0fed036 100755 --- a/bin/omarchy-remove-preinstalls +++ b/bin/omarchy-remove-preinstalls @@ -102,6 +102,12 @@ if gum confirm --affirmative "$(omarchy-i18n common.yes)" --negative "$(omarchy- rm -f ~/.local/bin/amp fi + # Preserve a user-managed Command Code launcher at the same path. + if [[ -f ~/.local/bin/commandcode && ! -L ~/.local/bin/commandcode ]] && + grep -Eq '^mise use -g .*"npm:command-code"' ~/.local/bin/commandcode; then + rm -f ~/.local/bin/commandcode + fi + omarchy-pkg-drop \ aether \ cliamp \ diff --git a/default/fonts/omarchy/README.md b/default/fonts/omarchy/README.md index e5a6758bb51..8a607cf70c6 100644 --- a/default/fonts/omarchy/README.md +++ b/default/fonts/omarchy/README.md @@ -27,6 +27,7 @@ The private-use glyphs in `omarchy.ttf` are: - `U+E916` — Kilo, from - `U+E917` — Devin, from the official mark published by Cognition at - `U+E918` — Amp, from , its three paths joined into one +- `U+E919` — Command Code, from , its traced paths joined into one The agent marks are monochrome so the menu can render them using the active theme's foreground and selection colors. diff --git a/default/fonts/omarchy/omarchy.ttf b/default/fonts/omarchy/omarchy.ttf index 71fc91429f9c1a505f437cf2132f241db9fd8fe8..0d664cd6f11a667c3701dae9adde21060a60bad1 100644 GIT binary patch delta 1025 zcmYjQ-)q}e6u!T7udXCZwj*81wQRRmoWzc4T&r@D(sbBs$LJWNY=1yZk~I%?)2-c_ zwHtmdjE#W{ed$A=O2zrOeRK0@_v!0!$E+n4Yu(gM5-`0&O4U?cwawVVj|>IsqBHa{-9_)9 zAJH#pAN_&;#u75)w%?9bto*m(3jY6ysVkbF1;U}W0uqUGEfSnX^Juj+8&9=isb93F& z(sH-Ayj!-cC!OTp>cK!^-j{#>K4Q%oog% zdAw@dI>UqtVVOYycO&y+1=g2Byjdj)$$VByGWm=FaNVV`uUXY)NjG9LmJPoJ5 zVA=Invzdj;N|lW;EC-flHf`JI;H_342Ci!`LckCvOisClbS_ISf$73bTAI|SDy6cl zD4#DepS9(rm8rz=BFBuq^uP`tup2$c|Own6s)yKe~I){OQ z!7KwPz`@9(1mpwd8B{WIODa+rIHH04A3%OfPJVJ?W?|7~1_renAm1f7v7&&{mdOCf zKLF$_gN!MJ5);%~zODDsnRd#Tbo^>pd"$mock_bin/mise" <<'SH' +#!/bin/bash +printf '%s\0' "$@" >>"$OMARCHY_TEST_MISE_LOG" +SH +cat >"$mock_bin/noop" <<'SH' +#!/bin/bash +exit 0 +SH +chmod +x "$mock_bin/mise" "$mock_bin/noop" +for command in gum hyprctl omarchy-webapp-remove-all omarchy-tui-remove-all omarchy-pkg-drop; do + ln -s noop "$mock_bin/$command" +done + +apply_commandcode() { + case $1 in + migration) source "$ROOT/migrations/1792034306.sh" >/dev/null ;; + removal) omarchy-remove-preinstalls >/dev/null ;; + esac +} + +apply_commandcode migration +[[ -x $commandcode && ! -L $commandcode ]] || fail "Command Code migration creates a missing wrapper" +[[ ! -s $OMARCHY_TEST_MISE_LOG ]] || fail "Command Code migration installs the package before first launch" +cp "$commandcode" "$test_tmp/wrapper" +apply_commandcode migration +cmp -s "$commandcode" "$test_tmp/wrapper" || fail "Command Code migration is idempotent" +"$commandcode" --prompt 'Review this project' +mapfile -d '' -t args <"$OMARCHY_TEST_MISE_LOG" +[[ ${args[*]} == "use -g --quiet $package x $package -- commandcode --prompt Review this project" ]] || + fail "Command Code migration wrapper installs the official package and forwards arguments" +pass "Command Code migration creates a lazy, repeatable wrapper" + +rm -f "$commandcode" +mkdir -p "$HOME/.local/state/omarchy" +touch "$HOME/.local/state/omarchy/preinstalls-removed" +apply_commandcode migration +[[ ! -e $commandcode ]] || fail "Command Code migration respects removed preinstalls with a missing command" +rm "$HOME/.local/state/omarchy/preinstalls-removed" +pass "Command Code migration respects the preinstall opt-out" + +for operation in migration removal; do + for kind in script binary symlink; do + rm -f "$commandcode" + case $kind in + script) printf '#!/bin/bash\necho user-managed commandcode\n' >"$commandcode" ;; + binary) cp /usr/bin/true "$commandcode" ;; + symlink) ln -s /usr/bin/true "$commandcode" ;; + esac + [[ $kind == "symlink" ]] || chmod +x "$commandcode" + rm -f "$test_tmp/saved-$kind" + cp -P "$commandcode" "$test_tmp/saved-$kind" + for attempt in 1 2; do + apply_commandcode "$operation" + cmp -s "$commandcode" "$test_tmp/saved-$kind" || fail "Command Code $operation preserves a user $kind" + if [[ $kind == "symlink" ]]; then + [[ -L $commandcode && $(readlink "$commandcode") == "/usr/bin/true" ]] || + fail "Command Code $operation preserves the user's symlink" + else + [[ ! -L $commandcode ]] || fail "Command Code $operation keeps the user's $kind as a regular file" + fi + done + pass "Command Code $operation preserves a user $kind on repeated runs" + done +done + +omarchy-mise-install "$package" commandcode +apply_commandcode removal +[[ ! -e $commandcode ]] || fail "Remove Preinstalls removes Omarchy's Command Code wrapper" +pass "Remove Preinstalls removes Omarchy's Command Code wrapper" diff --git a/test/shell.d/default-agent-test.sh b/test/shell.d/default-agent-test.sh index c58f9c160f3..2e9e3a765f2 100644 --- a/test/shell.d/default-agent-test.sh +++ b/test/shell.d/default-agent-test.sh @@ -150,6 +150,7 @@ zeroclaw_package="github:zeroclaw-labs/zeroclaw" kilo_package="npm:@kilocode/cli" devin_package="http:devin[platforms_linux_x64_url=https://static.devin.ai/cli/{{ version }}/devin-{{ version }}-x86_64-unknown-linux.tar.gz,platforms_linux_arm64_url=https://static.devin.ai/cli/{{ version }}/devin-{{ version }}-aarch64-unknown-linux.tar.gz,bin_path=bin,version_list_url=https://static.devin.ai/cli/current/manifest.json,version_json_path=.version]" amp_package="amp" +commandcode_package="npm:command-code" cursor_agent_package="cursor-agent" muse_package="http:muse[url=https://api.meta.ai/muse-launcher.sh,bin=muse,version_list_url=https://api.meta.ai/muse-code/channels/muse-stable,version_json_path=.version]" @@ -182,6 +183,7 @@ assert_lazy_stub "$zeroclaw_package" zeroclaw assert_lazy_stub "$kilo_package" kilo assert_lazy_stub "$devin_package" devin assert_lazy_stub "$amp_package" amp +assert_lazy_stub "$commandcode_package" commandcode assert_lazy_stub "$cursor_agent_package" cursor-agent assert_lazy_stub "$muse_package" muse pass "custom agent lazy stubs preserve their mise packages" @@ -205,6 +207,7 @@ grep -Fx "$freebuff_package freebuff" "$stub_log" >/dev/null || fail "user setup grep -Fx "$zeroclaw_package zeroclaw" "$stub_log" >/dev/null || fail "user setup creates the ZeroClaw lazy stub" grep -Fx "$kilo_package kilo" "$stub_log" >/dev/null || fail "user setup creates the Kilo lazy stub" grep -Fx "$devin_package devin" "$stub_log" >/dev/null || fail "user setup creates the Devin CLI lazy stub" +grep -Fx "$commandcode_package commandcode" "$stub_log" >/dev/null || fail "user setup creates the Command Code lazy stub" grep -Fx "$amp_package" "$stub_log" >/dev/null || fail "user setup creates the Amp lazy stub" pass "user setup creates the custom agent lazy stubs" @@ -223,6 +226,7 @@ grep -Fx "$freebuff_package freebuff" "$stub_log" >/dev/null || fail "user setup grep -Fx "$zeroclaw_package zeroclaw" "$stub_log" >/dev/null || fail "user setup recreates the ZeroClaw lazy stub" grep -Fx "$kilo_package kilo" "$stub_log" >/dev/null || fail "user setup recreates the Kilo lazy stub" grep -Fx "$devin_package devin" "$stub_log" >/dev/null || fail "user setup recreates the Devin CLI lazy stub" +grep -Fx "$commandcode_package commandcode" "$stub_log" >/dev/null || fail "user setup recreates the Command Code lazy stub" grep -Fx "$amp_package" "$stub_log" >/dev/null || fail "user setup recreates the Amp lazy stub" pass "user setup recreates wrappers even when the commands are present" @@ -284,6 +288,10 @@ grep -Fx "$kilo_package kilo" "$stub_log" >/dev/null || fail "Kilo migration cre OMARCHY_TEST_MISSING_COMMAND=devin source "$ROOT/migrations/1792034198.sh" >/dev/null grep -Fx "$devin_package devin" "$stub_log" >/dev/null || fail "Devin CLI migration creates its lazy stub" +: >"$stub_log" +OMARCHY_TEST_MISSING_COMMAND=commandcode source "$ROOT/migrations/1792034306.sh" >/dev/null +grep -Fx "$commandcode_package commandcode" "$stub_log" >/dev/null || fail "Command Code migration creates its lazy stub" + : >"$stub_log" OMARCHY_TEST_MISSING_COMMAND=amp source "$ROOT/migrations/1792034243.sh" >/dev/null grep -Fx "$amp_package" "$stub_log" >/dev/null || fail "Amp migration creates its lazy stub" @@ -402,6 +410,7 @@ OMARCHY_TEST_MISSING_COMMAND=zeroclaw source "$ROOT/migrations/1792034058.sh" >/ OMARCHY_TEST_MISSING_COMMAND=kilo source "$ROOT/migrations/1792034151.sh" >/dev/null OMARCHY_TEST_MISSING_COMMAND=devin source "$ROOT/migrations/1792034198.sh" >/dev/null OMARCHY_TEST_MISSING_COMMAND=amp source "$ROOT/migrations/1792034243.sh" >/dev/null +OMARCHY_TEST_MISSING_COMMAND=commandcode source "$ROOT/migrations/1792034306.sh" >/dev/null [[ ! -s $stub_log ]] || fail "agent migrations respect the preinstall opt-out" [[ ! -e $test_home/.local/bin/omp ]] || fail "agent migration removes the obsolete Oh My Pi wrapper after opt-out" @@ -546,9 +555,10 @@ pass "agent migrations install working wrappers without overriding the preinstal "$ROOT/bin/omarchy-mise-install" "$kilo_package" kilo "$ROOT/bin/omarchy-mise-install" "$devin_package" devin "$ROOT/bin/omarchy-mise-install" "$amp_package" amp +"$ROOT/bin/omarchy-mise-install" "$commandcode_package" commandcode touch "$test_home/.local/bin/agy" "$test_home/.local/bin/ori" omarchy-remove-preinstalls >/dev/null -for command in agy omp ori grok crush cursor-agent muse kimi qwen dsh mcode goose kiro-cli vibe freebuff zeroclaw kilo devin amp; do +for command in agy omp ori grok crush cursor-agent muse kimi qwen dsh mcode goose kiro-cli vibe freebuff zeroclaw kilo devin amp commandcode; do [[ ! -e $test_home/.local/bin/$command ]] || fail "Remove Preinstalls deletes the $command lazy stub" done pass "Remove Preinstalls deletes every optional agent lazy stub" @@ -654,6 +664,9 @@ declare -A expected_agents=( [devin-cli]="devin" [amp]="amp" [ampcode]="amp" + [commandcode]="commandcode" + [command-code]="commandcode" + [cmdc]="commandcode" [cursor]="cursor-agent" [cursor-agent]="cursor-agent" [muse]="muse" @@ -684,6 +697,7 @@ declare -A expected_packages=( [kilo]="$kilo_package" [devin]="$devin_package" [amp]="$amp_package" + [commandcode]="$commandcode_package" [cursor-agent]="$cursor_agent_package" [muse]="$muse_package" ) @@ -972,6 +986,7 @@ assert_launch zeroclaw zerocode assert_launch kilo kilo --auto "--prompt=Review this project" assert_launch devin devin --permission-mode dangerous --respect-workspace-trust false -- "Review this project" assert_launch amp bash -c 'printf "%s\n" "$1" | exec amp --dangerously-allow-all' amp "Review this project" +assert_launch commandcode commandcode --yolo -- "Review this project" [[ $(<"$clipboard_log") == "Review this project" ]] || fail "Freebuff launcher puts the prompt on the clipboard" mapfile -d '' -t notification_args <"$notification_history" [[ ${notification_args[*]} == *"Freebuff takes no starting prompt"* ]] || fail "Freebuff launcher says where the prompt went" @@ -1017,6 +1032,7 @@ assert_bypass zeroclaw zerocode assert_bypass kilo kilo --auto assert_bypass devin devin --permission-mode dangerous --respect-workspace-trust false assert_bypass amp amp --dangerously-allow-all +assert_bypass commandcode commandcode --yolo [[ ! -s $clipboard_log && ! -s $notification_history ]] || fail "Freebuff launcher leaves the clipboard alone without a prompt" pass "agent launcher skips permission prompts for every supported agent" diff --git a/test/shell.d/menu-test.sh b/test/shell.d/menu-test.sh index d1f7ce5b74e..08bd8b6fa68 100644 --- a/test/shell.d/menu-test.sh +++ b/test/shell.d/menu-test.sh @@ -264,6 +264,7 @@ const expectedAgents = { kilo: { icon: '\ue916', iconFont: 'omarchy', label: 'Kilo' }, devin: { icon: '\ue917', iconFont: 'omarchy', label: 'Devin CLI' }, amp: { icon: '\ue918', iconFont: 'omarchy', label: 'Amp' }, + commandcode: { icon: '\ue919', iconFont: 'omarchy', label: 'Command Code' }, hermes: { icon: '\ue90a', iconFont: 'omarchy', label: 'Hermes' }, openclaw: { icon: '\ue90c', iconFont: 'omarchy', label: 'OpenClaw' }, copilot: { icon: '', label: 'Copilot' }, @@ -289,7 +290,7 @@ assertDeepEqual( defaultItems .filter(item => item.parent === 'setup.default.agent') .map(item => item.label), - ['Amp', 'Antigravity', 'Claude', 'Codex', 'Copilot', 'Crush', 'Cursor CLI', 'DeepSeek Harness', 'Devin CLI', 'Freebuff', 'Goose', 'Grok', 'Hermes', 'Kilo', 'Kimi', 'Kiro CLI', 'MiniMax Code', 'Mistral Vibe', 'Muse Code', 'omp', 'OpenClaw', 'OpenCode', 'Ori', 'Pi', 'Qwen Code', 'ZeroClaw'], + ['Amp', 'Antigravity', 'Claude', 'Codex', 'Command Code', 'Copilot', 'Crush', 'Cursor CLI', 'DeepSeek Harness', 'Devin CLI', 'Freebuff', 'Goose', 'Grok', 'Hermes', 'Kilo', 'Kimi', 'Kiro CLI', 'MiniMax Code', 'Mistral Vibe', 'Muse Code', 'omp', 'OpenClaw', 'OpenCode', 'Ori', 'Pi', 'Qwen Code', 'ZeroClaw'], 'menu sorts coding agents alphabetically' ) const expectedDefaults = { @@ -886,5 +887,5 @@ assert( JS font_charset=$(fc-query --format='%{charset}' "$ROOT/default/fonts/omarchy/omarchy.ttf") -[[ $font_charset == *"e900-e918"* ]] || fail "Omarchy icon font includes every custom menu glyph" +[[ $font_charset == *"e900-e919"* ]] || fail "Omarchy icon font includes every custom menu glyph" pass "Omarchy icon font includes the official agent marks" diff --git a/test/shell.d/mise-provision-test.sh b/test/shell.d/mise-provision-test.sh index d3099f37f65..260357be9fd 100644 --- a/test/shell.d/mise-provision-test.sh +++ b/test/shell.d/mise-provision-test.sh @@ -17,7 +17,7 @@ SH chmod +x "$mock_bin/mise" source "$ROOT/install/user/mise.sh" -for agent in kimi cursor-agent muse qwen dsh mcode goose kiro-cli vibe freebuff zeroclaw kilo devin amp; do +for agent in kimi cursor-agent muse qwen dsh mcode goose kiro-cli vibe freebuff zeroclaw kilo devin amp commandcode; do [[ -x $HOME/.local/bin/$agent && ! -L $HOME/.local/bin/$agent ]] || fail "provisioning creates the missing $agent wrapper" cp "$HOME/.local/bin/$agent" "$test_tmp/$agent-wrapper" @@ -27,7 +27,7 @@ mapfile -d '' -t args <"$OMARCHY_TEST_MISE_LOG" fail "provisioning creates lazy wrappers without installing packages" pass "provisioning creates lazy Kimi, Cursor, and Muse wrappers" -for agent in kimi cursor-agent muse qwen dsh mcode goose kiro-cli vibe freebuff zeroclaw kilo devin amp; do +for agent in kimi cursor-agent muse qwen dsh mcode goose kiro-cli vibe freebuff zeroclaw kilo devin amp commandcode; do launcher="$HOME/.local/bin/$agent" for kind in script binary symlink; do rm -f "$launcher" From c1b00c4cb528d8439a2b50df646623806931f0ee Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:19:54 -0400 Subject: [PATCH 108/546] Add Phantombot as a selectable default coding agent Phantombot (phantombot) joins the default agent picker. mise installs its GitHub release binary for x86_64 or aarch64 in place of the PR's wrapper around the upstream install script, with a lazy stub and migration like the other agents. Permissions belong to the harness it runs underneath, so there is no bypass flag, and --prompt= opens its TUI with the prompt as the first turn. From #12251. Co-authored-by: Lena Park Co-authored-by: Andrew Hodges Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent | 7 +++ bin/omarchy-default-agent | 5 +- bin/omarchy-remove-preinstalls | 6 ++ default/fonts/omarchy/README.md | 1 + default/fonts/omarchy/omarchy.ttf | Bin 13996 -> 14256 bytes default/i18n/en/menu.json | 3 + default/omarchy/omarchy-menu.jsonc | 1 + install/user/mise.sh | 1 + manual/17-ai.md | 1 + migrations/1792034353.sh | 5 ++ test/shell.d/default-agent-test.sh | 16 ++++- test/shell.d/menu-test.sh | 5 +- test/shell.d/mise-provision-test.sh | 4 +- test/shell.d/phantombot-agent-test.sh | 85 ++++++++++++++++++++++++++ 14 files changed, 133 insertions(+), 7 deletions(-) create mode 100644 migrations/1792034353.sh create mode 100644 test/shell.d/phantombot-agent-test.sh diff --git a/bin/omarchy-agent b/bin/omarchy-agent index fb0150602dc..5c01d863b73 100755 --- a/bin/omarchy-agent +++ b/bin/omarchy-agent @@ -218,6 +218,13 @@ commandcode) command=(commandcode --yolo) [[ -n ${prompt:-} ]] && command+=(-- "$prompt") ;; +phantombot) + # Permissions belong to the harness Phantombot runs underneath, so it has + # no bypass flag. --prompt= opens the TUI with the prompt as its first turn; + # the bound form is the one that keeps a leading dash literal. + command=(phantombot) + [[ -n ${prompt:-} ]] && command+=("--prompt=$prompt") + ;; *) echo "Unsupported default agent: $agent" >&2 exit 1 diff --git a/bin/omarchy-default-agent b/bin/omarchy-default-agent index eebaee131e7..b863f814fce 100755 --- a/bin/omarchy-default-agent +++ b/bin/omarchy-default-agent @@ -1,7 +1,7 @@ #!/bin/bash # omarchy:summary=Set and launch the default coding agent -# omarchy:args=[pi|omp|opencode|ori|claude|codex|grok|openclaw|agy|hermes|copilot|crush|cursor-agent|muse|kimi|qwen|dsh|mcode|goose|kiro-cli|vibe|freebuff|zeroclaw|kilo|devin|amp|commandcode] +# omarchy:args=[pi|omp|opencode|ori|claude|codex|grok|openclaw|agy|hermes|copilot|crush|cursor-agent|muse|kimi|qwen|dsh|mcode|goose|kiro-cli|vibe|freebuff|zeroclaw|kilo|devin|amp|commandcode|phantombot] # omarchy:examples=omarchy default agent | omarchy default agent codex | omarchy default agent claude installing=false @@ -69,8 +69,9 @@ devin | devin-cli) ;; amp | ampcode) agent="amp"; name="Amp" ;; commandcode | command-code | cmdc) agent="commandcode"; name="Command Code"; agent_package="npm:command-code" ;; +phantombot) agent="phantombot"; name="Phantombot"; agent_package="github:phantomyard/phantombot" ;; *) - echo "Usage: omarchy-default-agent " + echo "Usage: omarchy-default-agent " exit 1 ;; esac diff --git a/bin/omarchy-remove-preinstalls b/bin/omarchy-remove-preinstalls index c12c0fed036..808bed5e3b0 100755 --- a/bin/omarchy-remove-preinstalls +++ b/bin/omarchy-remove-preinstalls @@ -108,6 +108,12 @@ if gum confirm --affirmative "$(omarchy-i18n common.yes)" --negative "$(omarchy- rm -f ~/.local/bin/commandcode fi + # Preserve a user-managed Phantombot launcher at the same path. + if [[ -f ~/.local/bin/phantombot && ! -L ~/.local/bin/phantombot ]] && + grep -Eq '^mise use -g .*"github:phantomyard/phantombot"' ~/.local/bin/phantombot; then + rm -f ~/.local/bin/phantombot + fi + omarchy-pkg-drop \ aether \ cliamp \ diff --git a/default/fonts/omarchy/README.md b/default/fonts/omarchy/README.md index 8a607cf70c6..a0270961f27 100644 --- a/default/fonts/omarchy/README.md +++ b/default/fonts/omarchy/README.md @@ -28,6 +28,7 @@ The private-use glyphs in `omarchy.ttf` are: - `U+E917` — Devin, from the official mark published by Cognition at - `U+E918` — Amp, from , its three paths joined into one - `U+E919` — Command Code, from , its traced paths joined into one +- `U+E91A` — Phantombot, redrawn as a monochrome ghost silhouette from the project logo at (`www/assets/img/logo.png`), since upstream publishes no monochrome SVG The agent marks are monochrome so the menu can render them using the active theme's foreground and selection colors. diff --git a/default/fonts/omarchy/omarchy.ttf b/default/fonts/omarchy/omarchy.ttf index 0d664cd6f11a667c3701dae9adde21060a60bad1..f9ee3828bcb24e3f3e640a5bd9d1ba4587ca11ab 100644 GIT binary patch delta 494 zcmZ3JyCHvqVSR;Q5Ca2a2Ll6>S9(rmnq_A2DFz1SD+~;(8!}Q8Q}otZ^)WDLwlFX- zm}LM3I2c(}fc!l`zDh=JNkt9=r!kQK0Lbsj$xlwKzSz2&fk7(=$oI%itSDf#XEFft zw*dJHd5O8H#@%;nfqbC(+#3t>i%WPEUqAB(nhkUSqw2;zvW$#UlW#G;<&*+y5Cw{{ zZNA9F!npYg^Fft*W`^ku4$PU%2N*ONEEzn3D#XQrP*6=>O<9SZjU7mZoELE|%>H8PZuVPsTNHa1dLVq#TP(l=02Vr5iPGBQz8Vq}z&F*K@o zboBOi{P$N%N>GrIQBY7yNS9(rm+S$Xs+ZY&_uP`vE9LPvbOwn6s)yKe~F@=GF z!7KwPz`@9(4CJ2y@>Mc&ODZxLIOYKPAAtOhoc!d(>Wf8}85lGxfPA;y#EJq&J0=4l z{{WD$ke8U7YTSLN7RZ0Vz`(t+AiubThpXypJp+R&&;mw@jeBGn86_v*VtmUf3Dh78 z6l2?bk%@(I^A+ZUD%^}fF-BwKtHNJXCQs1xXH=T}K(k|Vj+Pq-NCJqJHt*7U&kYn` LW#%&ZgV`nkt)xI` diff --git a/default/i18n/en/menu.json b/default/i18n/en/menu.json index 7e3952b61b9..d6802cb90a5 100644 --- a/default/i18n/en/menu.json +++ b/default/i18n/en/menu.json @@ -551,6 +551,9 @@ "setup.default.agent.ori": { "label": "Ori" }, + "setup.default.agent.phantombot": { + "label": "Phantombot" + }, "setup.default.agent.pi": { "label": "Pi" }, diff --git a/default/omarchy/omarchy-menu.jsonc b/default/omarchy/omarchy-menu.jsonc index b5a37a4ec35..d2cecee6527 100644 --- a/default/omarchy/omarchy-menu.jsonc +++ b/default/omarchy/omarchy-menu.jsonc @@ -214,6 +214,7 @@ "setup.default.agent.openclaw": {"icon":"","iconFont":"omarchy","label":"OpenClaw","checked":"[[ \"$(omarchy-default-agent)\" == \"openclaw\" ]]","action":"omarchy-default-agent openclaw"}, "setup.default.agent.opencode": {"icon":"","iconFont":"omarchy","label":"OpenCode","checked":"[[ \"$(omarchy-default-agent)\" == \"opencode\" ]]","action":"omarchy-default-agent opencode"}, "setup.default.agent.ori": {"icon":"","iconFont":"omarchy","label":"Ori","checked":"[[ \"$(omarchy-default-agent)\" == \"ori\" ]]","action":"omarchy-default-agent ori"}, + "setup.default.agent.phantombot": {"icon":"","iconFont":"omarchy","label":"Phantombot","checked":"[[ \"$(omarchy-default-agent)\" == \"phantombot\" ]]","action":"omarchy-default-agent phantombot"}, "setup.default.agent.pi": {"icon":"","iconFont":"omarchy","label":"Pi","checked":"[[ \"$(omarchy-default-agent)\" == \"pi\" ]]","action":"omarchy-default-agent pi"}, "setup.default.agent.qwen": {"icon":"","iconFont":"omarchy","label":"Qwen Code","checked":"[[ \"$(omarchy-default-agent)\" == \"qwen\" ]]","action":"omarchy-default-agent qwen"}, "setup.default.agent.zeroclaw": {"icon":"󰬡","label":"ZeroClaw","checked":"[[ \"$(omarchy-default-agent)\" == \"zeroclaw\" ]]","action":"omarchy-default-agent zeroclaw"}, diff --git a/install/user/mise.sh b/install/user/mise.sh index 7f13d26c911..992405296a2 100644 --- a/install/user/mise.sh +++ b/install/user/mise.sh @@ -26,6 +26,7 @@ omarchy-mise-install github:zeroclaw-labs/zeroclaw zeroclaw omarchy-mise-install npm:@kilocode/cli kilo omarchy-mise-install "http:devin[platforms_linux_x64_url=https://static.devin.ai/cli/{{ version }}/devin-{{ version }}-x86_64-unknown-linux.tar.gz,platforms_linux_arm64_url=https://static.devin.ai/cli/{{ version }}/devin-{{ version }}-aarch64-unknown-linux.tar.gz,bin_path=bin,version_list_url=https://static.devin.ai/cli/current/manifest.json,version_json_path=.version]" devin omarchy-mise-install npm:command-code commandcode +omarchy-mise-install github:phantomyard/phantombot phantombot omarchy-mise-install amp omarchy-mise-install cursor-agent omarchy-mise-install npm:@kitlangton/ghui ghui diff --git a/manual/17-ai.md b/manual/17-ai.md index ff162582f28..995407bbbb1 100644 --- a/manual/17-ai.md +++ b/manual/17-ai.md @@ -21,6 +21,7 @@ Omarchy treats AI coding agents as first-class citizens, but it doesn't pick a f | `kimi` | [Kimi Code](https://www.kimi.com/code/docs/en/) | | `kiro-cli` | [Kiro CLI](https://kiro.dev/docs/cli/) | | `mcode` | [MiniMax Code](https://github.com/MiniMax-AI/minimax-code), MiniMax's coding agent | +| `phantombot` | [Phantombot](https://phantombot.bot), a persona-first assistant over any harness | | `qwen` | [Qwen Code](https://github.com/QwenLM/qwen-code), Alibaba's coding agent | | `vibe` | [Mistral Vibe](https://docs.mistral.ai/vibe/code/overview), Mistral AI's coding agent | | `zeroclaw` | [ZeroClaw](https://github.com/zeroclaw-labs/zeroclaw), opened in its `zerocode` terminal UI | diff --git a/migrations/1792034353.sh b/migrations/1792034353.sh new file mode 100644 index 00000000000..79bf84667e1 --- /dev/null +++ b/migrations/1792034353.sh @@ -0,0 +1,5 @@ +echo "Install Phantombot coding agent mise wrapper" + +if [[ ! -f $HOME/.local/state/omarchy/preinstalls-removed ]] && omarchy-cmd-missing phantombot; then + omarchy-mise-install github:phantomyard/phantombot phantombot +fi diff --git a/test/shell.d/default-agent-test.sh b/test/shell.d/default-agent-test.sh index 2e9e3a765f2..145389b1b70 100644 --- a/test/shell.d/default-agent-test.sh +++ b/test/shell.d/default-agent-test.sh @@ -151,6 +151,7 @@ kilo_package="npm:@kilocode/cli" devin_package="http:devin[platforms_linux_x64_url=https://static.devin.ai/cli/{{ version }}/devin-{{ version }}-x86_64-unknown-linux.tar.gz,platforms_linux_arm64_url=https://static.devin.ai/cli/{{ version }}/devin-{{ version }}-aarch64-unknown-linux.tar.gz,bin_path=bin,version_list_url=https://static.devin.ai/cli/current/manifest.json,version_json_path=.version]" amp_package="amp" commandcode_package="npm:command-code" +phantombot_package="github:phantomyard/phantombot" cursor_agent_package="cursor-agent" muse_package="http:muse[url=https://api.meta.ai/muse-launcher.sh,bin=muse,version_list_url=https://api.meta.ai/muse-code/channels/muse-stable,version_json_path=.version]" @@ -184,6 +185,7 @@ assert_lazy_stub "$kilo_package" kilo assert_lazy_stub "$devin_package" devin assert_lazy_stub "$amp_package" amp assert_lazy_stub "$commandcode_package" commandcode +assert_lazy_stub "$phantombot_package" phantombot assert_lazy_stub "$cursor_agent_package" cursor-agent assert_lazy_stub "$muse_package" muse pass "custom agent lazy stubs preserve their mise packages" @@ -208,6 +210,7 @@ grep -Fx "$zeroclaw_package zeroclaw" "$stub_log" >/dev/null || fail "user setup grep -Fx "$kilo_package kilo" "$stub_log" >/dev/null || fail "user setup creates the Kilo lazy stub" grep -Fx "$devin_package devin" "$stub_log" >/dev/null || fail "user setup creates the Devin CLI lazy stub" grep -Fx "$commandcode_package commandcode" "$stub_log" >/dev/null || fail "user setup creates the Command Code lazy stub" +grep -Fx "$phantombot_package phantombot" "$stub_log" >/dev/null || fail "user setup creates the Phantombot lazy stub" grep -Fx "$amp_package" "$stub_log" >/dev/null || fail "user setup creates the Amp lazy stub" pass "user setup creates the custom agent lazy stubs" @@ -227,6 +230,7 @@ grep -Fx "$zeroclaw_package zeroclaw" "$stub_log" >/dev/null || fail "user setup grep -Fx "$kilo_package kilo" "$stub_log" >/dev/null || fail "user setup recreates the Kilo lazy stub" grep -Fx "$devin_package devin" "$stub_log" >/dev/null || fail "user setup recreates the Devin CLI lazy stub" grep -Fx "$commandcode_package commandcode" "$stub_log" >/dev/null || fail "user setup recreates the Command Code lazy stub" +grep -Fx "$phantombot_package phantombot" "$stub_log" >/dev/null || fail "user setup recreates the Phantombot lazy stub" grep -Fx "$amp_package" "$stub_log" >/dev/null || fail "user setup recreates the Amp lazy stub" pass "user setup recreates wrappers even when the commands are present" @@ -292,6 +296,10 @@ grep -Fx "$devin_package devin" "$stub_log" >/dev/null || fail "Devin CLI migrat OMARCHY_TEST_MISSING_COMMAND=commandcode source "$ROOT/migrations/1792034306.sh" >/dev/null grep -Fx "$commandcode_package commandcode" "$stub_log" >/dev/null || fail "Command Code migration creates its lazy stub" +: >"$stub_log" +OMARCHY_TEST_MISSING_COMMAND=phantombot source "$ROOT/migrations/1792034353.sh" >/dev/null +grep -Fx "$phantombot_package phantombot" "$stub_log" >/dev/null || fail "Phantombot migration creates its lazy stub" + : >"$stub_log" OMARCHY_TEST_MISSING_COMMAND=amp source "$ROOT/migrations/1792034243.sh" >/dev/null grep -Fx "$amp_package" "$stub_log" >/dev/null || fail "Amp migration creates its lazy stub" @@ -411,6 +419,7 @@ OMARCHY_TEST_MISSING_COMMAND=kilo source "$ROOT/migrations/1792034151.sh" >/dev/ OMARCHY_TEST_MISSING_COMMAND=devin source "$ROOT/migrations/1792034198.sh" >/dev/null OMARCHY_TEST_MISSING_COMMAND=amp source "$ROOT/migrations/1792034243.sh" >/dev/null OMARCHY_TEST_MISSING_COMMAND=commandcode source "$ROOT/migrations/1792034306.sh" >/dev/null +OMARCHY_TEST_MISSING_COMMAND=phantombot source "$ROOT/migrations/1792034353.sh" >/dev/null [[ ! -s $stub_log ]] || fail "agent migrations respect the preinstall opt-out" [[ ! -e $test_home/.local/bin/omp ]] || fail "agent migration removes the obsolete Oh My Pi wrapper after opt-out" @@ -556,9 +565,10 @@ pass "agent migrations install working wrappers without overriding the preinstal "$ROOT/bin/omarchy-mise-install" "$devin_package" devin "$ROOT/bin/omarchy-mise-install" "$amp_package" amp "$ROOT/bin/omarchy-mise-install" "$commandcode_package" commandcode +"$ROOT/bin/omarchy-mise-install" "$phantombot_package" phantombot touch "$test_home/.local/bin/agy" "$test_home/.local/bin/ori" omarchy-remove-preinstalls >/dev/null -for command in agy omp ori grok crush cursor-agent muse kimi qwen dsh mcode goose kiro-cli vibe freebuff zeroclaw kilo devin amp commandcode; do +for command in agy omp ori grok crush cursor-agent muse kimi qwen dsh mcode goose kiro-cli vibe freebuff zeroclaw kilo devin amp commandcode phantombot; do [[ ! -e $test_home/.local/bin/$command ]] || fail "Remove Preinstalls deletes the $command lazy stub" done pass "Remove Preinstalls deletes every optional agent lazy stub" @@ -667,6 +677,7 @@ declare -A expected_agents=( [commandcode]="commandcode" [command-code]="commandcode" [cmdc]="commandcode" + [phantombot]="phantombot" [cursor]="cursor-agent" [cursor-agent]="cursor-agent" [muse]="muse" @@ -698,6 +709,7 @@ declare -A expected_packages=( [devin]="$devin_package" [amp]="$amp_package" [commandcode]="$commandcode_package" + [phantombot]="$phantombot_package" [cursor-agent]="$cursor_agent_package" [muse]="$muse_package" ) @@ -987,6 +999,7 @@ assert_launch kilo kilo --auto "--prompt=Review this project" assert_launch devin devin --permission-mode dangerous --respect-workspace-trust false -- "Review this project" assert_launch amp bash -c 'printf "%s\n" "$1" | exec amp --dangerously-allow-all' amp "Review this project" assert_launch commandcode commandcode --yolo -- "Review this project" +assert_launch phantombot phantombot "--prompt=Review this project" [[ $(<"$clipboard_log") == "Review this project" ]] || fail "Freebuff launcher puts the prompt on the clipboard" mapfile -d '' -t notification_args <"$notification_history" [[ ${notification_args[*]} == *"Freebuff takes no starting prompt"* ]] || fail "Freebuff launcher says where the prompt went" @@ -1033,6 +1046,7 @@ assert_bypass kilo kilo --auto assert_bypass devin devin --permission-mode dangerous --respect-workspace-trust false assert_bypass amp amp --dangerously-allow-all assert_bypass commandcode commandcode --yolo +assert_bypass phantombot phantombot [[ ! -s $clipboard_log && ! -s $notification_history ]] || fail "Freebuff launcher leaves the clipboard alone without a prompt" pass "agent launcher skips permission prompts for every supported agent" diff --git a/test/shell.d/menu-test.sh b/test/shell.d/menu-test.sh index 08bd8b6fa68..c2a7549741a 100644 --- a/test/shell.d/menu-test.sh +++ b/test/shell.d/menu-test.sh @@ -265,6 +265,7 @@ const expectedAgents = { devin: { icon: '\ue917', iconFont: 'omarchy', label: 'Devin CLI' }, amp: { icon: '\ue918', iconFont: 'omarchy', label: 'Amp' }, commandcode: { icon: '\ue919', iconFont: 'omarchy', label: 'Command Code' }, + phantombot: { icon: '\ue91a', iconFont: 'omarchy', label: 'Phantombot' }, hermes: { icon: '\ue90a', iconFont: 'omarchy', label: 'Hermes' }, openclaw: { icon: '\ue90c', iconFont: 'omarchy', label: 'OpenClaw' }, copilot: { icon: '', label: 'Copilot' }, @@ -290,7 +291,7 @@ assertDeepEqual( defaultItems .filter(item => item.parent === 'setup.default.agent') .map(item => item.label), - ['Amp', 'Antigravity', 'Claude', 'Codex', 'Command Code', 'Copilot', 'Crush', 'Cursor CLI', 'DeepSeek Harness', 'Devin CLI', 'Freebuff', 'Goose', 'Grok', 'Hermes', 'Kilo', 'Kimi', 'Kiro CLI', 'MiniMax Code', 'Mistral Vibe', 'Muse Code', 'omp', 'OpenClaw', 'OpenCode', 'Ori', 'Pi', 'Qwen Code', 'ZeroClaw'], + ['Amp', 'Antigravity', 'Claude', 'Codex', 'Command Code', 'Copilot', 'Crush', 'Cursor CLI', 'DeepSeek Harness', 'Devin CLI', 'Freebuff', 'Goose', 'Grok', 'Hermes', 'Kilo', 'Kimi', 'Kiro CLI', 'MiniMax Code', 'Mistral Vibe', 'Muse Code', 'omp', 'OpenClaw', 'OpenCode', 'Ori', 'Phantombot', 'Pi', 'Qwen Code', 'ZeroClaw'], 'menu sorts coding agents alphabetically' ) const expectedDefaults = { @@ -887,5 +888,5 @@ assert( JS font_charset=$(fc-query --format='%{charset}' "$ROOT/default/fonts/omarchy/omarchy.ttf") -[[ $font_charset == *"e900-e919"* ]] || fail "Omarchy icon font includes every custom menu glyph" +[[ $font_charset == *"e900-e91a"* ]] || fail "Omarchy icon font includes every custom menu glyph" pass "Omarchy icon font includes the official agent marks" diff --git a/test/shell.d/mise-provision-test.sh b/test/shell.d/mise-provision-test.sh index 260357be9fd..96934ac6bfa 100644 --- a/test/shell.d/mise-provision-test.sh +++ b/test/shell.d/mise-provision-test.sh @@ -17,7 +17,7 @@ SH chmod +x "$mock_bin/mise" source "$ROOT/install/user/mise.sh" -for agent in kimi cursor-agent muse qwen dsh mcode goose kiro-cli vibe freebuff zeroclaw kilo devin amp commandcode; do +for agent in kimi cursor-agent muse qwen dsh mcode goose kiro-cli vibe freebuff zeroclaw kilo devin amp commandcode phantombot; do [[ -x $HOME/.local/bin/$agent && ! -L $HOME/.local/bin/$agent ]] || fail "provisioning creates the missing $agent wrapper" cp "$HOME/.local/bin/$agent" "$test_tmp/$agent-wrapper" @@ -27,7 +27,7 @@ mapfile -d '' -t args <"$OMARCHY_TEST_MISE_LOG" fail "provisioning creates lazy wrappers without installing packages" pass "provisioning creates lazy Kimi, Cursor, and Muse wrappers" -for agent in kimi cursor-agent muse qwen dsh mcode goose kiro-cli vibe freebuff zeroclaw kilo devin amp commandcode; do +for agent in kimi cursor-agent muse qwen dsh mcode goose kiro-cli vibe freebuff zeroclaw kilo devin amp commandcode phantombot; do launcher="$HOME/.local/bin/$agent" for kind in script binary symlink; do rm -f "$launcher" diff --git a/test/shell.d/phantombot-agent-test.sh b/test/shell.d/phantombot-agent-test.sh new file mode 100644 index 00000000000..d000197ce8b --- /dev/null +++ b/test/shell.d/phantombot-agent-test.sh @@ -0,0 +1,85 @@ +#!/bin/bash + +set -euo pipefail +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" + +test_tmp=$(mktemp -d) +trap 'rm -rf "$test_tmp"' EXIT +mock_bin="$test_tmp/bin" +export HOME="$test_tmp/home with spaces" +export PATH="$mock_bin:$HOME/.local/bin:$ROOT/bin:/usr/local/bin:/usr/bin" +export OMARCHY_PATH="$ROOT" +export OMARCHY_TEST_MISE_LOG="$test_tmp/mise-log" +mkdir -p "$mock_bin" "$HOME/.local/bin" +phantombot="$HOME/.local/bin/phantombot" +package="github:phantomyard/phantombot" + +cat >"$mock_bin/mise" <<'SH' +#!/bin/bash +printf '%s\0' "$@" >>"$OMARCHY_TEST_MISE_LOG" +SH +cat >"$mock_bin/noop" <<'SH' +#!/bin/bash +exit 0 +SH +chmod +x "$mock_bin/mise" "$mock_bin/noop" +for command in gum hyprctl omarchy-webapp-remove-all omarchy-tui-remove-all omarchy-pkg-drop; do + ln -s noop "$mock_bin/$command" +done + +apply_phantombot() { + case $1 in + migration) source "$ROOT/migrations/1792034353.sh" >/dev/null ;; + removal) omarchy-remove-preinstalls >/dev/null ;; + esac +} + +apply_phantombot migration +[[ -x $phantombot && ! -L $phantombot ]] || fail "Phantombot migration creates a missing wrapper" +[[ ! -s $OMARCHY_TEST_MISE_LOG ]] || fail "Phantombot migration installs the package before first launch" +cp "$phantombot" "$test_tmp/wrapper" +apply_phantombot migration +cmp -s "$phantombot" "$test_tmp/wrapper" || fail "Phantombot migration is idempotent" +"$phantombot" --prompt 'Review this project' +mapfile -d '' -t args <"$OMARCHY_TEST_MISE_LOG" +[[ ${args[*]} == "use -g --quiet $package x $package -- phantombot --prompt Review this project" ]] || + fail "Phantombot migration wrapper installs the official package and forwards arguments" +pass "Phantombot migration creates a lazy, repeatable wrapper" + +rm -f "$phantombot" +mkdir -p "$HOME/.local/state/omarchy" +touch "$HOME/.local/state/omarchy/preinstalls-removed" +apply_phantombot migration +[[ ! -e $phantombot ]] || fail "Phantombot migration respects removed preinstalls with a missing command" +rm "$HOME/.local/state/omarchy/preinstalls-removed" +pass "Phantombot migration respects the preinstall opt-out" + +for operation in migration removal; do + for kind in script binary symlink; do + rm -f "$phantombot" + case $kind in + script) printf '#!/bin/bash\necho user-managed phantombot\n' >"$phantombot" ;; + binary) cp /usr/bin/true "$phantombot" ;; + symlink) ln -s /usr/bin/true "$phantombot" ;; + esac + [[ $kind == "symlink" ]] || chmod +x "$phantombot" + rm -f "$test_tmp/saved-$kind" + cp -P "$phantombot" "$test_tmp/saved-$kind" + for attempt in 1 2; do + apply_phantombot "$operation" + cmp -s "$phantombot" "$test_tmp/saved-$kind" || fail "Phantombot $operation preserves a user $kind" + if [[ $kind == "symlink" ]]; then + [[ -L $phantombot && $(readlink "$phantombot") == "/usr/bin/true" ]] || + fail "Phantombot $operation preserves the user's symlink" + else + [[ ! -L $phantombot ]] || fail "Phantombot $operation keeps the user's $kind as a regular file" + fi + done + pass "Phantombot $operation preserves a user $kind on repeated runs" + done +done + +omarchy-mise-install "$package" phantombot +apply_phantombot removal +[[ ! -e $phantombot ]] || fail "Remove Preinstalls removes Omarchy's Phantombot wrapper" +pass "Remove Preinstalls removes Omarchy's Phantombot wrapper" From 5302a3559298895f90ad26f2c84d12743c3bc641 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:25:24 -0400 Subject: [PATCH 109/546] List the unavailable-source status in the collector contract Co-Authored-By: Claude Opus 5.5 --- docs/agent-usage-collectors.md | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/agent-usage-collectors.md b/docs/agent-usage-collectors.md index c849822cff8..1e721d100e2 100644 --- a/docs/agent-usage-collectors.md +++ b/docs/agent-usage-collectors.md @@ -45,6 +45,7 @@ The panel acts on a small vocabulary in `usageStatusText`, so the same situation | `"Sign-in expired"` | The credential was refused, or a refresh token lapsed; signing in again is needed | | `"Limits paused"` | An access token lapsed that the agent's own CLI renews the next time it runs; signing in again is not needed | | `" limits unavailable"` | A check failed with nothing kept to show | +| `" unavailable"` | The agent's stats source failed (an unreadable database, a billing API that answered nothing) with nothing kept to show | | `"Balance unavailable"` | A prepaid balance couldn't be read and none was kept | The detail goes in `authHelpText`: what failed and the one thing to do about it. It is `""` when healthy. A plan the service doesn't offer limits for (pay-as-you-go, no Go plan, a free tier) isn't trouble: the status stays `""` and local stats show alone. Any other status shows under the name in the urgent color, so collectors don't invent new ones. From 46a45c68bc3d15e1c6ddd03af7ae359179165794 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:28:00 -0400 Subject: [PATCH 110/546] Count Ollama cloud models reached through the local daemon A cloud-tagged model on pi's or omp's local `ollama` provider runs on Ollama's servers, so the local collector leaves it out; nothing counted it. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-ollama | 18 +++++++++++++++--- .../shell.d/agent-usage-ollama-scanner-test.sh | 17 +++++++++++++++++ 2 files changed, 32 insertions(+), 3 deletions(-) diff --git a/bin/omarchy-agent-usage-ollama b/bin/omarchy-agent-usage-ollama index 03eb4a735ba..c9a7e826172 100755 --- a/bin/omarchy-agent-usage-ollama +++ b/bin/omarchy-agent-usage-ollama @@ -214,7 +214,19 @@ def write_json(path: Path, payload: dict[str, Any]) -> None: # # 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. +# Ollama Cloud API, which serves Claude models alongside the open ones. A +# cloud-tagged model ("glm-4.6:cloud", "gpt-oss:120b-cloud") reached through +# the local daemon's `ollama` provider runs on Ollama's servers too, so it +# counts here; the local collector leaves exactly those out. + + +def is_cloud_model(model: str) -> bool: + tag = model.lower().rpartition(":")[2] if ":" in model else "" + return tag == "cloud" or tag.endswith("-cloud") + + +def billed_by_ollama_cloud(provider: str, model: str) -> bool: + return provider == PROVIDER or (provider == "ollama" and is_cloud_model(model)) # `omp --profile=` (and pi's equivalent) relocates the whole agent tree @@ -292,7 +304,7 @@ def scan_sessions(model_prefix: str) -> tuple[dict[str, Any], bool]: for root in pi_session_roots(): if not root.is_dir(): continue - lines, read_all = matching_lines(root, r'"provider"\s*:\s*"ollama-cloud"') + lines, read_all = matching_lines(root, r'"provider"\s*:\s*"ollama(-cloud)?"') complete = complete and read_all for path, line in lines: try: @@ -304,7 +316,7 @@ def scan_sessions(model_prefix: str) -> tuple[dict[str, Any], bool]: message = entry.get("message") if not isinstance(message, dict) or message.get("role") != "assistant": continue - if str(message.get("provider") or "") != PROVIDER: + if not billed_by_ollama_cloud(str(message.get("provider") or ""), str(message.get("model") or "")): continue usage = message.get("usage") if not isinstance(usage, dict): diff --git a/test/shell.d/agent-usage-ollama-scanner-test.sh b/test/shell.d/agent-usage-ollama-scanner-test.sh index 4343f0d158c..68174eab659 100755 --- a/test/shell.d/agent-usage-ollama-scanner-test.sh +++ b/test/shell.d/agent-usage-ollama-scanner-test.sh @@ -312,3 +312,20 @@ forced=$(HOME="$ROOTS_HOME" PI_CODING_AGENT_DIR="$ROOTS_HOME/moved" XDG_CONFIG_H [[ $(jq -r '.todayTotalTokens' <<<"$again") == "1112" && $(jq -r '.todayTotalTokens' <<<"$forced") == "1119" ]] || fail "Ollama collector reuses a seconds-old scan on a normal run and rescans on --force" "$again $forced" pass "Ollama collector reuses a seconds-old scan and rescans on --force" + +# A cloud-tagged model reached through the local daemon's `ollama` provider is +# billed by Ollama Cloud and counts here; a model the daemon runs itself is +# the local collector's. +TAGGED_HOME=$(mktemp -d) +trap 'rm -rf "$TEST_HOME" "$PI_HOME" "$CACHE_HOME" "$ROOTS_HOME" "$TAGGED_HOME"' EXIT +mkdir -p "$TAGGED_HOME/.pi/agent/sessions/project" +cat >"$TAGGED_HOME/.pi/agent/sessions/project/pi.jsonl" < Date: Wed, 14 Oct 2026 23:29:08 -0400 Subject: [PATCH 111/546] List the manual's coding agents in one order and one table format Co-Authored-By: Claude Opus 5.5 --- manual/17-ai.md | 52 ++++++++++++++++++++++++------------------------- 1 file changed, 26 insertions(+), 26 deletions(-) diff --git a/manual/17-ai.md b/manual/17-ai.md index 995407bbbb1..83cfcb5db8f 100644 --- a/manual/17-ai.md +++ b/manual/17-ai.md @@ -2,35 +2,35 @@ Omarchy treats AI coding agents as first-class citizens, but it doesn't pick a favorite for you. Instead, every major coding-agent CLI comes pre-wired as a lazy-loaded launcher. The launchers are tiny [mise](https://mise.jdx.dev/)-managed stubs in `~/.local/bin/`, so nothing is downloaded until the first time you actually run one. Hermes is the exception: nothing is stubbed for it, and choosing it installs Hermes Desktop, as _Desktop apps_ below explains. Invoke any of these and authenticate when prompted: -| Command | Agent | -|------------|------------------------------------------------------------------| -| `claude` | [Claude Code](https://code.claude.com/docs/en/overview) | -| `codex` | [OpenAI Codex](https://github.com/openai/codex) | -| `opencode` | [OpenCode](https://opencode.ai/) | -| `agy` | [Google Antigravity CLI](https://github.com/google-antigravity/antigravity-cli) | -| `amp` | [Amp](https://ampcode.com/) | -| `commandcode` | [Command Code](https://commandcode.ai) | -| `copilot` | [GitHub Copilot CLI](https://github.com/github/copilot-cli) | -| `crush` | [Crush](https://github.com/charmbracelet/crush) | -| `devin` | [Devin CLI](https://devin.ai/cli) from Cognition | -| `dsh` | [DeepSeek Harness](https://github.com/deepseek-ai/deepseek-harness) | +| Command | Agent | +|---|---| +| `claude` | [Claude Code](https://code.claude.com/docs/en/overview) | +| `codex` | [OpenAI Codex](https://github.com/openai/codex) | +| `opencode` | [OpenCode](https://opencode.ai/) | +| `agy` | [Google Antigravity CLI](https://github.com/google-antigravity/antigravity-cli) | +| `amp` | [Amp](https://ampcode.com/) | +| `commandcode` | [Command Code](https://commandcode.ai) | +| `copilot` | [GitHub Copilot CLI](https://github.com/github/copilot-cli) | +| `crush` | [Crush](https://github.com/charmbracelet/crush) | +| `cursor-agent` | [Cursor CLI](https://cursor.com/cli) | +| `devin` | [Devin CLI](https://devin.ai/cli) from Cognition | +| `dsh` | [DeepSeek Harness](https://github.com/deepseek-ai/deepseek-harness) | | `freebuff` | [Freebuff](https://freebuff.com), Codebuff's free, ad-supported coding agent | -| `goose` | [Goose](https://github.com/aaif-goose/goose) | -| `grok` | [Grok Build](https://x.ai/build) from xAI | -| `kilo` | [Kilo](https://kilo.ai/docs/code-with-ai/platforms/cli), an OpenCode-based coding agent | -| `kimi` | [Kimi Code](https://www.kimi.com/code/docs/en/) | -| `kiro-cli` | [Kiro CLI](https://kiro.dev/docs/cli/) | -| `mcode` | [MiniMax Code](https://github.com/MiniMax-AI/minimax-code), MiniMax's coding agent | +| `goose` | [Goose](https://github.com/aaif-goose/goose) | +| `grok` | [Grok Build](https://x.ai/build) from xAI | +| `hermes` | [Hermes](https://hermes-agent.nousresearch.com/), Nous Research's agent | +| `kilo` | [Kilo](https://kilo.ai/docs/code-with-ai/platforms/cli), an OpenCode-based coding agent | +| `kimi` | [Kimi Code](https://www.kimi.com/code/docs/en/) | +| `kiro-cli` | [Kiro CLI](https://kiro.dev/docs/cli/) | +| `mcode` | [MiniMax Code](https://github.com/MiniMax-AI/minimax-code), MiniMax's coding agent | +| `muse` | [Muse Code](https://dev.meta.ai), Meta's coding agent | +| `omp` | [Oh My Pi](https://github.com/can1357/oh-my-pi) | +| `ori` | [Ori](https://openrouter.ai/docs/guides/ori/harness), OpenRouter's harness | | `phantombot` | [Phantombot](https://phantombot.bot), a persona-first assistant over any harness | -| `qwen` | [Qwen Code](https://github.com/QwenLM/qwen-code), Alibaba's coding agent | -| `vibe` | [Mistral Vibe](https://docs.mistral.ai/vibe/code/overview), Mistral AI's coding agent | +| `pi` | [Mario Zechner's Pi](https://github.com/badlogic/pi-mono) | +| `qwen` | [Qwen Code](https://github.com/QwenLM/qwen-code), Alibaba's coding agent | +| `vibe` | [Mistral Vibe](https://docs.mistral.ai/vibe/code/overview), Mistral AI's coding agent | | `zeroclaw` | [ZeroClaw](https://github.com/zeroclaw-labs/zeroclaw), opened in its `zerocode` terminal UI | -| `pi` | [Mario Zechner's Pi](https://github.com/badlogic/pi-mono) | -| `omp` | [Oh My Pi](https://github.com/can1357/oh-my-pi) | -| `ori` | [Ori](https://openrouter.ai/docs/guides/ori/harness), OpenRouter's harness | -| `hermes` | [Hermes](https://hermes-agent.nousresearch.com/), Nous Research's agent | -| `muse` | [Muse Code](https://dev.meta.ai), Meta's coding agent | -| `cursor-agent` | [Cursor CLI](https://cursor.com/cli) | `ori` is the odd one out: it runs the other harnesses against OpenRouter's whole model catalog, so `ori claude`, `ori codex`, or `ori opencode` start those agents on whichever model you point them at, and `ori code` is Ori's own agent. From eba97869326c8cfe63ff861593615a08a9525049 Mon Sep 17 00:00:00 2001 From: Josh Christopher Algdal Date: Wed, 14 Oct 2026 23:12:33 -0400 Subject: [PATCH 112/546] Add a Mistral Vibe usage collector Read the per-session meta.json files Vibe writes under ~/.vibe/logs/session/ into the agents-panel record, so the panel gains a Mistral Vibe section with tokens by day, the model breakdown, and totals. Vibe exposes no rate-limit windows locally, so the record carries local stats only. The plan comes from ~/.vibe/whoami_cache.json. Generated by Mistral Vibe. Co-Authored-By: Mistral Vibe Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-vibe | 156 +++++++++++++++++++++++++++++++++++ 1 file changed, 156 insertions(+) create mode 100755 bin/omarchy-agent-usage-vibe diff --git a/bin/omarchy-agent-usage-vibe b/bin/omarchy-agent-usage-vibe new file mode 100755 index 00000000000..17983026eb3 --- /dev/null +++ b/bin/omarchy-agent-usage-vibe @@ -0,0 +1,156 @@ +#!/usr/bin/python3 +# omarchy:summary=Print the Mistral Vibe usage record as JSON +# omarchy:args=[--force] [--limits-only] +# omarchy:hidden=true +"""Collect Mistral Vibe usage into one display-ready JSON record. + +Local stats come from per-session meta.json files that Vibe writes under +~/.vibe/logs/session/. Each session carries prompt, completion, and cached +token counts plus the model it ran on, so the panel gets tokens by day and the +all-time model breakdown without a usage API. Vibe exposes no rate-limit +windows locally, so the record carries no limits — the panel shows local +stats only, the same shape Fireworks falls back to without a balance. The plan +tier comes from ~/.vibe/whoami_cache.json. The agents panel only ever reads +the JSON this prints. +""" + +import json +import os +import sys +from datetime import datetime, timedelta, timezone +from pathlib import Path + +AGENT_ID = "vibe" +AGENT_NAME = "Mistral Vibe" + +HOME = Path.home() +VIBE_HOME = Path(os.environ.get("VIBE_HOME", str(HOME / ".vibe"))) +SESSIONS_DIR = VIBE_HOME / "logs" / "session" +WHOAMI_CACHE = VIBE_HOME / "whoami_cache.json" + + +def num(value): + try: + return int(value or 0) + except Exception: + return 0 + + +def empty_bucket(): + return { + "inputTokens": 0, + "outputTokens": 0, + "cacheReadInputTokens": 0, + "cacheCreationInputTokens": 0, + } + + +def session_day(meta): + text = str(meta.get("start_time") or "") + try: + dt = datetime.fromisoformat(text) + if dt.tzinfo is not None: + dt = dt.astimezone() + return dt.strftime("%Y-%m-%d") + except Exception: + return text[:10] if text else datetime.now().strftime("%Y-%m-%d") + + +def session_model(meta): + config = meta.get("config") or {} + model = config.get("active_model") or config.get("routed_default_model") or "" + return str(model or "vibe") or "vibe" + + +def tier_label(): + try: + cache = json.load(open(WHOAMI_CACHE)) + except Exception: + return "" + for entry in cache.values(): + payload = entry.get("payload") if isinstance(entry, dict) else None + if not isinstance(payload, dict): + continue + plan = str(payload.get("plan_name") or "").strip() + if plan: + return plan.capitalize() + plan_type = str(payload.get("plan_type") or "").strip() + if plan_type: + return plan_type.capitalize() + return "" + + +def scan(): + today = datetime.now().astimezone().date() + recent_dates = [(today - timedelta(days=offset)).isoformat() for offset in range(6, -1, -1)] + recent = {day: 0 for day in recent_dates} + today_by_model = {} + model_usage = {} + active_dates = set() + total_sessions = 0 + today_sessions = 0 + + if SESSIONS_DIR.is_dir(): + for meta_path in sorted(SESSIONS_DIR.glob("session_*/meta.json")): + try: + meta = json.load(open(meta_path)) + except Exception: + continue + stats = meta.get("stats") or {} + prompt = num(stats.get("session_prompt_tokens")) + cached = num(stats.get("session_cached_tokens")) + completion = num(stats.get("session_completion_tokens")) + uncached = max(0, prompt - cached) + total = uncached + cached + completion + if total <= 0 and not meta.get("total_messages"): + continue + + day = session_day(meta) + model = session_model(meta) + + bucket = model_usage.setdefault(model, empty_bucket()) + bucket["inputTokens"] += uncached + bucket["outputTokens"] += completion + bucket["cacheReadInputTokens"] += cached + + if day: + active_dates.add(day) + if day in recent: + recent[day] += total + if day == today.isoformat(): + today_by_model[model] = today_by_model.get(model, 0) + total + today_sessions += 1 + total_sessions += 1 + + return { + "schemaVersion": 1, + "id": AGENT_ID, + "name": AGENT_NAME, + "updatedAt": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"), + "ready": total_sessions > 0, + "hasLocalStats": True, + "hasPromptStats": False, + "tierLabel": tier_label(), + "usageStatusText": "" if total_sessions > 0 else "No Vibe sessions found", + "authHelpText": "", + "limits": [], + "todayPrompts": 0, + "todaySessions": today_sessions, + "todayTotalTokens": sum(today_by_model.values()), + "todayTokensByModel": today_by_model, + "recentDays": [{"date": day, "messageCount": recent[day]} for day in recent_dates], + "totalPrompts": 0, + "totalSessions": total_sessions, + "activeDays": len(active_dates), + "activeDates": sorted(active_dates), + "modelUsage": model_usage, + } + + +def main(): + print(json.dumps(scan(), separators=(",", ":"))) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) From 4f372fea1f9dae3b229b33eee94109294a842386 Mon Sep 17 00:00:00 2001 From: Josh Christopher Algdal Date: Tue, 8 Sep 2026 18:34:44 +0200 Subject: [PATCH 113/546] Add Mistral Vibe brand marks for the agents panel Gradient M icon (vibe.svg) for dark surfaces and monochrome version (vibe-light.svg) for light surfaces, sourced from the Mistral brand page. The panel resolves these by convention from assets/.svg. Generated by Mistral Vibe. Co-Authored-By: Mistral Vibe --- shell/plugins/agents/assets/vibe-light.svg | 1 + shell/plugins/agents/assets/vibe.svg | 1 + 2 files changed, 2 insertions(+) create mode 100644 shell/plugins/agents/assets/vibe-light.svg create mode 100644 shell/plugins/agents/assets/vibe.svg diff --git a/shell/plugins/agents/assets/vibe-light.svg b/shell/plugins/agents/assets/vibe-light.svg new file mode 100644 index 00000000000..94aea0d2752 --- /dev/null +++ b/shell/plugins/agents/assets/vibe-light.svg @@ -0,0 +1 @@ +Mistral Vibe \ No newline at end of file diff --git a/shell/plugins/agents/assets/vibe.svg b/shell/plugins/agents/assets/vibe.svg new file mode 100644 index 00000000000..fd905ddb705 --- /dev/null +++ b/shell/plugins/agents/assets/vibe.svg @@ -0,0 +1 @@ +Mistral Vibe \ No newline at end of file From 286976215ec9919feff6b6d4caf7fd5129415d91 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:18:07 -0400 Subject: [PATCH 114/546] Count only Mistral-served Vibe usage, plus pi and opencode on mistral Vibe also drives other providers, a local llama.cpp server out of the box, so a session counts only when the model it resolved to is served by Mistral's API. Subagent sessions count toward the session that spawned them, prompts are the user messages Vibe didn't inject, and VIBE_HOME and config.toml's session_logging are honored. pi, omp, and opencode messages on the mistral provider join the record, which no other collector claims. The collector now takes the standard flags, caches its scan per day under flock with a per-session index, rejects untrusted times and token fields, and names the plan as Vibe's own resolve_user_plan does. Mistral publishes no limits, so the status stays empty and there are no meters. The colored mark reads on light surfaces too, so it needs no light twin. Co-authored-by: Josh Christopher Algdal Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-vibe | 828 ++++++++++++++++++--- shell/plugins/agents/assets/vibe-light.svg | 1 - test/shell.d/agent-usage-vibe-test.sh | 234 ++++++ 3 files changed, 972 insertions(+), 91 deletions(-) delete mode 100644 shell/plugins/agents/assets/vibe-light.svg create mode 100755 test/shell.d/agent-usage-vibe-test.sh diff --git a/bin/omarchy-agent-usage-vibe b/bin/omarchy-agent-usage-vibe index 17983026eb3..ae7985f8051 100755 --- a/bin/omarchy-agent-usage-vibe +++ b/bin/omarchy-agent-usage-vibe @@ -4,36 +4,135 @@ # omarchy:hidden=true """Collect Mistral Vibe usage into one display-ready JSON record. -Local stats come from per-session meta.json files that Vibe writes under -~/.vibe/logs/session/. Each session carries prompt, completion, and cached -token counts plus the model it ran on, so the panel gets tokens by day and the -all-time model breakdown without a usage API. Vibe exposes no rate-limit -windows locally, so the record carries no limits — the panel shows local -stats only, the same shape Fireworks falls back to without a balance. The plan -tier comes from ~/.vibe/whoami_cache.json. The agents panel only ever reads +Local stats come from the per-session meta.json files Vibe writes under +~/.vibe/logs/session/ (VIBE_HOME moves it, and config.toml's +session_logging.save_dir moves the sessions), plus pi/omp and opencode +sessions that ran on the `mistral` provider. Vibe also drives other providers +(a local llama.cpp server out of the box), so a Vibe session only counts when +the model it ran on is served by Mistral. The plan comes from the whoami +answer Vibe caches in whoami_cache.json. Mistral publishes no usage limits +to read, so the record carries no meters. The agents panel only ever reads the JSON this prints. """ +import argparse +import fcntl +import hashlib import json import os +import shutil +import sqlite3 +import subprocess import sys +import tempfile +import time +import tomllib from datetime import datetime, timedelta, timezone from pathlib import Path +from urllib.parse import urlparse AGENT_ID = "vibe" AGENT_NAME = "Mistral Vibe" -HOME = Path.home() -VIBE_HOME = Path(os.environ.get("VIBE_HOME", str(HOME / ".vibe"))) -SESSIONS_DIR = VIBE_HOME / "logs" / "session" -WHOAMI_CACHE = VIBE_HOME / "whoami_cache.json" +# The provider id pi, omp, and opencode give Mistral's own API. A model reached +# through a gateway (OpenRouter, a proxy) is billed there and counted there. +PROVIDER = "mistral" +# The alias an unpinned Vibe config resolves to when it is configured, as +# Vibe's own resolve_default_model_alias does. +DEFAULT_MODEL_ALIAS = "mistral-medium-3.5" -def num(value): - try: - return int(value or 0) - except Exception: +# 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). --limits-only promises only fresh limits, and Vibe has none, so it may +# reuse a scan for 15 minutes. +SCAN_REUSE_SECONDS = 20 +LIMITS_ONLY_REUSE_SECONDS = 900 + +CACHE_SCHEMA_VERSION = 1 +INDEX_VERSION = 1 + +# Counts and times read from disk are untrusted. Anything past these bounds is +# a corrupt or crafted field, not usage. +MAX_TOKENS = 10**12 +MIN_EPOCH_SECONDS = 946_684_800 # 2000-01-01 +MAX_MODEL_NAME_LEN = 120 + + +def token_count(value): + """A token field as a whole number, or 0 when it isn't a plausible count. + + bool is an int subclass, and a string, NaN, negative, or absurd value is a + bad field: it counts as nothing rather than skewing the totals. + """ + if isinstance(value, bool) or not isinstance(value, (int, float)): + return 0 + if value != value or value <= 0 or value > MAX_TOKENS: return 0 + return int(value) + + +def epoch_seconds(value): + """Seconds since the epoch for an ISO string or a seconds/milliseconds + number, or None when the value can't be a real moment of use.""" + if value is None or isinstance(value, bool): + return None + if isinstance(value, (int, float)): + seconds = float(value) + elif isinstance(value, str) and value.strip(): + text = value.strip() + try: + seconds = float(text) + except ValueError: + try: + parsed = datetime.fromisoformat(text[:-1] + "+00:00" if text.endswith("Z") else text) + except ValueError: + return None + if parsed.tzinfo is None: + parsed = parsed.astimezone() + seconds = parsed.timestamp() + else: + return None + if seconds != seconds: + return None + if seconds > 10_000_000_000: + seconds /= 1000 + # A day of slack past now covers a clock a little behind the writer's. + if seconds < MIN_EPOCH_SECONDS or seconds > time.time() + 86_400: + return None + return seconds + + +def local_day(seconds): + return datetime.fromtimestamp(seconds).strftime("%Y-%m-%d") + + +def model_label(raw): + """A model id fit for the record: its last path segment, without a + thinking-level suffix, or None when it is too long to be a real id.""" + value = str(raw or "").rstrip("/").split("/")[-1].split(":")[0] + if len(value) > MAX_MODEL_NAME_LEN: + return None + return value or "mistral" + + +def obj(value): + return value if isinstance(value, dict) else {} + + +def runtime_env(): + home = str(Path.home()) + path_parts = [ + os.environ.get("PATH", ""), + f"{home}/.local/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 empty_bucket(): @@ -45,112 +144,661 @@ def empty_bucket(): } -def session_day(meta): - text = str(meta.get("start_time") or "") +class Stats: + """Token and activity totals shared by every local source. + + Sessions and days are sets, so a day active in both Vibe and pi counts once + toward activeDays. + """ + + def __init__(self): + 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() + self.today_tokens_by_model = {} + self.total_prompts = 0 + self.total_sessions = set() + self.active_days = set() + self.model_usage = {} + + def add(self, session, day, model, input_tokens, output_tokens, cache_read, cache_write, prompts): + total = input_tokens + output_tokens + cache_read + cache_write + if total <= 0 and prompts <= 0: + return + self.total_prompts += prompts + self.total_sessions.add(session) + self.active_days.add(day) + if total > 0: + bucket = self.model_usage.setdefault(model, empty_bucket()) + bucket["inputTokens"] += input_tokens + bucket["outputTokens"] += output_tokens + bucket["cacheReadInputTokens"] += cache_read + bucket["cacheCreationInputTokens"] += cache_write + if day in self.recent: + # recentDays.messageCount holds tokens, despite the shared legacy name. + self.recent[day] += total + if day == self.today: + self.today_prompts += prompts + self.today_sessions.add(session) + self.today_total_tokens += total + if total > 0: + self.today_tokens_by_model[model] = self.today_tokens_by_model.get(model, 0) + total + + def record(self): + 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), + # The dates travel too: merging snapshots from several machines needs + # their union, which a count alone cannot give. + "activeDays": len(self.active_days), + "activeDates": sorted(self.active_days), + "modelUsage": self.model_usage, + } + + +# ------------------------------------------------------------------ Vibe + + +def vibe_home(): + configured = os.environ.get("VIBE_HOME", "").strip() + if configured: + return Path(configured).expanduser().resolve() + return Path.home() / ".vibe" + + +def session_logging(): + """Where Vibe saves sessions and the prefix it names them with, honoring + [session_logging] in Vibe's own config.toml as Vibe does.""" + home = vibe_home() + save_dir = home / "logs" / "session" + prefix = "session" try: - dt = datetime.fromisoformat(text) - if dt.tzinfo is not None: - dt = dt.astimezone() - return dt.strftime("%Y-%m-%d") - except Exception: - return text[:10] if text else datetime.now().strftime("%Y-%m-%d") + with (home / "config.toml").open("rb") as handle: + section = obj(tomllib.load(handle).get("session_logging")) + except (OSError, tomllib.TOMLDecodeError, UnicodeDecodeError): + section = {} + configured_dir = section.get("save_dir") + if isinstance(configured_dir, str) and configured_dir.strip(): + path = Path(configured_dir.strip()).expanduser() + # Vibe resolves a relative save_dir against wherever it was started, which + # this collector can't know; such a config keeps the default location. + if path.is_absolute(): + save_dir = path + configured_prefix = section.get("session_prefix") + if isinstance(configured_prefix, str) and configured_prefix.strip() and "/" not in configured_prefix: + prefix = configured_prefix.strip() + return save_dir, prefix + + +def resolve_model(config): + """The model entry a session ran on, resolved the way Vibe resolves it: the + pinned alias, else the experiment-routed one, else the default alias, else + the first configured model.""" + models = config.get("models") + # Current Vibe dumps the models keyed by alias; older releases as a list. + if isinstance(models, dict): + entries = list(models.values()) + elif isinstance(models, list): + entries = models + else: + return None + by_alias = {} + for entry in entries: + if isinstance(entry, dict): + alias = str(entry.get("alias") or entry.get("name") or "") + if alias and alias not in by_alias: + by_alias[alias] = entry + for alias in (config.get("active_model"), config.get("routed_default_model"), DEFAULT_MODEL_ALIAS): + if isinstance(alias, str) and alias in by_alias: + return by_alias[alias] + return next(iter(by_alias.values()), None) + + +def is_mistral_provider(name, providers): + """Whether the named provider is Mistral's own API. + + Vibe marks Mistral's providers with the `mistral` backend; a generic one + pointed at a mistral.ai host is the same API. A provider the config doesn't + describe only counts under Vibe's default `mistral` name. + """ + for provider in providers if isinstance(providers, list) else []: + if not isinstance(provider, dict) or str(provider.get("name") or "") != name: + continue + if str(provider.get("backend") or "").lower() == "mistral": + return True + host = (urlparse(str(provider.get("api_base") or "")).hostname or "").lower() + return host == "mistral.ai" or host.endswith(".mistral.ai") + return name == PROVIDER def session_model(meta): - config = meta.get("config") or {} - model = config.get("active_model") or config.get("routed_default_model") or "" - return str(model or "vibe") or "vibe" + """The model label a session ran on when Mistral served it, else None.""" + config = obj(meta.get("config")) + model = resolve_model(config) + if not isinstance(model, dict): + return None + if not is_mistral_provider(str(model.get("provider") or ""), config.get("providers")): + return None + return model_label(model.get("alias") or model.get("name")) + + +def count_prompts(messages_path): + """The prompts typed into a session: user messages Vibe did not inject + itself, such as reminders and compaction summaries.""" + prompts = 0 + with messages_path.open("rb") as handle: + for raw in handle: + if b'"user"' not in raw: + continue + try: + message = json.loads(raw) + except ValueError: + continue + if not isinstance(message, dict) or message.get("role") != "user": + continue + if message.get("injected") is True or message.get("context_boundary"): + continue + prompts += 1 + return prompts + + +def file_signature(path): + try: + stat = path.stat() + except FileNotFoundError: + return None + return [stat.st_size, stat.st_mtime_ns] + + +def parse_meta(meta_path, count_session_prompts): + """One session's index row: [startSeconds, model, input, output, cacheRead, + prompts], or None when it ran on another provider or used nothing. + + Vibe keeps running totals per session, so the whole session is booked on + the day it started. + """ + meta = json.loads(meta_path.read_text(encoding="utf-8")) + if not isinstance(meta, dict): + return None + model = session_model(meta) + started = epoch_seconds(meta.get("start_time")) + if model is None or started is None: + return None + stats = obj(meta.get("stats")) + prompt_tokens = token_count(stats.get("session_prompt_tokens")) + # Mistral reports cached prompt tokens inside prompt_tokens, as OpenAI does. + cached = min(token_count(stats.get("session_cached_tokens")), prompt_tokens) + completion = token_count(stats.get("session_completion_tokens")) + prompts = 0 + messages = meta_path.parent / "messages.jsonl" + if count_session_prompts and messages.is_file(): + prompts = count_prompts(messages) + if prompt_tokens + completion <= 0 and prompts <= 0: + return None + return [started, model, prompt_tokens - cached, completion, cached, prompts] + + +def load_index(path): + try: + data = json.loads(path.read_text(encoding="utf-8")) + except Exception: + return {} + if isinstance(data, dict) and data.get("version") == INDEX_VERSION and isinstance(data.get("files"), dict): + return data["files"] + return {} + + +def scan_vibe_sessions(stats, index_file): + """Add Vibe sessions that ran on Mistral. Returns whether every session was + read, so a partial scan is never cached. + + A session's subagents keep their own meta.json under its agents/ directory + with their own totals; they count toward the session that spawned them, and + their prompts are the parent's instructions rather than anything typed. + """ + save_dir, prefix = session_logging() + previous = load_index(index_file) if index_file else {} + current = {} + complete = True + try: + session_dirs = sorted(path for path in save_dir.glob(f"{prefix}_*") if path.is_dir()) + except OSError: + return False + + for session_dir in session_dirs: + metas = [session_dir / "meta.json"] + agents = session_dir / "agents" + if agents.is_dir(): + metas.extend(sorted(agents.rglob("meta.json"))) + for meta_path in metas: + key = str(meta_path) + top_level = meta_path.parent == session_dir + try: + signature = [file_signature(meta_path), file_signature(meta_path.parent / "messages.jsonl") if top_level else None] + if signature[0] is None: + continue + held = previous.get(key) + if isinstance(held, dict) and held.get("sig") == signature: + entry = held + else: + entry = {"sig": signature, "row": parse_meta(meta_path, top_level)} + except (OSError, ValueError, UnicodeDecodeError) as exc: + # Mid-write or unreadable: skip it this run, uncached. + print(f"omarchy-agent-usage-vibe: skipping {meta_path}: {type(exc).__name__}", file=sys.stderr) + complete = False + continue + current[key] = entry + row = entry.get("row") + if row is None: + continue + try: + started, model, input_tokens, output_tokens, cache_read, prompts = row + day = local_day(float(started)) + counts = [int(input_tokens), int(output_tokens), int(cache_read), int(prompts)] + except (TypeError, ValueError, OverflowError, OSError): + continue + stats.add(str(session_dir), day, str(model), counts[0], counts[1], counts[2], 0, counts[3]) + + if index_file and current != previous: + try: + write_json(index_file, {"version": INDEX_VERSION, "files": current}) + except Exception as exc: + print(f"omarchy-agent-usage-vibe: could not write session index ({exc})", file=sys.stderr) + return complete def tier_label(): + """The plan from the newest whoami answer Vibe cached, named as Vibe's own + resolve_user_plan names it. A plan Vibe has no name for stays unnamed.""" try: - cache = json.load(open(WHOAMI_CACHE)) + entries = json.loads((vibe_home() / "whoami_cache.json").read_text(encoding="utf-8")) except Exception: return "" - for entry in cache.values(): - payload = entry.get("payload") if isinstance(entry, dict) else None - if not isinstance(payload, dict): + newest = None + for entry in entries.values() if isinstance(entries, dict) else []: + entry = obj(entry) + stored = entry.get("stored_at_timestamp") + payload = obj(entry.get("payload")) + if isinstance(stored, bool) or not isinstance(stored, (int, float)) or not payload: continue - plan = str(payload.get("plan_name") or "").strip() - if plan: - return plan.capitalize() - plan_type = str(payload.get("plan_type") or "").strip() - if plan_type: - return plan_type.capitalize() + if newest is None or stored > newest[0]: + newest = (stored, payload) + if newest is None: + return "" + kind = str(newest[1].get("plan_type") or "").strip().lower() + name = str(newest[1].get("plan_name") or "").strip().upper() + if kind == "chat": + return {"FREE": "Free", "INDIVIDUAL": "Pro", "EDU": "Student", "TEAM": "Team"}.get(name, "") + if kind == "mistral_code": + return {"F": "Free Codestral", "E": "Code Enterprise"}.get(name, "") + if kind == "api" and name: + return "Free API" if "FREE" in name else "PAYG API" return "" -def scan(): - today = datetime.now().astimezone().date() - recent_dates = [(today - timedelta(days=offset)).isoformat() for offset in range(6, -1, -1)] - recent = {day: 0 for day in recent_dates} - today_by_model = {} - model_usage = {} - active_dates = set() - total_sessions = 0 - today_sessions = 0 +# ----------------------------------------------------------- pi and omp - if SESSIONS_DIR.is_dir(): - for meta_path in sorted(SESSIONS_DIR.glob("session_*/meta.json")): + +def find_command(name): + return shutil.which(name, path=ENV.get("PATH")) + + +# `omp --profile=` (and pi's equivalent) relocates the whole agent tree +# under /profiles//, so a subscription driven entirely through a +# profile leaves the default root empty. Sessions are keyed by path, so a +# profile adds sessions rather than double-counting the default root. +def pi_session_roots(): + pi_agent = os.environ.get("PI_CODING_AGENT_DIR", "").strip() + roots = [Path(pi_agent).expanduser() / "sessions" if pi_agent else Path.home() / ".pi" / "agent" / "sessions"] + roots.append(Path.home() / ".omp" / "agent" / "sessions") + for base in (Path.home() / ".pi", Path.home() / ".omp"): + try: + # Sorted so the scan order does not depend on directory order. + roots.extend(sorted(child / "agent" / "sessions" for child in (base / "profiles").iterdir() if child.is_dir())) + except OSError: + # No profiles directory, or it is unreadable: the default root stands. + pass + unique = [] + for root in roots: + if root not in unique: + unique.append(root) + return unique + + +def scan_pi_sessions(stats): + """Add pi/omp assistant messages that ran on the mistral provider. Returns + whether the scan could run at all.""" + roots = [root for root in pi_session_roots() if root.is_dir()] + if not roots: + return True + rg = find_command("rg") + if not rg: + print("omarchy-agent-usage-vibe: rg is missing; pi sessions not counted", file=sys.stderr) + return False + seen = set() + for root in roots: + try: + proc = subprocess.Popen( + # --no-ignore: these are data files, not a source tree. When $HOME is + # itself a git checkout, its ignore rules would otherwise hide every + # session file. + [rg, "--json", "--no-ignore", "-e", r'"provider"\s*:\s*"mistral"', str(root)], + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + text=True, + errors="replace", + env=ENV, + ) + except OSError: + return False + + assert proc.stdout is not None + for raw in proc.stdout: try: - meta = json.load(open(meta_path)) + event = json.loads(raw) + if event.get("type") != "match": + continue + data = obj(event.get("data")) + line = obj(data.get("lines")).get("text", "") + path = os.path.abspath(str(obj(data.get("path")).get("text") or "pi-session")) + entry = json.loads(line) except Exception: continue - stats = meta.get("stats") or {} - prompt = num(stats.get("session_prompt_tokens")) - cached = num(stats.get("session_cached_tokens")) - completion = num(stats.get("session_completion_tokens")) - uncached = max(0, prompt - cached) - total = uncached + cached + completion - if total <= 0 and not meta.get("total_messages"): + if not isinstance(entry, dict) or entry.get("type") != "message": + continue + message = obj(entry.get("message")) + if message.get("role") != "assistant" or str(message.get("provider") or "") != PROVIDER: continue + # A fork copies its parent's messages with their ids and times, so the + # pair names a message across files. A line without an id always counts. + message_id = str(entry.get("id") or "") + stamp = str(entry.get("timestamp") or "") + if message_id and stamp: + key = (message_id, stamp) + if key in seen: + continue + seen.add(key) + started = epoch_seconds(entry.get("timestamp") or message.get("timestamp")) + model = model_label(message.get("model")) + if started is None or model is None: + continue + usage = obj(message.get("usage")) + input_tokens = token_count(usage.get("input")) + output_tokens = token_count(usage.get("output")) + cache_read = token_count(usage.get("cacheRead")) + cache_write = token_count(usage.get("cacheWrite")) + if not (input_tokens or output_tokens or cache_read or cache_write): + input_tokens = token_count(usage.get("totalTokens")) + if not (input_tokens or output_tokens or cache_read or cache_write): + continue + stats.add(path, local_day(started), model, input_tokens, output_tokens, cache_read, cache_write, 1) - day = session_day(meta) - model = session_model(meta) + try: + proc.wait(timeout=1) + except Exception: + proc.kill() + return True - bucket = model_usage.setdefault(model, empty_bucket()) - bucket["inputTokens"] += uncached - bucket["outputTokens"] += completion - bucket["cacheReadInputTokens"] += cached - if day: - active_dates.add(day) - if day in recent: - recent[day] += total - if day == today.isoformat(): - today_by_model[model] = today_by_model.get(model, 0) + total - today_sessions += 1 - total_sessions += 1 +# ------------------------------------------------------------- opencode - return { + +def opencode_db(): + # OPENCODE_DB points opencode at another database, as the CLI itself reads + # it; a relative name lives beside the default one. Only a build whose + # channel was literally "v2" wrote opencode-v2.db. + root = Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) / "opencode" + configured = os.environ.get("OPENCODE_DB", "").strip() + if configured: + path = Path(configured).expanduser() + return path if path.is_absolute() else root / path + normal = root / "opencode.db" + legacy = root / "opencode-v2.db" + return legacy if not normal.is_file() and legacy.is_file() else normal + + +def table_columns(conn, table): + try: + return {row[1] for row in conn.execute(f"PRAGMA table_info({table})")} + except sqlite3.Error: + return set() + + +def add_opencode_row(stats, message_id, session_id, entry, row_type, created, seen): + # The original `message` table keeps the provider and model at the top + # level; the V2 `session_message` table nests them under `model` and keeps + # the role in its `type` column, with a compaction's tokens spent like any + # reply's. + if not isinstance(entry, dict): + return + if (entry.get("role") or row_type) not in ("assistant", "compaction"): + return + nested = obj(entry.get("model")) + # Exact match: opencode provider ids are free-form, and a custom gateway + # named like Mistral is not Mistral's API. + if str(entry.get("providerID") or nested.get("providerID") or "") != PROVIDER: + return + tokens = obj(entry.get("tokens")) + cache = obj(tokens.get("cache")) + input_tokens = token_count(tokens.get("input")) + # opencode keeps thinking tokens out of output; both are generated. + output_tokens = token_count(tokens.get("output")) + token_count(tokens.get("reasoning")) + cache_read = token_count(cache.get("read")) + cache_write = token_count(cache.get("write")) + if not (input_tokens or output_tokens or cache_read or cache_write): + return + # A migrated message can sit in both stores under the same id. + if message_id in seen: + return + started = epoch_seconds(obj(entry.get("time")).get("created") or created) + model = model_label(entry.get("modelID") or nested.get("id") or nested.get("modelID")) + if started is None or model is None: + return + seen.add(message_id) + stats.add("opencode:" + str(session_id), local_day(started), model, input_tokens, output_tokens, cache_read, cache_write, 1) + + +def scan_opencode_db(stats): + """Add opencode messages on the mistral provider. Read-only: opencode may be + writing right now. Returns whether the scan ran to completion; a scan cut + short still contributes what it read, but isn't cached.""" + db = opencode_db() + 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") + tables = {name for (name,) in conn.execute("SELECT name FROM sqlite_master WHERE type = 'table'")} + seen = set() + # The LIKE gates only skip rows whose JSON cannot name the provider; each + # row is still checked after parsing. + columns = table_columns(conn, "session_message") + if "session_message" in tables and {"id", "session_id", "type", "data"} <= columns: + created = "m.time_created" if "time_created" in columns else "NULL" + source = "session_message m" + fork_filter = "" + if "time_created" in columns and "session_v2" in tables and {"id", "fork_session_id", "time_created"} <= table_columns(conn, "session_v2"): + # A fork copies its source session's messages under new ids; only + # what came after the fork was spent in it. + source += " LEFT JOIN session_v2 s ON s.id = m.session_id" + fork_filter = " AND (s.fork_session_id IS NULL OR m.time_created >= s.time_created)" + for message_id, session_id, row_type, time_created, raw in conn.execute( + f"SELECT m.id, m.session_id, m.type, {created}, m.data FROM {source}" + " WHERE m.type IN ('assistant', 'compaction')" + " AND m.data LIKE '%\"providerID\"%:%\"mistral\"%'" + fork_filter + ): + try: + add_opencode_row(stats, str(message_id), session_id, json.loads(raw), str(row_type or ""), time_created, seen) + except Exception: + continue + columns = table_columns(conn, "message") + if "message" in tables and {"session_id", "data"} <= columns: + row_id = "id" if "id" in columns else "rowid" + for message_id, session_id, raw in conn.execute( + f"SELECT {row_id}, session_id, data FROM message" + " WHERE data LIKE '%\"providerID\"%:%\"mistral\"%'" + ): + try: + add_opencode_row(stats, str(message_id), session_id, json.loads(raw), "", None, seen) + except Exception: + continue + except sqlite3.Error: + # Transient lock, schema migration, corruption: the numbers stop here, + # incomplete. + return False + finally: + conn.close() + return True + + +# ---------------------------------------------------------------- cache + + +def cache_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 cache_paths(): + # The digest covers every data path the scan reads, so another home or + # Vibe directory never reuses this one's numbers. + save_dir, prefix = session_logging() + sources = [str(Path.home()), str(save_dir), prefix, str(opencode_db()), os.environ.get("PI_CODING_AGENT_DIR", "")] + digest = hashlib.sha1("\n".join(sources).encode("utf-8")).hexdigest()[:16] + root = cache_root() + return root / f"vibe-scan-{digest}.json", root / f"vibe-scan-{digest}.lock", root / f"vibe-index-{digest}.json" + + +def read_fresh_json(path, max_age_seconds): + if max_age_seconds <= 0: + 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, payload): + # A temp name unique to this writer: 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 + + +STATS_KEYS = ("todayPrompts", "todaySessions", "todayTotalTokens", "todayTokensByModel", "recentDays", "totalPrompts", "totalSessions", "activeDays", "activeDates", "modelUsage") + + +def read_cached_stats(cache_file, max_age_seconds): + cached = read_fresh_json(cache_file, max_age_seconds) + if not isinstance(cached, dict) or cached.get("schemaVersion") != CACHE_SCHEMA_VERSION: + return None + # today* fields only mean "today" on the day they were scanned. + if cached.get("scanDate") != datetime.now().strftime("%Y-%m-%d"): + return None + stats = cached.get("stats") + if not isinstance(stats, dict) or not all(key in stats for key in STATS_KEYS): + return None + return stats + + +def run_scans(index_file): + stats = Stats() + complete = scan_vibe_sessions(stats, index_file) + complete = scan_pi_sessions(stats) and complete + complete = scan_opencode_db(stats) and complete + return stats.record(), complete + + +def cached_local_stats(max_age): + """Local stats, with the cache as a pure optimization: any cache failure + degrades to a direct scan and a warning on stderr.""" + try: + cache_file, lock_file, index_file = cache_paths() + cached = read_cached_stats(cache_file, max_age) + if cached is not None: + return cached + with lock_file.open("w") as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + cached = read_cached_stats(cache_file, max_age) + if cached is not None: + return cached + stats, complete = run_scans(index_file) + # An interrupted scan still serves this run, but caching it would hide + # the missing usage from every reader until the cache expires. + if complete: + try: + write_json(cache_file, {"schemaVersion": CACHE_SCHEMA_VERSION, "scanDate": datetime.now().strftime("%Y-%m-%d"), "stats": stats}) + except Exception as exc: + print(f"omarchy-agent-usage-vibe: could not write usage cache ({exc})", file=sys.stderr) + return stats + except Exception as exc: + print(f"omarchy-agent-usage-vibe: cache unavailable ({exc}); scanning directly", file=sys.stderr) + stats, _ = run_scans(None) + return stats + + +def main(): + parser = argparse.ArgumentParser(description="Print the Mistral Vibe usage record as JSON") + # --force rescans everything. --limits-only only promises fresh limits, and + # Mistral offers none, so it reuses a scan far longer than a normal run. + parser.add_argument("--force", action="store_true") + parser.add_argument("--limits-only", action="store_true") + args = parser.parse_args() + + max_age = 0 if args.force else (LIMITS_ONLY_REUSE_SECONDS if args.limits_only else SCAN_REUSE_SECONDS) + stats = cached_local_stats(max_age) + tier = tier_label() + + # Mistral publishes no usage limits, so there is nothing to wait for or + # fail at: the status stays empty and local stats show alone. A machine + # with neither Mistral usage nor a cached plan prints a record the panel + # skips. + record = { "schemaVersion": 1, "id": AGENT_ID, "name": AGENT_NAME, - "updatedAt": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"), - "ready": total_sessions > 0, + "updatedAt": datetime.now(timezone.utc).isoformat(), + "ready": stats["totalSessions"] > 0 or bool(tier), "hasLocalStats": True, - "hasPromptStats": False, - "tierLabel": tier_label(), - "usageStatusText": "" if total_sessions > 0 else "No Vibe sessions found", + "hasPromptStats": True, + "tierLabel": tier, + "usageStatusText": "", "authHelpText": "", "limits": [], - "todayPrompts": 0, - "todaySessions": today_sessions, - "todayTotalTokens": sum(today_by_model.values()), - "todayTokensByModel": today_by_model, - "recentDays": [{"date": day, "messageCount": recent[day]} for day in recent_dates], - "totalPrompts": 0, - "totalSessions": total_sessions, - "activeDays": len(active_dates), - "activeDates": sorted(active_dates), - "modelUsage": model_usage, + "limitsStale": False, + "limitsFetchedAt": 0, } - - -def main(): - print(json.dumps(scan(), separators=(",", ":"))) + record.update(stats) + print(json.dumps(record, separators=(",", ":"))) return 0 if __name__ == "__main__": - sys.exit(main()) + raise SystemExit(main()) diff --git a/shell/plugins/agents/assets/vibe-light.svg b/shell/plugins/agents/assets/vibe-light.svg deleted file mode 100644 index 94aea0d2752..00000000000 --- a/shell/plugins/agents/assets/vibe-light.svg +++ /dev/null @@ -1 +0,0 @@ -Mistral Vibe \ No newline at end of file diff --git a/test/shell.d/agent-usage-vibe-test.sh b/test/shell.d/agent-usage-vibe-test.sh new file mode 100755 index 00000000000..dec0ee96d5b --- /dev/null +++ b/test/shell.d/agent-usage-vibe-test.sh @@ -0,0 +1,234 @@ +#!/bin/bash + +source "$(dirname "$0")/base-test.sh" + +require_command jq +require_command python3 +require_command rg + +COLLECTOR="$ROOT/bin/omarchy-agent-usage-vibe" +TMP_ROOT=$(mktemp -d) +trap 'rm -rf "$TMP_ROOT"' EXIT + +unset VIBE_HOME PI_CODING_AGENT_DIR OPENCODE_DB XDG_DATA_HOME XDG_CACHE_HOME + +today=$(date +%Y-%m-%d) +now_iso=$(date -u +%Y-%m-%dT%H:%M:%S+00:00) +two_days_ago=$(date -d '2 days ago' +%Y-%m-%d) +two_days_ago_iso=$(date -u -d '2 days ago' +%Y-%m-%dT%H:%M:%S+00:00) + +new_home() { + local home="$TMP_ROOT/$1" + mkdir -p "$home" + printf '%s\n' "$home" +} + +collect() { + local home="$1" + shift + HOME="$home" XDG_CACHE_HOME="$home/.cache" XDG_DATA_HOME="$home/.local/share" "$COLLECTOR" "$@" +} + +# write_meta +write_meta() { + mkdir -p "$1" + cat >"$1/meta.json" <"$sessions/session_20260101_120000_aaaa/messages.jsonl" <<'EOF' +{"role": "user", "content": "Fix the bug"} +{"role": "assistant", "content": "Done"} +{"role": "user", "content": "", "injected": true} +{"role": "user", "content": "summary", "context_boundary": "compaction"} +{"role": "tool", "content": "ok"} +{"role": "user", "content": "Thanks"} +EOF +# A subagent's totals belong to the session that spawned it; its task is not +# a prompt anyone typed. +write_meta "$sessions/session_20260101_120000_aaaa/agents/explore_20260101_120100_bbbb" "$now_iso" 300 0 50 "$MISTRAL_CONFIG" +printf '%s\n' '{"role": "user", "content": "Look around"}' >"$sessions/session_20260101_120000_aaaa/agents/explore_20260101_120100_bbbb/messages.jsonl" +# A session on the local llama.cpp provider is not Mistral usage. +write_meta "$sessions/session_20260101_130000_cccc" "$now_iso" 9000 0 9000 "$LOCAL_CONFIG" +printf '%s\n' '{"role": "user", "content": "local"}' >"$sessions/session_20260101_130000_cccc/messages.jsonl" +# Running totals are booked on the day the session started. +write_meta "$sessions/session_20260101_140000_dddd" "$two_days_ago_iso" 100 0 10 "$MISTRAL_CONFIG" + +record=$(collect "$home" --force) +[[ $(jq -c '{ready, todayPrompts, todaySessions, todayTotalTokens, totalPrompts, totalSessions, activeDays}' <<<"$record") == '{"ready":true,"todayPrompts":2,"todaySessions":1,"todayTotalTokens":1550,"totalPrompts":2,"totalSessions":2,"activeDays":2}' ]] || + fail "Vibe collector counts Mistral sessions, their subagents, and typed prompts only" "$record" +pass "Vibe collector counts Mistral sessions, their subagents, and typed prompts only" + +[[ $(jq -c '.modelUsage' <<<"$record") == '{"mistral-medium-3.5":{"inputTokens":1000,"outputTokens":260,"cacheReadInputTokens":400,"cacheCreationInputTokens":0}}' ]] || + fail "Vibe collector splits cached prompt tokens out of input and skips other providers" "$record" +pass "Vibe collector splits cached prompt tokens out of input and skips other providers" + +[[ $(jq -r --arg day "$two_days_ago" '.recentDays[] | select(.date == $day) | .messageCount' <<<"$record") == "110" ]] || + fail "Vibe collector books a session on the day it started" "$record" +pass "Vibe collector books a session on the day it started" + +# ---------------------------------------------------------------- model resolution + +home=$(new_home resolve) +sessions="$home/.vibe/logs/session" +# Unpinned: Vibe falls back to its default alias. +write_meta "$sessions/session_1" "$now_iso" 10 0 0 '{"active_model":"","models":{"local":{"name":"q","provider":"llamacpp","alias":"local"},"mistral-medium-3.5":{"name":"m","provider":"mistral","alias":"mistral-medium-3.5"}},"providers":[{"name":"mistral","backend":"mistral"},{"name":"llamacpp","backend":"generic"}]}' +# Older releases dump the models as a list. +write_meta "$sessions/session_2" "$now_iso" 20 0 0 '{"active_model":"devstral","models":[{"name":"devstral-2","provider":"mistral","alias":"devstral"}],"providers":[{"name":"mistral","backend":"mistral"}]}' +# A generic provider aimed at Mistral's API is still Mistral. +write_meta "$sessions/session_3" "$now_iso" 40 0 0 '{"active_model":"codestral","models":{"codestral":{"name":"codestral-latest","provider":"la-plateforme","alias":"codestral"}},"providers":[{"name":"la-plateforme","api_base":"https://codestral.mistral.ai/v1","backend":"generic"}]}' +# A lookalike host is not. +write_meta "$sessions/session_4" "$now_iso" 80 0 0 '{"active_model":"x","models":{"x":{"name":"x","provider":"proxy","alias":"x"}},"providers":[{"name":"proxy","api_base":"https://mistral.ai.example.com/v1","backend":"generic"}]}' +record=$(collect "$home" --force) +[[ $(jq -c '[.modelUsage | to_entries[] | {(.key): .value.inputTokens}] | add' <<<"$record") == '{"mistral-medium-3.5":10,"devstral":20,"codestral":40}' ]] || + fail "Vibe collector resolves the session model and its provider as Vibe does" "$record" +pass "Vibe collector resolves the session model and its provider as Vibe does" + +# ---------------------------------------------------------------- untrusted fields + +home=$(new_home untrusted) +sessions="$home/.vibe/logs/session" +write_meta "$sessions/session_bad_time" "not a time" 500 0 500 "$MISTRAL_CONFIG" +write_meta "$sessions/session_future" "2999-01-01T00:00:00+00:00" 500 0 500 "$MISTRAL_CONFIG" +write_meta "$sessions/session_odd_tokens" "$now_iso" '"700"' true NaN "$MISTRAL_CONFIG" +write_meta "$sessions/session_negative" "$now_iso" -5 0 7 "$MISTRAL_CONFIG" +mkdir -p "$sessions/session_not_json" && printf '[1, 2' >"$sessions/session_not_json/meta.json" +record=$(collect "$home" --force 2>/dev/null) +[[ $(jq -c '{ready, todayTotalTokens, totalSessions}' <<<"$record") == '{"ready":true,"todayTotalTokens":7,"totalSessions":1}' ]] || + fail "Vibe collector skips bad times and bad token fields without losing the record" "$record" +pass "Vibe collector skips bad times and bad token fields without losing the record" + +# A meta.json that fails to parse leaves the scan incomplete, so it isn't cached. +[[ -z $(find "$home/.cache/omarchy/agent-usage" -name 'vibe-scan-*.json' 2>/dev/null) ]] || + fail "Vibe collector does not cache a scan that could not read every session" +pass "Vibe collector does not cache a scan that could not read every session" + +# ---------------------------------------------------------------- plan + +home=$(new_home plan) +mkdir -p "$home/.vibe" +cat >"$home/.vibe/whoami_cache.json" <<'EOF' +{"old":{"stored_at_timestamp":100,"payload":{"plan_type":"api","plan_name":"FREE_TRIAL"}},"new":{"stored_at_timestamp":200,"payload":{"plan_type":"chat","plan_name":"individual"}}} +EOF +record=$(collect "$home" --force) +[[ $(jq -c '{ready, tierLabel}' <<<"$record") == '{"ready":true,"tierLabel":"Pro"}' ]] || + fail "Vibe collector names the newest cached plan as Vibe does" "$record" +pass "Vibe collector names the newest cached plan as Vibe does" + +printf '%s\n' '{"a":{"stored_at_timestamp":1,"payload":{"plan_type":"chat","plan_name":"MYSTERY"}}}' >"$home/.vibe/whoami_cache.json" +record=$(collect "$home" --force) +[[ $(jq -c '{ready, tierLabel}' <<<"$record") == '{"ready":false,"tierLabel":""}' ]] || + fail "Vibe collector leaves a plan Vibe cannot name unnamed" "$record" +pass "Vibe collector leaves a plan Vibe cannot name unnamed" + +# ---------------------------------------------------------------- VIBE_HOME and save_dir + +home=$(new_home relocated) +mkdir -p "$home/vibe-home" +cat >"$home/vibe-home/config.toml" <"$home/.pi/agent/sessions/p/a.jsonl" <"$home/.omp/profiles/work/agent/sessions/p/fork.jsonl" <"$home/custom-pi/sessions/p/a.jsonl" +# A $HOME that is a git checkout must not hide session files from rg. +printf '*\n' >"$home/.gitignore" +git -C "$home" init -q +record=$(HOME="$home" XDG_CACHE_HOME="$home/.cache" PI_CODING_AGENT_DIR="$home/custom-pi" "$COLLECTOR" --force) +[[ $(jq -r '.todayTotalTokens' <<<"$record") == "9" ]] || + fail "Vibe collector honors PI_CODING_AGENT_DIR and ignores git ignore rules" "$record" +pass "Vibe collector honors PI_CODING_AGENT_DIR and ignores git ignore rules" + +# ---------------------------------------------------------------- scan cache + +home=$(new_home cache) +sessions="$home/.vibe/logs/session" +write_meta "$sessions/session_1" "$now_iso" 10 0 0 "$MISTRAL_CONFIG" +first=$(collect "$home") +write_meta "$sessions/session_2" "$now_iso" 5 0 0 "$MISTRAL_CONFIG" +[[ $(collect "$home" | jq -r '.todayTotalTokens') == "10" ]] || + fail "Vibe collector reuses a scan from the last few seconds" +[[ $(collect "$home" --limits-only | jq -r '.todayTotalTokens') == "10" ]] || + fail "Vibe collector reuses a recent scan for --limits-only" +[[ $(collect "$home" --force | jq -r '.todayTotalTokens') == "15" ]] || + fail "Vibe collector rescans with --force" +pass "Vibe collector reuses a recent scan and rescans with --force" + +cache_file=$(find "$home/.cache/omarchy/agent-usage" -name 'vibe-scan-*.json' -print -quit) +[[ $(jq -r '.scanDate' "$cache_file") == "$today" ]] || + fail "Vibe collector stamps its scan cache with the day it was made" "$(cat "$cache_file")" +jq '.scanDate = "2000-01-01" | .stats.todayTotalTokens = 1' "$cache_file" >"$cache_file.new" && mv "$cache_file.new" "$cache_file" +[[ $(collect "$home" | jq -r '.todayTotalTokens') == "15" ]] || + fail "Vibe collector ignores a scan cache from another day" +pass "Vibe collector ignores a scan cache from another day" + +# The per-session index is only an optimization: a changed session is reread. +write_meta "$sessions/session_1" "$now_iso" 50 0 0 "$MISTRAL_CONFIG" +[[ $(collect "$home" --force | jq -r '.todayTotalTokens') == "55" ]] || + fail "Vibe collector rereads a session whose meta.json changed" +pass "Vibe collector rereads a session whose meta.json changed" From 55aa5cc266ecb0af462627a02635871a94937a38 Mon Sep 17 00:00:00 2001 From: kenth-dev Date: Wed, 14 Oct 2026 23:18:30 -0400 Subject: [PATCH 115/546] Add a Kiro CLI usage collector Read the session files Kiro CLI writes under ~/.kiro/sessions/cli/ into the agents-panel record: each user turn's model, request count, and the credits its metering_usage reports. No rate-limit endpoint exists for Kiro subscriptions, so the record carries local stats only. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-kiro | 262 +++++++++++++++++++++++++++++++++++ 1 file changed, 262 insertions(+) create mode 100755 bin/omarchy-agent-usage-kiro diff --git a/bin/omarchy-agent-usage-kiro b/bin/omarchy-agent-usage-kiro new file mode 100755 index 00000000000..9d676cf8f71 --- /dev/null +++ b/bin/omarchy-agent-usage-kiro @@ -0,0 +1,262 @@ +#!/usr/bin/python3 +# omarchy:summary=Print the Kiro usage record as JSON +# omarchy:args=[--force] +# omarchy:hidden=true +"""Collect Kiro CLI usage into one display-ready JSON record. + +Local stats come from Kiro CLI session files in ~/.kiro/sessions/cli/. +Each session JSON contains per-turn metering_usage (credits consumed), +model information, and timestamps. Kiro does not expose per-request token +counts, so the record reports credits as the primary metric and uses +request counts as the messageCount equivalent for daily charts. + +No rate-limit endpoint exists for Kiro subscriptions, so the record +carries local stats only. +""" + +from __future__ import annotations + +import argparse +import datetime as dt +import fcntl +import hashlib +import json +import os +import sys +import tempfile +import time +from pathlib import Path +from typing import Any + +AGENT_ID = "kiro" +AGENT_NAME = "Kiro" + +SCAN_REUSE_SECONDS = 20 + + +def sessions_dir() -> Path: + return Path(os.environ.get("KIRO_HOME") or (Path.home() / ".kiro")) / "sessions" / "cli" + + +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 number(value: Any) -> float: + try: + return float(value or 0) + except Exception: + return 0.0 + + +def local_date_from_iso(value: str | None) -> str: + """Parse an ISO-8601 timestamp to a local YYYY-MM-DD date string.""" + if not value: + return dt.datetime.now().strftime("%Y-%m-%d") + try: + parsed = dt.datetime.fromisoformat(value.replace("Z", "+00:00")) + if parsed.tzinfo is not None: + parsed = parsed.astimezone() + return parsed.strftime("%Y-%m-%d") + except Exception: + return dt.datetime.now().strftime("%Y-%m-%d") + + +def recent_date_strings() -> list[str]: + today = dt.datetime.now().date() + return [(today - dt.timedelta(days=offset)).strftime("%Y-%m-%d") for offset in range(6, -1, -1)] + + +def empty_bucket() -> dict[str, float]: + return { + "credits": 0.0, + "requests": 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: + 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") + tmp.chmod(0o644) + tmp.replace(path) + except BaseException: + tmp.unlink(missing_ok=True) + raise + + +def scan_sessions(sessions_path: Path) -> dict[str, Any]: + today = dt.datetime.now().strftime("%Y-%m-%d") + 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_credits_by_model: dict[str, float] = {} + credits_by_model: dict[str, dict[str, float]] = {} + total_prompts = 0 + today_prompt_count = 0 + today_total_credits = 0.0 + total_credits = 0.0 + + if not sessions_path.is_dir(): + return _empty_stats(recent_dates) + + for path in sessions_path.glob("*.json"): + try: + data = json.loads(path.read_text(encoding="utf-8")) + except Exception: + continue + + session_id = data.get("session_id") or path.stem + created_at = data.get("created_at") + state = data.get("session_state") + if not isinstance(state, dict): + continue + meta = state.get("conversation_metadata") + if not isinstance(meta, dict): + continue + turns = meta.get("user_turn_metadatas") + if not isinstance(turns, list): + continue + + for turn in turns: + if not isinstance(turn, dict): + continue + + # Determine the day: prefer end_timestamp, fall back to created_at + timestamp = turn.get("end_timestamp") or created_at + day = local_date_from_iso(timestamp) + + model = str(turn.get("model") or "auto") + request_count = int(number(turn.get("total_request_count")) or 1) + + # Sum credits from metering_usage array + metering = turn.get("metering_usage") + turn_credits = 0.0 + if isinstance(metering, list): + for entry in metering: + if isinstance(entry, dict): + turn_credits += number(entry.get("value")) + + if turn_credits <= 0 and request_count <= 0: + continue + + total_prompts += 1 + sessions.add(session_id) + active_days.add(day) + total_credits += turn_credits + + bucket = credits_by_model.setdefault(model, empty_bucket()) + bucket["credits"] += turn_credits + bucket["requests"] += request_count + + if day in recent: + # Use request count as the daily chart metric (analogous to + # messageCount in other collectors) + recent[day]["messageCount"] += request_count + + if day == today: + today_prompt_count += 1 + today_sessions.add(session_id) + today_total_credits += turn_credits + today_credits_by_model[model] = today_credits_by_model.get(model, 0.0) + turn_credits + + return { + "todayPrompts": today_prompt_count, + "todaySessions": len(today_sessions), + "todayTotalCredits": round(today_total_credits, 6), + "todayCreditsByModel": {k: round(v, 6) for k, v in today_credits_by_model.items()}, + "recentDays": [recent[day] for day in recent_dates], + "modelUsage": { + model: {"credits": round(b["credits"], 6), "requests": int(b["requests"])} + for model, b in credits_by_model.items() + }, + "totalPrompts": total_prompts, + "totalSessions": len(sessions), + "totalCredits": round(total_credits, 6), + "activeDays": len(active_days), + "activeDates": sorted(active_days), + } + + +def _empty_stats(recent_dates: list[str]) -> dict[str, Any]: + return { + "todayPrompts": 0, + "todaySessions": 0, + "todayTotalCredits": 0.0, + "todayCreditsByModel": {}, + "recentDays": [{"date": day, "messageCount": 0} for day in recent_dates], + "modelUsage": {}, + "totalPrompts": 0, + "totalSessions": 0, + "totalCredits": 0.0, + "activeDays": 0, + "activeDates": [], + } + + +def cached_scan(sessions_path: Path, max_age_seconds: float) -> dict[str, Any]: + digest = hashlib.sha1(str(sessions_path).encode("utf-8")).hexdigest()[:16] + cache_file = cache_root() / f"kiro-scan-{digest}.json" + lock_file = cache_root() / f"kiro-scan-{digest}.lock" + + cached = read_fresh_json(cache_file, max_age_seconds) + if cached is not None: + return cached + + with lock_file.open("w") as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + cached = read_fresh_json(cache_file, max_age_seconds) + if cached is not None: + return cached + summary = scan_sessions(sessions_path) + write_json(cache_file, summary) + return summary + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--force", action="store_true", help="rescan sessions ignoring caches") + parser.add_argument("--limits-only", action="store_true", help="reuse recent scan (no limits to fetch for Kiro)") + args = parser.parse_args() + + scan_age = 0 if args.force else SCAN_REUSE_SECONDS + stats = cached_scan(sessions_dir(), scan_age) + + record: dict[str, Any] = { + "schemaVersion": 1, + "id": AGENT_ID, + "name": AGENT_NAME, + "updatedAt": dt.datetime.now(dt.timezone.utc).isoformat(), + "ready": stats.get("totalPrompts", 0) > 0, + "hasLocalStats": True, + "hasPromptStats": True, + "tierLabel": "", + "usageStatusText": "", + "authHelpText": "", + "limits": [], + } + record.update(stats) + print(json.dumps(record, separators=(",", ":"), sort_keys=True)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) From 770861b81a10f5e2c583be7296d38cfb2c4d09df Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:21:05 -0400 Subject: [PATCH 116/546] Bring the Kiro collector to the record contract Kiro's session files do carry the service's token counts per user turn (input, output, cache read, cache write), so those fill the standard model buckets and daily token totals; a turn stored with zeros, as Kiro's Auto agent does, still counts as a prompt with no tokens, and nothing is estimated. The credit fields that replaced the contract's token fields are gone, and recentDays now holds tokens rather than request counts. The scan cache is stamped with its day and schema, guarded by 0 <= age < N, skipped for --force, reused longer for --limits-only, and not written when a session file couldn't be read; a per-file index rereads only sessions that changed. Untrusted times and counts are rejected field by field. The unverified KIRO_HOME override is dropped. Kiro keeps neither its monthly credit allowance nor the month's use on disk, so there is no meter. Co-authored-by: kenth-dev Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-kiro | 560 ++++++++++++++++---------- test/shell.d/agent-usage-kiro-test.sh | 128 ++++++ 2 files changed, 478 insertions(+), 210 deletions(-) create mode 100755 test/shell.d/agent-usage-kiro-test.sh diff --git a/bin/omarchy-agent-usage-kiro b/bin/omarchy-agent-usage-kiro index 9d676cf8f71..957db1bc54d 100755 --- a/bin/omarchy-agent-usage-kiro +++ b/bin/omarchy-agent-usage-kiro @@ -1,23 +1,21 @@ #!/usr/bin/python3 -# omarchy:summary=Print the Kiro usage record as JSON -# omarchy:args=[--force] +# omarchy:summary=Print the Kiro CLI usage record as JSON +# omarchy:args=[--force] [--limits-only] # omarchy:hidden=true """Collect Kiro CLI usage into one display-ready JSON record. -Local stats come from Kiro CLI session files in ~/.kiro/sessions/cli/. -Each session JSON contains per-turn metering_usage (credits consumed), -model information, and timestamps. Kiro does not expose per-request token -counts, so the record reports credits as the primary metric and uses -request counts as the messageCount equivalent for daily charts. - -No rate-limit endpoint exists for Kiro subscriptions, so the record -carries local stats only. +Local stats come from the session files Kiro CLI writes under +~/.kiro/sessions/cli/: one JSON document per session whose +conversation_metadata keeps a row per user turn, with the model, the time the +turn ended, and the token counts the service reported. Those counts are often +zero (Kiro's Auto agent stores zeros), and such a turn still counts as a +prompt with no tokens; nothing is estimated from text length. Kiro bills in +credits against a monthly allowance, but neither the allowance nor the +month's use is kept on disk, so the record carries no meter. The agents panel +only ever reads the JSON this prints. """ -from __future__ import annotations - import argparse -import datetime as dt import fcntl import hashlib import json @@ -25,238 +23,380 @@ import os import sys import tempfile import time +from datetime import datetime, timedelta, timezone from pathlib import Path -from typing import Any AGENT_ID = "kiro" AGENT_NAME = "Kiro" +# 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). --limits-only promises only fresh limits, and Kiro has none to read, +# so it may reuse a scan for 15 minutes. SCAN_REUSE_SECONDS = 20 +LIMITS_ONLY_REUSE_SECONDS = 900 + +CACHE_SCHEMA_VERSION = 1 +INDEX_VERSION = 1 +# Counts and times read from disk are untrusted. Anything past these bounds is +# a corrupt or crafted field, not usage. +MAX_TOKENS = 10**12 +MIN_EPOCH_SECONDS = 946_684_800 # 2000-01-01 +MAX_MODEL_NAME_LEN = 120 -def sessions_dir() -> Path: - return Path(os.environ.get("KIRO_HOME") or (Path.home() / ".kiro")) / "sessions" / "cli" +def token_count(value): + """A token field as a whole number, or 0 when it isn't a plausible count. -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 + bool is an int subclass, and a string, NaN, negative, or absurd value is a + bad field: it counts as nothing rather than skewing the totals. + """ + if isinstance(value, bool) or not isinstance(value, (int, float)): + return 0 + if value != value or value <= 0 or value > MAX_TOKENS: + return 0 + return int(value) -def number(value: Any) -> float: +def epoch_seconds(value): + """Seconds since the epoch for an ISO string or a seconds/milliseconds + number, or None when the value can't be a real moment of use.""" + if value is None or isinstance(value, bool): + return None + if isinstance(value, (int, float)): + seconds = float(value) + elif isinstance(value, str) and value.strip(): + text = value.strip() try: - return float(value or 0) - except Exception: - return 0.0 + seconds = float(text) + except ValueError: + try: + parsed = datetime.fromisoformat(text[:-1] + "+00:00" if text.endswith("Z") else text) + except ValueError: + return None + if parsed.tzinfo is None: + parsed = parsed.astimezone() + seconds = parsed.timestamp() + else: + return None + if seconds != seconds: + return None + if seconds > 10_000_000_000: + seconds /= 1000 + # A day of slack past now covers a clock a little behind the writer's. + if seconds < MIN_EPOCH_SECONDS or seconds > time.time() + 86_400: + return None + return seconds -def local_date_from_iso(value: str | None) -> str: - """Parse an ISO-8601 timestamp to a local YYYY-MM-DD date string.""" - if not value: - return dt.datetime.now().strftime("%Y-%m-%d") - try: - parsed = dt.datetime.fromisoformat(value.replace("Z", "+00:00")) - if parsed.tzinfo is not None: - parsed = parsed.astimezone() - return parsed.strftime("%Y-%m-%d") - except Exception: - return dt.datetime.now().strftime("%Y-%m-%d") +def local_day(seconds): + return datetime.fromtimestamp(seconds).strftime("%Y-%m-%d") -def recent_date_strings() -> list[str]: - today = dt.datetime.now().date() - return [(today - dt.timedelta(days=offset)).strftime("%Y-%m-%d") for offset in range(6, -1, -1)] +def obj(value): + return value if isinstance(value, dict) else {} -def empty_bucket() -> dict[str, float]: +def model_label(raw): + """A model id fit for the record, or None when it is too long to be one.""" + value = str(raw or "").strip() + if len(value) > MAX_MODEL_NAME_LEN: + return None + return value or "auto" + + +def empty_bucket(): + return { + "inputTokens": 0, + "outputTokens": 0, + "cacheReadInputTokens": 0, + "cacheCreationInputTokens": 0, + } + + +class Stats: + def __init__(self): + 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() + self.today_tokens_by_model = {} + self.total_prompts = 0 + self.total_sessions = set() + self.active_days = set() + self.model_usage = {} + + def add(self, session, day, model, input_tokens, output_tokens, cache_read, cache_write): + """One user turn: a prompt, with whatever tokens Kiro recorded for it.""" + total = input_tokens + output_tokens + cache_read + cache_write + self.total_prompts += 1 + self.total_sessions.add(session) + self.active_days.add(day) + if total > 0: + bucket = self.model_usage.setdefault(model, empty_bucket()) + bucket["inputTokens"] += input_tokens + bucket["outputTokens"] += output_tokens + bucket["cacheReadInputTokens"] += cache_read + bucket["cacheCreationInputTokens"] += cache_write + if day in self.recent: + # recentDays.messageCount holds tokens, despite the shared legacy name. + self.recent[day] += total + if day == self.today: + self.today_prompts += 1 + self.today_sessions.add(session) + self.today_total_tokens += total + if total > 0: + self.today_tokens_by_model[model] = self.today_tokens_by_model.get(model, 0) + total + + def record(self): return { - "credits": 0.0, - "requests": 0, + "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), + # The dates travel too: merging snapshots from several machines needs + # their union, which a count alone cannot give. + "activeDays": len(self.active_days), + "activeDates": sorted(self.active_days), + "modelUsage": self.model_usage, } -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 +def sessions_dir(): + return Path.home() / ".kiro" / "sessions" / "cli" + + +def parse_session(path): + """One session file's index entry: its id and a row per user turn, + [endSeconds, model, input, output, cacheRead, cacheWrite]. + + A turn that never ran (no request, no tokens) is no prompt. A turn whose + end time is missing or unreadable falls back to the session's creation + time, and is skipped when that is unreadable too. + """ + data = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(data, dict): + return {"session": path.stem, "rows": []} + state = obj(data.get("session_state")) + turns = obj(state.get("conversation_metadata")).get("user_turn_metadatas") + session_model = obj(obj(state.get("rts_model_state")).get("model_info")).get("model_id") + created = epoch_seconds(data.get("created_at")) + rows = [] + for turn in turns if isinstance(turns, list) else []: + if not isinstance(turn, dict): + continue + input_tokens = token_count(turn.get("input_token_count")) + output_tokens = token_count(turn.get("output_token_count")) + cache_read = token_count(turn.get("cache_read_input_token_count")) + cache_write = token_count(turn.get("cache_write_input_token_count")) + requests = token_count(turn.get("total_request_count")) + if not (requests or input_tokens or output_tokens or cache_read or cache_write): + continue + ended = epoch_seconds(turn.get("end_timestamp")) + if ended is None: + ended = created + # A turn without a model of its own ran on the session's. + model = turn.get("model") + if not isinstance(model, str) or not model.strip(): + model = session_model if isinstance(session_model, str) else "" + model = model_label(model) + if ended is None or model is None: + continue + rows.append([ended, model, input_tokens, output_tokens, cache_read, cache_write]) + session = data.get("session_id") + return {"session": str(session) if isinstance(session, str) and session else path.stem, "rows": rows} + + +def load_index(path): + try: + data = json.loads(path.read_text(encoding="utf-8")) + except Exception: + return {} + if isinstance(data, dict) and data.get("version") == INDEX_VERSION and isinstance(data.get("files"), dict): + return data["files"] + return {} + + +def scan_sessions(stats, index_file): + """Add every session's turns. Kiro rewrites a session file whole on every + turn, so the index keeps each file's parsed turns by size and mtime and + only rereads files that changed. Returns whether every file was read.""" + root = sessions_dir() + previous = load_index(index_file) if index_file else {} + current = {} + complete = True + try: + paths = sorted(root.glob("*.json")) if root.is_dir() else [] + except OSError: + return False + for path in paths: + key = str(path) try: - if time.time() - path.stat().st_mtime <= max_age_seconds: - return json.loads(path.read_text(encoding="utf-8")) - except Exception: - return None + stat = path.stat() + signature = [stat.st_size, stat.st_mtime_ns] + held = previous.get(key) + if isinstance(held, dict) and held.get("sig") == signature: + entry = held + else: + entry = dict(parse_session(path), sig=signature) + except (OSError, ValueError, UnicodeDecodeError) as exc: + # Mid-write or unreadable: skip it this run, uncached. + print(f"omarchy-agent-usage-kiro: skipping {path}: {type(exc).__name__}", file=sys.stderr) + complete = False + continue + current[key] = entry + for row in entry.get("rows") or []: + try: + ended, model, input_tokens, output_tokens, cache_read, cache_write = row + day = local_day(float(ended)) + counts = [int(input_tokens), int(output_tokens), int(cache_read), int(cache_write)] + except (TypeError, ValueError, OverflowError, OSError): + continue + stats.add(str(entry.get("session") or key), day, str(model), *counts) + + if index_file and current != previous: + try: + write_json(index_file, {"version": INDEX_VERSION, "files": current}) + except Exception as exc: + print(f"omarchy-agent-usage-kiro: could not write session index ({exc})", file=sys.stderr) + return complete + + +# ---------------------------------------------------------------- cache + + +def cache_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 cache_paths(): + digest = hashlib.sha1(str(sessions_dir()).encode("utf-8")).hexdigest()[:16] + root = cache_root() + return root / f"kiro-scan-{digest}.json", root / f"kiro-scan-{digest}.lock", root / f"kiro-index-{digest}.json" + + +def read_fresh_json(path, max_age_seconds): + if max_age_seconds <= 0: + 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: dict[str, Any]) -> 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=(",", ":"), sort_keys=True) + "\n") - tmp.chmod(0o644) - tmp.replace(path) - except BaseException: - tmp.unlink(missing_ok=True) - raise - - -def scan_sessions(sessions_path: Path) -> dict[str, Any]: - today = dt.datetime.now().strftime("%Y-%m-%d") - 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_credits_by_model: dict[str, float] = {} - credits_by_model: dict[str, dict[str, float]] = {} - total_prompts = 0 - today_prompt_count = 0 - today_total_credits = 0.0 - total_credits = 0.0 - - if not sessions_path.is_dir(): - return _empty_stats(recent_dates) - - for path in sessions_path.glob("*.json"): - try: - data = json.loads(path.read_text(encoding="utf-8")) - except Exception: - continue - - session_id = data.get("session_id") or path.stem - created_at = data.get("created_at") - state = data.get("session_state") - if not isinstance(state, dict): - continue - meta = state.get("conversation_metadata") - if not isinstance(meta, dict): - continue - turns = meta.get("user_turn_metadatas") - if not isinstance(turns, list): - continue - - for turn in turns: - if not isinstance(turn, dict): - continue - - # Determine the day: prefer end_timestamp, fall back to created_at - timestamp = turn.get("end_timestamp") or created_at - day = local_date_from_iso(timestamp) - - model = str(turn.get("model") or "auto") - request_count = int(number(turn.get("total_request_count")) or 1) - - # Sum credits from metering_usage array - metering = turn.get("metering_usage") - turn_credits = 0.0 - if isinstance(metering, list): - for entry in metering: - if isinstance(entry, dict): - turn_credits += number(entry.get("value")) - - if turn_credits <= 0 and request_count <= 0: - continue - - total_prompts += 1 - sessions.add(session_id) - active_days.add(day) - total_credits += turn_credits - - bucket = credits_by_model.setdefault(model, empty_bucket()) - bucket["credits"] += turn_credits - bucket["requests"] += request_count - - if day in recent: - # Use request count as the daily chart metric (analogous to - # messageCount in other collectors) - recent[day]["messageCount"] += request_count - - if day == today: - today_prompt_count += 1 - today_sessions.add(session_id) - today_total_credits += turn_credits - today_credits_by_model[model] = today_credits_by_model.get(model, 0.0) + turn_credits +def write_json(path, payload): + # A temp name unique to this writer: 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 - return { - "todayPrompts": today_prompt_count, - "todaySessions": len(today_sessions), - "todayTotalCredits": round(today_total_credits, 6), - "todayCreditsByModel": {k: round(v, 6) for k, v in today_credits_by_model.items()}, - "recentDays": [recent[day] for day in recent_dates], - "modelUsage": { - model: {"credits": round(b["credits"], 6), "requests": int(b["requests"])} - for model, b in credits_by_model.items() - }, - "totalPrompts": total_prompts, - "totalSessions": len(sessions), - "totalCredits": round(total_credits, 6), - "activeDays": len(active_days), - "activeDates": sorted(active_days), - } +STATS_KEYS = ("todayPrompts", "todaySessions", "todayTotalTokens", "todayTokensByModel", "recentDays", "totalPrompts", "totalSessions", "activeDays", "activeDates", "modelUsage") -def _empty_stats(recent_dates: list[str]) -> dict[str, Any]: - return { - "todayPrompts": 0, - "todaySessions": 0, - "todayTotalCredits": 0.0, - "todayCreditsByModel": {}, - "recentDays": [{"date": day, "messageCount": 0} for day in recent_dates], - "modelUsage": {}, - "totalPrompts": 0, - "totalSessions": 0, - "totalCredits": 0.0, - "activeDays": 0, - "activeDates": [], - } + +def read_cached_stats(cache_file, max_age_seconds): + cached = read_fresh_json(cache_file, max_age_seconds) + if not isinstance(cached, dict) or cached.get("schemaVersion") != CACHE_SCHEMA_VERSION: + return None + # today* fields only mean "today" on the day they were scanned. + if cached.get("scanDate") != datetime.now().strftime("%Y-%m-%d"): + return None + stats = cached.get("stats") + if not isinstance(stats, dict) or not all(key in stats for key in STATS_KEYS): + return None + return stats -def cached_scan(sessions_path: Path, max_age_seconds: float) -> dict[str, Any]: - digest = hashlib.sha1(str(sessions_path).encode("utf-8")).hexdigest()[:16] - cache_file = cache_root() / f"kiro-scan-{digest}.json" - lock_file = cache_root() / f"kiro-scan-{digest}.lock" +def run_scan(index_file): + stats = Stats() + complete = scan_sessions(stats, index_file) + return stats.record(), complete - cached = read_fresh_json(cache_file, max_age_seconds) - if cached is not None: - return cached +def cached_local_stats(max_age): + """Local stats, with the cache as a pure optimization: any cache failure + degrades to a direct scan and a warning on stderr.""" + try: + cache_file, lock_file, index_file = cache_paths() + cached = read_cached_stats(cache_file, max_age) + if cached is not None: + return cached with lock_file.open("w") as lock: - fcntl.flock(lock, fcntl.LOCK_EX) - cached = read_fresh_json(cache_file, max_age_seconds) - if cached is not None: - return cached - summary = scan_sessions(sessions_path) - write_json(cache_file, summary) - return summary - - -def main() -> int: - parser = argparse.ArgumentParser() - parser.add_argument("--force", action="store_true", help="rescan sessions ignoring caches") - parser.add_argument("--limits-only", action="store_true", help="reuse recent scan (no limits to fetch for Kiro)") - args = parser.parse_args() - - scan_age = 0 if args.force else SCAN_REUSE_SECONDS - stats = cached_scan(sessions_dir(), scan_age) - - record: dict[str, Any] = { - "schemaVersion": 1, - "id": AGENT_ID, - "name": AGENT_NAME, - "updatedAt": dt.datetime.now(dt.timezone.utc).isoformat(), - "ready": stats.get("totalPrompts", 0) > 0, - "hasLocalStats": True, - "hasPromptStats": True, - "tierLabel": "", - "usageStatusText": "", - "authHelpText": "", - "limits": [], - } - record.update(stats) - print(json.dumps(record, separators=(",", ":"), sort_keys=True)) - return 0 + fcntl.flock(lock, fcntl.LOCK_EX) + cached = read_cached_stats(cache_file, max_age) + if cached is not None: + return cached + stats, complete = run_scan(index_file) + # An interrupted scan still serves this run, but caching it would hide + # the missing usage from every reader until the cache expires. + if complete: + try: + write_json(cache_file, {"schemaVersion": CACHE_SCHEMA_VERSION, "scanDate": datetime.now().strftime("%Y-%m-%d"), "stats": stats}) + except Exception as exc: + print(f"omarchy-agent-usage-kiro: could not write usage cache ({exc})", file=sys.stderr) + return stats + except Exception as exc: + print(f"omarchy-agent-usage-kiro: cache unavailable ({exc}); scanning directly", file=sys.stderr) + stats, _ = run_scan(None) + return stats + + +def main(): + parser = argparse.ArgumentParser(description="Print the Kiro CLI usage record as JSON") + # --force rescans everything. --limits-only only promises fresh limits, and + # Kiro keeps none on disk, so it reuses a scan far longer than a normal run. + parser.add_argument("--force", action="store_true") + parser.add_argument("--limits-only", action="store_true") + args = parser.parse_args() + + max_age = 0 if args.force else (LIMITS_ONLY_REUSE_SECONDS if args.limits_only else SCAN_REUSE_SECONDS) + stats = cached_local_stats(max_age) + + # There is no allowance to read, so nothing waits on a sign-in: the status + # stays empty and local stats show alone. A machine that never ran Kiro + # prints a record the panel skips. + record = { + "schemaVersion": 1, + "id": AGENT_ID, + "name": AGENT_NAME, + "updatedAt": datetime.now(timezone.utc).isoformat(), + "ready": stats["totalPrompts"] > 0, + "hasLocalStats": True, + "hasPromptStats": True, + "tierLabel": "", + "usageStatusText": "", + "authHelpText": "", + "limits": [], + "limitsStale": False, + "limitsFetchedAt": 0, + } + record.update(stats) + print(json.dumps(record, separators=(",", ":"))) + return 0 if __name__ == "__main__": - raise SystemExit(main()) + raise SystemExit(main()) diff --git a/test/shell.d/agent-usage-kiro-test.sh b/test/shell.d/agent-usage-kiro-test.sh new file mode 100755 index 00000000000..a945785d644 --- /dev/null +++ b/test/shell.d/agent-usage-kiro-test.sh @@ -0,0 +1,128 @@ +#!/bin/bash + +source "$(dirname "$0")/base-test.sh" + +require_command jq +require_command python3 + +COLLECTOR="$ROOT/bin/omarchy-agent-usage-kiro" +TMP_ROOT=$(mktemp -d) +trap 'rm -rf "$TMP_ROOT"' EXIT + +unset XDG_CACHE_HOME + +today=$(date +%Y-%m-%d) +now_iso=$(date -u +%Y-%m-%dT%H:%M:%SZ) +two_days_ago=$(date -d '2 days ago' +%Y-%m-%d) +two_days_ago_iso=$(date -u -d '2 days ago' +%Y-%m-%dT%H:%M:%SZ) + +new_home() { + local home="$TMP_ROOT/$1" + mkdir -p "$home/.kiro/sessions/cli" + printf '%s\n' "$home" +} + +collect() { + local home="$1" + shift + HOME="$home" XDG_CACHE_HOME="$home/.cache" "$COLLECTOR" "$@" +} + +# write_session +write_session() { + cat >"$1/.kiro/sessions/cli/$2.json" <"$home/.kiro/sessions/cli/torn.json" +record=$(collect "$home" --force 2>/dev/null) +# An unreadable or future end time falls back to the session's creation time; +# a turn whose every count is bad never ran. +[[ $(jq -c '{todayPrompts, todayTotalTokens}' <<<"$record") == '{"todayPrompts":3,"todayTotalTokens":27}' ]] || + fail "Kiro collector skips bad times and bad token fields without losing the record" "$record" +pass "Kiro collector skips bad times and bad token fields without losing the record" + +[[ -z $(find "$home/.cache/omarchy/agent-usage" -name 'kiro-scan-*.json' 2>/dev/null) ]] || + fail "Kiro collector does not cache a scan that could not read every session" +pass "Kiro collector does not cache a scan that could not read every session" + +# ---------------------------------------------------------------- scan cache + +home=$(new_home cache) +write_session "$home" s1 "[{\"end_timestamp\":\"$now_iso\",\"total_request_count\":1,\"input_token_count\":10}]" +collect "$home" >/dev/null +write_session "$home" s2 "[{\"end_timestamp\":\"$now_iso\",\"total_request_count\":1,\"input_token_count\":5}]" +[[ $(collect "$home" | jq -r '.todayTotalTokens') == "10" ]] || + fail "Kiro collector reuses a scan from the last few seconds" +[[ $(collect "$home" --limits-only | jq -r '.todayTotalTokens') == "10" ]] || + fail "Kiro collector reuses a recent scan for --limits-only" +[[ $(collect "$home" --force | jq -r '.todayTotalTokens') == "15" ]] || + fail "Kiro collector rescans with --force" +pass "Kiro collector reuses a recent scan and rescans with --force" + +cache_file=$(find "$home/.cache/omarchy/agent-usage" -name 'kiro-scan-*.json' -print -quit) +[[ $(jq -r '.scanDate' "$cache_file") == "$today" ]] || + fail "Kiro collector stamps its scan cache with the day it was made" "$(cat "$cache_file")" +jq '.scanDate = "2000-01-01" | .stats.todayTotalTokens = 1' "$cache_file" >"$cache_file.new" && mv "$cache_file.new" "$cache_file" +[[ $(collect "$home" | jq -r '.todayTotalTokens') == "15" ]] || + fail "Kiro collector ignores a scan cache from another day" +pass "Kiro collector ignores a scan cache from another day" + +# The cache stamped in the future is distrusted rather than kept forever. +jq '.stats.todayTotalTokens = 1' "$cache_file" >"$cache_file.new" && mv "$cache_file.new" "$cache_file" +touch -d '+1 hour' "$cache_file" +[[ $(collect "$home" | jq -r '.todayTotalTokens') == "15" ]] || + fail "Kiro collector distrusts a scan cache dated in the future" +pass "Kiro collector distrusts a scan cache dated in the future" + +# Kiro rewrites a session file whole each turn; the index rereads it. +write_session "$home" s1 "[{\"end_timestamp\":\"$now_iso\",\"total_request_count\":1,\"input_token_count\":10},{\"end_timestamp\":\"$now_iso\",\"total_request_count\":1,\"input_token_count\":20}]" +[[ $(collect "$home" --force | jq -r '"\(.todayPrompts):\(.todayTotalTokens)"') == "3:35" ]] || + fail "Kiro collector rereads a session file that changed" +pass "Kiro collector rereads a session file that changed" From c3eb4341f93364c642e6f5b3093e687fb0d1f4e4 Mon Sep 17 00:00:00 2001 From: David Heinemeier Hansson Date: Wed, 14 Oct 2026 23:21:10 -0400 Subject: [PATCH 117/546] Add the Kiro mark to the agents panel The ghost mark from Kiro's own KiroCrew repository (Apache-2.0), in white with a dark twin for light surfaces. Co-Authored-By: Claude Opus 5.5 --- shell/plugins/agents/assets/kiro-light.svg | 1 + shell/plugins/agents/assets/kiro.svg | 1 + 2 files changed, 2 insertions(+) create mode 100644 shell/plugins/agents/assets/kiro-light.svg create mode 100644 shell/plugins/agents/assets/kiro.svg diff --git a/shell/plugins/agents/assets/kiro-light.svg b/shell/plugins/agents/assets/kiro-light.svg new file mode 100644 index 00000000000..44767f48784 --- /dev/null +++ b/shell/plugins/agents/assets/kiro-light.svg @@ -0,0 +1 @@ +Kiro diff --git a/shell/plugins/agents/assets/kiro.svg b/shell/plugins/agents/assets/kiro.svg new file mode 100644 index 00000000000..c3f2c16a102 --- /dev/null +++ b/shell/plugins/agents/assets/kiro.svg @@ -0,0 +1 @@ +Kiro From 44e2b9412370ca43757789f5703fb64e1ff93593 Mon Sep 17 00:00:00 2001 From: Paul Lamb Date: Wed, 14 Oct 2026 23:21:20 -0400 Subject: [PATCH 118/546] Add a Freebuff usage collector omarchy-agent-usage-freebuff collects the agents panel record: exact prompts, sessions, and active days from the chat stores under ~/.config/manicode, tokens estimated from message text because the free tier meters credits server-side and no per-message token usage is recorded. Tests cover the collector (chats, days, tokens, sign-in, mid-write files). Generated with Codebuff Co-Authored-By: Codebuff Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-freebuff | 340 ++++++++++++++++++++++ test/shell.d/agent-usage-freebuff-test.sh | 116 ++++++++ 2 files changed, 456 insertions(+) create mode 100755 bin/omarchy-agent-usage-freebuff create mode 100755 test/shell.d/agent-usage-freebuff-test.sh diff --git a/bin/omarchy-agent-usage-freebuff b/bin/omarchy-agent-usage-freebuff new file mode 100755 index 00000000000..5669882d739 --- /dev/null +++ b/bin/omarchy-agent-usage-freebuff @@ -0,0 +1,340 @@ +#!/usr/bin/python3 +# omarchy:summary=Print the Freebuff usage record as JSON +# omarchy:args=[--force] [--limits-only] +# omarchy:hidden=true +"""Collect Freebuff usage into one display-ready JSON record. + +Freebuff keeps chats under ~/.config/manicode/projects//chats//: chat-messages.json (a JSON array of user prompts and AI replies), +chat-meta.json, and a debug log. It records no per-message token usage — the +free tier meters credits server-side ("creditsCharged": "server") — so token +counts here are estimates from message text (~4 characters per token), +attributed to a single "freebuff" model bucket. Prompts, sessions, and active +days are exact counts, dated from the epoch milliseconds embedded in every +message id (user-1790830285538). The agents panel only ever reads the JSON +this prints. +""" + +import argparse +import fcntl +import hashlib +import json +import os +import re +import sys +import tempfile +import time +from datetime import datetime, timedelta, timezone +from pathlib import Path + +AGENT_ID = "freebuff" +AGENT_NAME = "Freebuff" +AUTH_HELP = "Start Freebuff, or run `freebuff login`, to sign in." + +# 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, however low +# refreshIntervalSec is set. --limits-only promises only fresh limits, so it +# may reuse a scan for up to 15 minutes. Freebuff exposes no limits endpoint, +# so both windows exist purely to keep the rescan cost down. +SCAN_REUSE_SECONDS = 20 +LIMITS_ONLY_REUSE_SECONDS = 900 + +CHARS_PER_TOKEN = 4 +# Message ids look like user-1790830285538 / ai-1790830285726-8f3c26a13464e: +# the number after the variant is the creation time in epoch milliseconds +# (10-digit values would be seconds; accept both). +MESSAGE_ID_TS = re.compile(r"^(?:user|ai)-(\d{10,})") +# Escape-safe fallback for a chat-messages.json that fails to parse (a file +# caught mid-write): user messages in the real objects are never string-escaped, +# so this only ever undercounts transiently. +RAW_USER_VARIANT = re.compile(r'"variant"\s*:\s*"user"') + +now = datetime.now() +today = now.strftime("%Y-%m-%d") +recent_dates = [(now - timedelta(days=offset)).strftime("%Y-%m-%d") for offset in range(6, -1, -1)] +recent = {day: {"date": day, "messageCount": 0} for day in recent_dates} + + +def manicode_dir(): + # No official override env var is known for Freebuff's config dir; honor + # MANICODE_CONFIG_DIR so the path stays overridable for testing. + return Path(os.environ.get("MANICODE_CONFIG_DIR") or (Path.home() / ".config" / "manicode")) + + +def local_day_from_epoch(value): + try: + seconds = value / 1000.0 if value > 10_000_000_000 else float(value) + return datetime.fromtimestamp(seconds).strftime("%Y-%m-%d") + except (OverflowError, OSError, ValueError): + return None + + +def message_day(message): + match = MESSAGE_ID_TS.match(str(message.get("id") or "")) + if match: + day = local_day_from_epoch(int(match.group(1))) + if day: + return day + return None + + +def message_text(message): + """User prompts are plain strings; AI replies keep text in typed blocks.""" + parts = [] + content = message.get("content") + if isinstance(content, str): + parts.append(content) + blocks = message.get("blocks") + if isinstance(blocks, list): + for block in blocks: + if isinstance(block, dict) and block.get("type") == "text" and isinstance(block.get("content"), str): + parts.append(block["content"]) + return "".join(parts) + + +def chat_fallback_day(path): + """Day for messages whose id carries no timestamp: the file's last write.""" + try: + return datetime.fromtimestamp(path.stat().st_mtime).strftime("%Y-%m-%d") + except OSError: + return today + + +def scan_chat_file(path): + """Return {day: {"prompts": n, "inChars": n, "outChars": n}} for one chat. + + Returns {} for a chat with no datable messages; the caller decides whether + that still counts as a session. A parse failure falls back to a raw prompt + count on the file's last-write day so a session never vanishes from the + totals while Freebuff is rewriting it. + """ + fallback_day = chat_fallback_day(path) + per_day = {} + try: + data = json.loads(path.read_text(encoding="utf-8", errors="replace")) + except Exception: + data = None + + if isinstance(data, list): + for message in data: + if not isinstance(message, dict): + continue + variant = message.get("variant") + if variant not in ("user", "ai"): + continue + text = message_text(message) + if not text: + continue + day = message_day(message) or fallback_day + bucket = per_day.setdefault(day, {"prompts": 0, "inChars": 0, "outChars": 0}) + if variant == "user": + bucket["prompts"] += 1 + bucket["inChars"] += len(text) + else: + bucket["outChars"] += len(text) + return per_day + + # Mid-write or foreign file: prompts only, no token estimates. + try: + raw = path.read_text(encoding="utf-8", errors="replace") + prompts = len(RAW_USER_VARIANT.findall(raw)) + except OSError: + return {} + if prompts: + return {fallback_day: {"prompts": prompts, "inChars": 0, "outChars": 0}} + return {} + + +def chat_roots(): + root = manicode_dir() / "projects" + if not root.is_dir(): + return [] + try: + return sorted(root.glob("*/chats/*/chat-messages.json")) + except OSError: + return [] + + +def cache_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 cache_paths(): + digest = hashlib.sha1((str(Path.home()) + "\n" + str(manicode_dir())).encode("utf-8")).hexdigest()[:16] + root = cache_root() + return root / f"freebuff-scan-{digest}.json", root / f"freebuff-scan-{digest}.lock" + + +def read_fresh_cache(path, max_age_seconds): + 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: + cached = json.loads(path.read_text(encoding="utf-8")) + if isinstance(cached, dict) and cached.get("schemaVersion") == 1 and cached.get("scanDate") == today: + files = cached.get("files") + if isinstance(files, dict): + return files + except Exception: + return None + return None + + +def write_cache(path, entries): + 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: + payload = { + "schemaVersion": 1, + "scanDate": today, + "updatedAt": datetime.now(timezone.utc).isoformat(), + "files": entries, + } + handle.write(json.dumps(payload, separators=(",", ":")) + "\n") + tmp.chmod(0o644) + tmp.replace(path) + except BaseException: + tmp.unlink(missing_ok=True) + raise + + +def collect_entries(max_age): + """Per-chat usage entries, reusing unchanged files across scans. + + Each entry is {"mtime": float, "size": int, "perDay": {...}} keyed by the + chat-messages.json path. A file whose stat still matches its entry is reused + as-is; only new or rewritten chats are re-read, so the mid-session chat that + Freebuff rewrites on every message is the only recurring parse cost. + """ + cache_file, lock_file = cache_paths() + + with lock_file.open("w") as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + fresh = read_fresh_cache(cache_file, max_age) + if fresh is not None: + return fresh + + old_entries = {} + try: + cached = json.loads(cache_file.read_text(encoding="utf-8")) if cache_file.exists() else None + if isinstance(cached, dict) and cached.get("schemaVersion") == 1 and isinstance(cached.get("files"), dict): + old_entries = cached["files"] + except Exception: + old_entries = {} + + entries = {} + for path in chat_roots(): + key = str(path) + try: + stat = path.stat() + except OSError: + continue + old = old_entries.get(key) + if isinstance(old, dict) and old.get("size") == stat.st_size and abs(float(old.get("mtime") or -1) - stat.st_mtime) < 1e-6: + entries[key] = old + continue + per_day = scan_chat_file(path) + if per_day: + entries[key] = {"mtime": stat.st_mtime, "size": stat.st_size, "perDay": per_day} + + try: + write_cache(cache_file, entries) + except Exception as exc: + print(f"omarchy-agent-usage-freebuff: could not write usage cache ({exc})", file=sys.stderr) + return entries + + +def signed_in(): + return (manicode_dir() / "credentials.json").is_file() + + +def build_record(entries): + total_prompts = 0 + total_input_tokens = 0 + total_output_tokens = 0 + total_sessions = set() + active_days = set() + today_prompts = 0 + today_sessions = set() + today_total_tokens = 0 + today_tokens_by_model = {} + model_usage_bucket = {"inputTokens": 0, "outputTokens": 0, "cacheReadInputTokens": 0, "cacheCreationInputTokens": 0} + + for key, entry in entries.items(): + per_day = entry.get("perDay") if isinstance(entry, dict) else None + if not isinstance(per_day, dict) or not per_day: + continue + total_sessions.add(key) + for day, bucket in per_day.items(): + if not isinstance(bucket, dict): + continue + prompts = int(bucket.get("prompts") or 0) + input_tokens = int(bucket.get("inChars") or 0) // CHARS_PER_TOKEN + output_tokens = int(bucket.get("outChars") or 0) // CHARS_PER_TOKEN + total = input_tokens + output_tokens + total_prompts += prompts + total_input_tokens += input_tokens + total_output_tokens += output_tokens + model_usage_bucket["inputTokens"] += input_tokens + model_usage_bucket["outputTokens"] += output_tokens + active_days.add(day) + if day in recent: + recent[day]["messageCount"] += total + if day == today: + today_prompts += prompts + today_total_tokens += total + today_sessions.add(key) + today_tokens_by_model[AGENT_ID] = today_tokens_by_model.get(AGENT_ID, 0) + total + + authenticated = signed_in() + has_tokens = (total_input_tokens + total_output_tokens) > 0 + record = { + "schemaVersion": 1, + "id": AGENT_ID, + "name": AGENT_NAME, + "updatedAt": datetime.now(timezone.utc).isoformat(), + "ready": authenticated, + "hasLocalStats": True, + "hasPromptStats": True, + "tierLabel": "Free" if authenticated else "", + "usageStatusText": "" if authenticated else "Freebuff not signed in", + "authHelpText": "" if authenticated else AUTH_HELP, + "limits": [], + "todayPrompts": today_prompts, + "todaySessions": len(today_sessions), + "todayTotalTokens": today_total_tokens, + "todayTokensByModel": today_tokens_by_model, + "recentDays": [recent[day] for day in recent_dates], + "totalPrompts": total_prompts, + "totalSessions": len(total_sessions), + "activeDays": len(active_days), + "activeDates": sorted(active_days), + "modelUsage": {AGENT_ID: model_usage_bucket} if has_tokens else {}, + } + return record + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--force", action="store_true") + parser.add_argument("--limits-only", action="store_true") + args = parser.parse_args() + + max_age = 0 if args.force else (LIMITS_ONLY_REUSE_SECONDS if args.limits_only else SCAN_REUSE_SECONDS) + try: + entries = collect_entries(max_age) + except Exception as exc: + print(f"omarchy-agent-usage-freebuff: scan unavailable ({exc})", file=sys.stderr) + entries = {} + print(json.dumps(build_record(entries), separators=(",", ":"))) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/test/shell.d/agent-usage-freebuff-test.sh b/test/shell.d/agent-usage-freebuff-test.sh new file mode 100755 index 00000000000..34f147010e8 --- /dev/null +++ b/test/shell.d/agent-usage-freebuff-test.sh @@ -0,0 +1,116 @@ +#!/bin/bash + +set -euo pipefail + +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" + +require_command python3 + +test_tmp=$(mktemp -d) +trap 'rm -rf "$test_tmp"' EXIT + +export HOME="$test_tmp/home" +export XDG_STATE_HOME="$test_tmp/state" +export XDG_CACHE_HOME="$test_tmp/cache" + +collector="$ROOT/bin/omarchy-agent-usage-freebuff" +projects="$HOME/.config/manicode/projects/testproj/chats" + +# Freebuff stores each chat as a directory named for its UTC start time, +# holding a chat-messages.json array. Message ids carry the epoch +# milliseconds of their creation; user messages are plain strings and AI +# replies keep text in typed blocks. +make_chat() { + local dir="$1" + mkdir -p "$dir" + cat >"$dir/chat-messages.json" +} + +today=$(python3 -c 'import datetime as dt; print(dt.datetime.now().strftime("%Y-%m-%d"))') +yesterday=$(python3 -c 'import datetime as dt; print((dt.datetime.now() - dt.timedelta(days=1)).strftime("%Y-%m-%d"))') +now_ms=$(python3 -c 'import time; print(int(time.time() * 1000))') +yesterday_ms=$(python3 -c 'import time; print(int((time.time() - 86400) * 1000))') + +# One chat from yesterday, one from today: two prompts and two replies total. +make_chat "$projects/2026-09-30T10-00-00.000Z" < 0)] | length') == 2 ]] || fail "recentDays" +pass "weekly chart has two non-zero days" + +# Unsigned in: no credentials file, so the record says how to fix it. +[[ $(printf '%s' "$record" | jq '.ready') == "false" ]] || fail "ready" +[[ $(printf '%s' "$record" | jq -r '.authHelpText') == "Start Freebuff, or run \`freebuff login\`, to sign in." ]] || fail "authHelpText" +pass "unsigned-in record points at sign-in" + +mkdir -p "$HOME/.config/manicode" +printf '{}' >"$HOME/.config/manicode/credentials.json" +record=$("$collector") +[[ $(printf '%s' "$record" | jq '.ready') == "true" ]] || fail "ready signed-in" +[[ $(printf '%s' "$record" | jq -r '.tierLabel') == "Free" ]] || fail "tierLabel" +[[ $(printf '%s' "$record" | jq -r '.usageStatusText') == "" ]] || fail "usageStatusText" +pass "signed-in record reports the Free tier" + +# A chat caught mid-write (invalid JSON) must not take the scan down: its +# prompts are still counted, from the escaped user-variant fallback. +mkdir -p "$projects/2026-10-01T16-00-00.000Z" +printf '[{"id": "user-%s", "variant": "user", "content": "half written' "$now_ms" \ + >"$projects/2026-10-01T16-00-00.000Z/chat-messages.json" +# --force: the collector dedups concurrent runs through a recent-scan cache, +# so a normal rerun inside the test would legitimately reuse the scan above. +record=$("$collector" --force) +[[ $(printf '%s' "$record" | jq '.totalPrompts') == 4 ]] || fail "midwrite prompts" +[[ $(printf '%s' "$record" | jq '.totalSessions') == 3 ]] || fail "midwrite sessions" +pass "mid-write chat still counts its prompts" + +# The update wrapper writes the record where the panel reads it. It scans for +# collectors under OMARCHY_PATH, which the running shell points at the +# installed tree, so the test points it at this checkout the same way. +export PATH="$ROOT/bin:$PATH" +export OMARCHY_PATH="$ROOT" +"$ROOT/bin/omarchy-agent-usage-update" freebuff +[[ -f "$XDG_STATE_HOME/omarchy/agents/usage/freebuff.json" ]] || fail "wrapper record" +jq -e '.id == "freebuff"' "$XDG_STATE_HOME/omarchy/agents/usage/freebuff.json" >/dev/null || fail "wrapper json" +pass "usage-update writes the panel's freebuff.json" + +# And the stock agents still collect through the same wrapper. +"$ROOT/bin/omarchy-agent-usage-update" codex +[[ -f "$XDG_STATE_HOME/omarchy/agents/usage/codex.json" ]] || fail "codex via wrapper" +pass "usage-update still serves the stock collectors" From 056b08b18c39b66e09c3c1bcad4e44f9b76a4682 Mon Sep 17 00:00:00 2001 From: Paul Lamb Date: Wed, 14 Oct 2026 23:21:26 -0400 Subject: [PATCH 119/546] Harden Freebuff usage reporting Keep a malformed or mid-write chat out of the cache instead of guessing its prompts, honor FREEBUFF_CONFIG_DIR, and report a keychain sign-in as unverifiable rather than signed in. Generated with Codebuff Co-Authored-By: Codebuff Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-freebuff | 143 ++++++++++++++-------- test/shell.d/agent-usage-freebuff-test.sh | 43 +++++-- 2 files changed, 123 insertions(+), 63 deletions(-) diff --git a/bin/omarchy-agent-usage-freebuff b/bin/omarchy-agent-usage-freebuff index 5669882d739..fd2706ee248 100755 --- a/bin/omarchy-agent-usage-freebuff +++ b/bin/omarchy-agent-usage-freebuff @@ -45,10 +45,6 @@ CHARS_PER_TOKEN = 4 # the number after the variant is the creation time in epoch milliseconds # (10-digit values would be seconds; accept both). MESSAGE_ID_TS = re.compile(r"^(?:user|ai)-(\d{10,})") -# Escape-safe fallback for a chat-messages.json that fails to parse (a file -# caught mid-write): user messages in the real objects are never string-escaped, -# so this only ever undercounts transiently. -RAW_USER_VARIANT = re.compile(r'"variant"\s*:\s*"user"') now = datetime.now() today = now.strftime("%Y-%m-%d") @@ -57,9 +53,15 @@ recent = {day: {"date": day, "messageCount": 0} for day in recent_dates} def manicode_dir(): - # No official override env var is known for Freebuff's config dir; honor - # MANICODE_CONFIG_DIR so the path stays overridable for testing. - return Path(os.environ.get("MANICODE_CONFIG_DIR") or (Path.home() / ".config" / "manicode")) + # Freebuff's supported override; it must be absolute. The compatibility + # variable remains useful for isolated collector tests. + configured = os.environ.get("FREEBUFF_CONFIG_DIR") or os.environ.get("MANICODE_CONFIG_DIR") + if configured: + path = Path(configured).expanduser() + if path.is_absolute(): + return path + print("omarchy-agent-usage-freebuff: ignoring relative config directory", file=sys.stderr) + return Path.home() / ".config" / "manicode" def local_day_from_epoch(value): @@ -102,48 +104,47 @@ def chat_fallback_day(path): def scan_chat_file(path): - """Return {day: {"prompts": n, "inChars": n, "outChars": n}} for one chat. + """Return (per-day usage, complete) for one chat. - Returns {} for a chat with no datable messages; the caller decides whether - that still counts as a session. A parse failure falls back to a raw prompt - count on the file's last-write day so a session never vanishes from the - totals while Freebuff is rewriting it. + A malformed, unreadable, or concurrently modified file is omitted from this + run and marks the scan incomplete, so a partial snapshot never enters cache. """ - fallback_day = chat_fallback_day(path) - per_day = {} try: - data = json.loads(path.read_text(encoding="utf-8", errors="replace")) - except Exception: - data = None + before = path.stat() + raw = path.read_text(encoding="utf-8") + after = path.stat() + except (OSError, UnicodeError): + return {}, False + + if before.st_size != after.st_size or before.st_mtime_ns != after.st_mtime_ns: + return {}, False - if isinstance(data, list): - for message in data: - if not isinstance(message, dict): - continue - variant = message.get("variant") - if variant not in ("user", "ai"): - continue - text = message_text(message) - if not text: - continue - day = message_day(message) or fallback_day - bucket = per_day.setdefault(day, {"prompts": 0, "inChars": 0, "outChars": 0}) - if variant == "user": - bucket["prompts"] += 1 - bucket["inChars"] += len(text) - else: - bucket["outChars"] += len(text) - return per_day - - # Mid-write or foreign file: prompts only, no token estimates. try: - raw = path.read_text(encoding="utf-8", errors="replace") - prompts = len(RAW_USER_VARIANT.findall(raw)) - except OSError: - return {} - if prompts: - return {fallback_day: {"prompts": prompts, "inChars": 0, "outChars": 0}} - return {} + data = json.loads(raw) + except (json.JSONDecodeError, UnicodeError): + return {}, False + if not isinstance(data, list): + return {}, False + + fallback_day = chat_fallback_day(path) + per_day = {} + for message in data: + if not isinstance(message, dict): + continue + variant = message.get("variant") + if variant not in ("user", "ai"): + continue + text = message_text(message) + if not text: + continue + day = message_day(message) or fallback_day + bucket = per_day.setdefault(day, {"prompts": 0, "inChars": 0, "outChars": 0}) + if variant == "user": + bucket["prompts"] += 1 + bucket["inChars"] += len(text) + else: + bucket["outChars"] += len(text) + return per_day, True def chat_roots(): @@ -230,29 +231,62 @@ def collect_entries(max_age): old_entries = {} entries = {} + complete = True for path in chat_roots(): key = str(path) try: stat = path.stat() except OSError: + complete = False continue old = old_entries.get(key) if isinstance(old, dict) and old.get("size") == stat.st_size and abs(float(old.get("mtime") or -1) - stat.st_mtime) < 1e-6: entries[key] = old continue - per_day = scan_chat_file(path) + per_day, file_complete = scan_chat_file(path) + complete = complete and file_complete if per_day: entries[key] = {"mtime": stat.st_mtime, "size": stat.st_size, "perDay": per_day} - try: - write_cache(cache_file, entries) - except Exception as exc: - print(f"omarchy-agent-usage-freebuff: could not write usage cache ({exc})", file=sys.stderr) + if complete: + try: + write_cache(cache_file, entries) + except Exception as exc: + print(f"omarchy-agent-usage-freebuff: could not write usage cache ({exc})", file=sys.stderr) + else: + # The old snapshot predates this incomplete observation and can be + # reused by --limits-only; discard it so the next request must rescan. + try: + cache_file.unlink(missing_ok=True) + except OSError as exc: + print(f"omarchy-agent-usage-freebuff: could not invalidate usage cache ({exc})", file=sys.stderr) + print("omarchy-agent-usage-freebuff: incomplete scan; not caching partial usage", file=sys.stderr) return entries -def signed_in(): - return (manicode_dir() / "credentials.json").is_file() +def authentication_status(): + """Return (status, help text) without treating a keychain pointer as a token. + + Freebuff may move authToken into the OS keychain, which this collector cannot + query. Report that state as unverifiable rather than claiming sign-in. + """ + path = manicode_dir() / "credentials.json" + try: + data = json.loads(path.read_text(encoding="utf-8")) + except FileNotFoundError: + return "signed-out", AUTH_HELP + except (OSError, UnicodeError, json.JSONDecodeError): + return "unverified", "Could not read Freebuff credentials; sign-in status is unknown." + + profile = data.get("default") if isinstance(data, dict) else None + if not isinstance(profile, dict): + return "signed-out", AUTH_HELP + token = profile.get("authToken") + if isinstance(token, str) and token.strip(): + return "signed-in", "" + if profile.get("tokenStore") == "keychain": + return "unverified", "Freebuff credentials are stored in the system keychain; sign-in status cannot be verified here." + return "signed-out", AUTH_HELP def build_record(entries): @@ -293,7 +327,8 @@ def build_record(entries): today_sessions.add(key) today_tokens_by_model[AGENT_ID] = today_tokens_by_model.get(AGENT_ID, 0) + total - authenticated = signed_in() + auth_status, auth_help = authentication_status() + authenticated = auth_status == "signed-in" has_tokens = (total_input_tokens + total_output_tokens) > 0 record = { "schemaVersion": 1, @@ -304,8 +339,8 @@ def build_record(entries): "hasLocalStats": True, "hasPromptStats": True, "tierLabel": "Free" if authenticated else "", - "usageStatusText": "" if authenticated else "Freebuff not signed in", - "authHelpText": "" if authenticated else AUTH_HELP, + "usageStatusText": "" if authenticated else ("Freebuff not signed in" if auth_status == "signed-out" else "Freebuff sign-in status unknown"), + "authHelpText": auth_help, "limits": [], "todayPrompts": today_prompts, "todaySessions": len(today_sessions), diff --git a/test/shell.d/agent-usage-freebuff-test.sh b/test/shell.d/agent-usage-freebuff-test.sh index 34f147010e8..f0f13479c66 100755 --- a/test/shell.d/agent-usage-freebuff-test.sh +++ b/test/shell.d/agent-usage-freebuff-test.sh @@ -80,25 +80,50 @@ pass "weekly chart has two non-zero days" [[ $(printf '%s' "$record" | jq -r '.authHelpText') == "Start Freebuff, or run \`freebuff login\`, to sign in." ]] || fail "authHelpText" pass "unsigned-in record points at sign-in" -mkdir -p "$HOME/.config/manicode" -printf '{}' >"$HOME/.config/manicode/credentials.json" -record=$("$collector") +credentials="$HOME/.config/manicode/credentials.json" +mkdir -p "$(dirname "$credentials")" + +for invalid in '{}' '{"default":{}}' '{"default":{"authToken":""}}' 'not-json'; do + printf '%s' "$invalid" >"$credentials" + record=$("$collector" --force) + [[ $(printf '%s' "$record" | jq '.ready') == "false" ]] || fail "invalid credentials not ready: $invalid" +done +pass "empty, unusable, and malformed credentials are not signed in" + +printf '{"default":{"authToken":"test-token"}}' >"$credentials" +record=$("$collector" --force) [[ $(printf '%s' "$record" | jq '.ready') == "true" ]] || fail "ready signed-in" [[ $(printf '%s' "$record" | jq -r '.tierLabel') == "Free" ]] || fail "tierLabel" [[ $(printf '%s' "$record" | jq -r '.usageStatusText') == "" ]] || fail "usageStatusText" -pass "signed-in record reports the Free tier" +pass "usable token credentials report the Free tier" -# A chat caught mid-write (invalid JSON) must not take the scan down: its -# prompts are still counted, from the escaped user-variant fallback. +# A keychain pointer contains no token for this standalone collector. +printf '{"default":{"tokenStore":"keychain"}}' >"$credentials" +record=$("$collector" --force) +[[ $(printf '%s' "$record" | jq '.ready') == "false" ]] || fail "keychain pointer is not verified" +[[ $(printf '%s' "$record" | jq -r '.usageStatusText') == "Freebuff sign-in status unknown" ]] || fail "keychain status" +[[ $(printf '%s' "$record" | jq -r '.authHelpText') == *"cannot be verified"* ]] || fail "keychain help" +pass "keychain-backed sign-in is reported as unverifiable" + +# A chat caught mid-write (invalid JSON) must not take the scan down or be +# counted as a complete prompt until it can be read successfully. mkdir -p "$projects/2026-10-01T16-00-00.000Z" printf '[{"id": "user-%s", "variant": "user", "content": "half written' "$now_ms" \ >"$projects/2026-10-01T16-00-00.000Z/chat-messages.json" # --force: the collector dedups concurrent runs through a recent-scan cache, # so a normal rerun inside the test would legitimately reuse the scan above. record=$("$collector" --force) -[[ $(printf '%s' "$record" | jq '.totalPrompts') == 4 ]] || fail "midwrite prompts" -[[ $(printf '%s' "$record" | jq '.totalSessions') == 3 ]] || fail "midwrite sessions" -pass "mid-write chat still counts its prompts" +[[ $(printf '%s' "$record" | jq '.totalPrompts') == 3 ]] || fail "partial prompt omitted" +[[ $(printf '%s' "$record" | jq '.totalSessions') == 2 ]] || fail "partial session omitted" +pass "incomplete chat is omitted from displayed usage" + +cat >"$projects/2026-10-01T16-00-00.000Z/chat-messages.json" < Date: Wed, 14 Oct 2026 23:23:53 -0400 Subject: [PATCH 120/546] Report Freebuff's real counts only and mark it ready on local use Freebuff's chat files carry no token usage, so the four-characters-per-token estimate is gone: prompts, sessions, and active days stay exact, and the token fields, model buckets, and daily chart stay at zero. The record is ready once there are local chats or a sign-in, and a keychain sign-in now counts, since Freebuff leaves its profile in credentials.json when it moves the token there. Freebuff has no limits, so a missing sign-in is nothing to fix: the off-vocabulary statuses are gone, and the plan stays unnamed because Freebuff Pro exists and nothing on disk says which one this is. The scan cache keeps its per-file reuse, gains a schema bump for the new entry shape and the 0 <= age < N guard, rejects message-id times outside a plausible range, and drops the MANICODE_CONFIG_DIR override Freebuff doesn't read. Co-authored-by: Paul Lamb Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-freebuff | 219 +++++++++------------- test/shell.d/agent-usage-freebuff-test.sh | 102 +++++++--- 2 files changed, 165 insertions(+), 156 deletions(-) diff --git a/bin/omarchy-agent-usage-freebuff b/bin/omarchy-agent-usage-freebuff index fd2706ee248..3d40665ffe9 100755 --- a/bin/omarchy-agent-usage-freebuff +++ b/bin/omarchy-agent-usage-freebuff @@ -5,14 +5,13 @@ """Collect Freebuff usage into one display-ready JSON record. Freebuff keeps chats under ~/.config/manicode/projects//chats//: chat-messages.json (a JSON array of user prompts and AI replies), -chat-meta.json, and a debug log. It records no per-message token usage — the -free tier meters credits server-side ("creditsCharged": "server") — so token -counts here are estimates from message text (~4 characters per token), -attributed to a single "freebuff" model bucket. Prompts, sessions, and active -days are exact counts, dated from the epoch milliseconds embedded in every -message id (user-1790830285538). The agents panel only ever reads the JSON -this prints. +timestamp>/ (FREEBUFF_CONFIG_DIR moves the whole directory), each with a +chat-messages.json array of user prompts and AI replies. Prompts, sessions, +and active days are exact counts, dated from the epoch milliseconds embedded +in every message id (user-1790830285538). The files record no token usage, +which Freebuff meters on its servers, so the record's token fields stay at +zero rather than guessing from message text. Freebuff exposes no limits to +read. The agents panel only ever reads the JSON this prints. """ import argparse @@ -29,7 +28,6 @@ from pathlib import Path AGENT_ID = "freebuff" AGENT_NAME = "Freebuff" -AUTH_HELP = "Start Freebuff, or run `freebuff login`, to sign in." # 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 @@ -40,22 +38,24 @@ AUTH_HELP = "Start Freebuff, or run `freebuff login`, to sign in." SCAN_REUSE_SECONDS = 20 LIMITS_ONLY_REUSE_SECONDS = 900 -CHARS_PER_TOKEN = 4 +# Version 2 entries hold prompt counts per day; version 1 held the character +# counts the token estimate came from. +CACHE_SCHEMA_VERSION = 2 + # Message ids look like user-1790830285538 / ai-1790830285726-8f3c26a13464e: # the number after the variant is the creation time in epoch milliseconds # (10-digit values would be seconds; accept both). -MESSAGE_ID_TS = re.compile(r"^(?:user|ai)-(\d{10,})") +MESSAGE_ID_TS = re.compile(r"^user-(\d{10,13})(?:\D|$)") +MIN_EPOCH_SECONDS = 946_684_800 # 2000-01-01 now = datetime.now() today = now.strftime("%Y-%m-%d") recent_dates = [(now - timedelta(days=offset)).strftime("%Y-%m-%d") for offset in range(6, -1, -1)] -recent = {day: {"date": day, "messageCount": 0} for day in recent_dates} def manicode_dir(): - # Freebuff's supported override; it must be absolute. The compatibility - # variable remains useful for isolated collector tests. - configured = os.environ.get("FREEBUFF_CONFIG_DIR") or os.environ.get("MANICODE_CONFIG_DIR") + # Freebuff's own override, which it requires to be absolute. + configured = os.environ.get("FREEBUFF_CONFIG_DIR") if configured: path = Path(configured).expanduser() if path.is_absolute(): @@ -65,46 +65,39 @@ def manicode_dir(): def local_day_from_epoch(value): - try: - seconds = value / 1000.0 if value > 10_000_000_000 else float(value) - return datetime.fromtimestamp(seconds).strftime("%Y-%m-%d") - except (OverflowError, OSError, ValueError): + """The local day for an epoch in seconds or milliseconds, or None for a + time that can't be a real moment of use.""" + seconds = value / 1000.0 if value > 10_000_000_000 else float(value) + # A day of slack past now covers a clock a little behind the writer's. + if seconds < MIN_EPOCH_SECONDS or seconds > time.time() + 86_400: return None + return datetime.fromtimestamp(seconds).strftime("%Y-%m-%d") def message_day(message): match = MESSAGE_ID_TS.match(str(message.get("id") or "")) - if match: - day = local_day_from_epoch(int(match.group(1))) - if day: - return day - return None + return local_day_from_epoch(int(match.group(1))) if match else None -def message_text(message): - """User prompts are plain strings; AI replies keep text in typed blocks.""" - parts = [] +def is_prompt(message): + """A user message that says something; Freebuff stores prompts as plain + strings.""" + if message.get("variant") != "user": + return False content = message.get("content") - if isinstance(content, str): - parts.append(content) - blocks = message.get("blocks") - if isinstance(blocks, list): - for block in blocks: - if isinstance(block, dict) and block.get("type") == "text" and isinstance(block.get("content"), str): - parts.append(block["content"]) - return "".join(parts) + return isinstance(content, str) and bool(content.strip()) def chat_fallback_day(path): - """Day for messages whose id carries no timestamp: the file's last write.""" + """Day for prompts whose id carries no usable time: the file's last write.""" try: - return datetime.fromtimestamp(path.stat().st_mtime).strftime("%Y-%m-%d") + return local_day_from_epoch(path.stat().st_mtime) or today except OSError: return today def scan_chat_file(path): - """Return (per-day usage, complete) for one chat. + """Return ({day: prompts}, complete) for one chat. A malformed, unreadable, or concurrently modified file is omitted from this run and marks the scan incomplete, so a partial snapshot never enters cache. @@ -129,21 +122,10 @@ def scan_chat_file(path): fallback_day = chat_fallback_day(path) per_day = {} for message in data: - if not isinstance(message, dict): - continue - variant = message.get("variant") - if variant not in ("user", "ai"): - continue - text = message_text(message) - if not text: + if not isinstance(message, dict) or not is_prompt(message): continue day = message_day(message) or fallback_day - bucket = per_day.setdefault(day, {"prompts": 0, "inChars": 0, "outChars": 0}) - if variant == "user": - bucket["prompts"] += 1 - bucket["inChars"] += len(text) - else: - bucket["outChars"] += len(text) + per_day[day] = per_day.get(day, 0) + 1 return per_day, True @@ -169,22 +151,32 @@ def cache_paths(): return root / f"freebuff-scan-{digest}.json", root / f"freebuff-scan-{digest}.lock" +def read_cache(path): + try: + cached = json.loads(path.read_text(encoding="utf-8")) + except Exception: + return None + if isinstance(cached, dict) and cached.get("schemaVersion") == CACHE_SCHEMA_VERSION and isinstance(cached.get("files"), dict): + return cached + return None + + def read_fresh_cache(path, max_age_seconds): - if max_age_seconds <= 0 or not path.exists(): + if max_age_seconds <= 0: 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: - cached = json.loads(path.read_text(encoding="utf-8")) - if isinstance(cached, dict) and cached.get("schemaVersion") == 1 and cached.get("scanDate") == today: - files = cached.get("files") - if isinstance(files, dict): - return files - except Exception: + except OSError: return None - return None + if not 0 <= age < max_age_seconds: + return None + cached = read_cache(path) + # today* fields only mean "today" on the day they were scanned. + if cached is None or cached.get("scanDate") != today: + return None + return cached["files"] def write_cache(path, entries): @@ -193,7 +185,7 @@ def write_cache(path, entries): try: with os.fdopen(handle_fd, "w", encoding="utf-8") as handle: payload = { - "schemaVersion": 1, + "schemaVersion": CACHE_SCHEMA_VERSION, "scanDate": today, "updatedAt": datetime.now(timezone.utc).isoformat(), "files": entries, @@ -209,10 +201,11 @@ def write_cache(path, entries): def collect_entries(max_age): """Per-chat usage entries, reusing unchanged files across scans. - Each entry is {"mtime": float, "size": int, "perDay": {...}} keyed by the - chat-messages.json path. A file whose stat still matches its entry is reused - as-is; only new or rewritten chats are re-read, so the mid-session chat that - Freebuff rewrites on every message is the only recurring parse cost. + Each entry is {"mtime": float, "size": int, "perDay": {day: prompts}} + keyed by the chat-messages.json path. A file whose stat still matches its + entry is reused as-is; only new or rewritten chats are re-read, so the + mid-session chat that Freebuff rewrites on every message is the only + recurring parse cost. """ cache_file, lock_file = cache_paths() @@ -222,13 +215,7 @@ def collect_entries(max_age): if fresh is not None: return fresh - old_entries = {} - try: - cached = json.loads(cache_file.read_text(encoding="utf-8")) if cache_file.exists() else None - if isinstance(cached, dict) and cached.get("schemaVersion") == 1 and isinstance(cached.get("files"), dict): - old_entries = cached["files"] - except Exception: - old_entries = {} + old_entries = (read_cache(cache_file) or {}).get("files") or {} entries = {} complete = True @@ -240,7 +227,7 @@ def collect_entries(max_age): complete = False continue old = old_entries.get(key) - if isinstance(old, dict) and old.get("size") == stat.st_size and abs(float(old.get("mtime") or -1) - stat.st_mtime) < 1e-6: + if isinstance(old, dict) and old.get("size") == stat.st_size and old.get("mtime") == stat.st_mtime: entries[key] = old continue per_day, file_complete = scan_chat_file(path) @@ -264,100 +251,76 @@ def collect_entries(max_age): return entries -def authentication_status(): - """Return (status, help text) without treating a keychain pointer as a token. - - Freebuff may move authToken into the OS keychain, which this collector cannot - query. Report that state as unverifiable rather than claiming sign-in. - """ - path = manicode_dir() / "credentials.json" +def signed_in(): + """Whether Freebuff holds a sign-in here: a token in credentials.json, or + the profile Freebuff leaves there after moving the token to the keychain.""" try: - data = json.loads(path.read_text(encoding="utf-8")) - except FileNotFoundError: - return "signed-out", AUTH_HELP - except (OSError, UnicodeError, json.JSONDecodeError): - return "unverified", "Could not read Freebuff credentials; sign-in status is unknown." - + data = json.loads((manicode_dir() / "credentials.json").read_text(encoding="utf-8")) + except (OSError, UnicodeError, ValueError): + return False profile = data.get("default") if isinstance(data, dict) else None if not isinstance(profile, dict): - return "signed-out", AUTH_HELP + return False token = profile.get("authToken") - if isinstance(token, str) and token.strip(): - return "signed-in", "" - if profile.get("tokenStore") == "keychain": - return "unverified", "Freebuff credentials are stored in the system keychain; sign-in status cannot be verified here." - return "signed-out", AUTH_HELP + return (isinstance(token, str) and bool(token.strip())) or profile.get("tokenStore") == "keychain" def build_record(entries): total_prompts = 0 - total_input_tokens = 0 - total_output_tokens = 0 total_sessions = set() active_days = set() today_prompts = 0 today_sessions = set() - today_total_tokens = 0 - today_tokens_by_model = {} - model_usage_bucket = {"inputTokens": 0, "outputTokens": 0, "cacheReadInputTokens": 0, "cacheCreationInputTokens": 0} for key, entry in entries.items(): per_day = entry.get("perDay") if isinstance(entry, dict) else None - if not isinstance(per_day, dict) or not per_day: + if not isinstance(per_day, dict): continue - total_sessions.add(key) - for day, bucket in per_day.items(): - if not isinstance(bucket, dict): + for day, prompts in per_day.items(): + if isinstance(prompts, bool) or not isinstance(prompts, int) or prompts <= 0: continue - prompts = int(bucket.get("prompts") or 0) - input_tokens = int(bucket.get("inChars") or 0) // CHARS_PER_TOKEN - output_tokens = int(bucket.get("outChars") or 0) // CHARS_PER_TOKEN - total = input_tokens + output_tokens total_prompts += prompts - total_input_tokens += input_tokens - total_output_tokens += output_tokens - model_usage_bucket["inputTokens"] += input_tokens - model_usage_bucket["outputTokens"] += output_tokens + total_sessions.add(key) active_days.add(day) - if day in recent: - recent[day]["messageCount"] += total if day == today: today_prompts += prompts - today_total_tokens += total today_sessions.add(key) - today_tokens_by_model[AGENT_ID] = today_tokens_by_model.get(AGENT_ID, 0) + total - auth_status, auth_help = authentication_status() - authenticated = auth_status == "signed-in" - has_tokens = (total_input_tokens + total_output_tokens) > 0 - record = { + # Freebuff has no limits, so a missing sign-in is nothing to fix here: the + # status stays empty, and a machine with neither chats nor a sign-in prints + # a record the panel skips. Freebuff has a paid tier too, and nothing on + # disk says which one this is, so the plan stays unnamed. + return { "schemaVersion": 1, "id": AGENT_ID, "name": AGENT_NAME, "updatedAt": datetime.now(timezone.utc).isoformat(), - "ready": authenticated, + "ready": total_prompts > 0 or signed_in(), "hasLocalStats": True, "hasPromptStats": True, - "tierLabel": "Free" if authenticated else "", - "usageStatusText": "" if authenticated else ("Freebuff not signed in" if auth_status == "signed-out" else "Freebuff sign-in status unknown"), - "authHelpText": auth_help, + "tierLabel": "", + "usageStatusText": "", + "authHelpText": "", "limits": [], + "limitsStale": False, + "limitsFetchedAt": 0, "todayPrompts": today_prompts, "todaySessions": len(today_sessions), - "todayTotalTokens": today_total_tokens, - "todayTokensByModel": today_tokens_by_model, - "recentDays": [recent[day] for day in recent_dates], + "todayTotalTokens": 0, + "todayTokensByModel": {}, + "recentDays": [{"date": day, "messageCount": 0} for day in recent_dates], "totalPrompts": total_prompts, "totalSessions": len(total_sessions), + # The dates travel too: merging snapshots from several machines needs + # their union, which a count alone cannot give. "activeDays": len(active_days), "activeDates": sorted(active_days), - "modelUsage": {AGENT_ID: model_usage_bucket} if has_tokens else {}, + "modelUsage": {}, } - return record def main(): - parser = argparse.ArgumentParser() + parser = argparse.ArgumentParser(description="Print the Freebuff usage record as JSON") parser.add_argument("--force", action="store_true") parser.add_argument("--limits-only", action="store_true") args = parser.parse_args() @@ -372,4 +335,4 @@ def main(): if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/test/shell.d/agent-usage-freebuff-test.sh b/test/shell.d/agent-usage-freebuff-test.sh index f0f13479c66..edb2e56cd84 100755 --- a/test/shell.d/agent-usage-freebuff-test.sh +++ b/test/shell.d/agent-usage-freebuff-test.sh @@ -63,47 +63,66 @@ pass "prompts, sessions, and active days counted across chats" [[ $(printf '%s' "$record" | jq '.todayPrompts') == 1 ]] || fail "todayPrompts" [[ $(printf '%s' "$record" | jq '.todaySessions') == 1 ]] || fail "todaySessions" -[[ $(printf '%s' "$record" | jq '.todayTotalTokens') -gt 0 ]] || fail "todayTotalTokens" pass "today's numbers cover only today's chat" -# Tokens are estimated from message text (~4 chars/token) in one bucket. -[[ $(printf '%s' "$record" | jq '.modelUsage.freebuff.inputTokens') -gt 0 ]] || fail "inputTokens" -[[ $(printf '%s' "$record" | jq '.modelUsage.freebuff.outputTokens') -gt 0 ]] || fail "outputTokens" -[[ $(printf '%s' "$record" | jq '.modelUsage.freebuff.cacheReadInputTokens') == 0 ]] || fail "cacheRead" -pass "token estimates land in one freebuff model bucket" +# The chat files record no token usage, and the record invents none. +[[ $(printf '%s' "$record" | jq -c '{todayTotalTokens, todayTokensByModel, modelUsage, chart: ([.recentDays[].messageCount] | add), days: (.recentDays | length)}') == '{"todayTotalTokens":0,"todayTokensByModel":{},"modelUsage":{},"chart":0,"days":7}' ]] || + fail "token fields stay at zero" "$record" +pass "token fields stay at zero rather than estimated from text" -[[ $(printf '%s' "$record" | jq '[.recentDays[] | select(.messageCount > 0)] | length') == 2 ]] || fail "recentDays" -pass "weekly chart has two non-zero days" +# Local chats alone make Freebuff in use, signed in or not, and a missing +# sign-in is nothing to fix: Freebuff has no limits to read. +[[ $(printf '%s' "$record" | jq -c '{ready, tierLabel, usageStatusText, authHelpText, limits}') == '{"ready":true,"tierLabel":"","usageStatusText":"","authHelpText":"","limits":[]}' ]] || + fail "local usage makes the record ready" "$record" +pass "local chats make the record ready without a sign-in" -# Unsigned in: no credentials file, so the record says how to fix it. -[[ $(printf '%s' "$record" | jq '.ready') == "false" ]] || fail "ready" -[[ $(printf '%s' "$record" | jq -r '.authHelpText') == "Start Freebuff, or run \`freebuff login\`, to sign in." ]] || fail "authHelpText" -pass "unsigned-in record points at sign-in" +# The sign-in checks run on an empty config home, where only the sign-in +# decides readiness. +signin_home="$test_tmp/signin" +credentials="$signin_home/credentials.json" +mkdir -p "$signin_home" -credentials="$HOME/.config/manicode/credentials.json" -mkdir -p "$(dirname "$credentials")" +record=$(FREEBUFF_CONFIG_DIR="$signin_home" "$collector" --force) +[[ $(printf '%s' "$record" | jq -c '{ready, tierLabel, usageStatusText, authHelpText, limits, balance: has("balance"), totalPrompts, totalSessions, activeDays, activeDates}') == '{"ready":false,"tierLabel":"","usageStatusText":"","authHelpText":"","limits":[],"balance":false,"totalPrompts":0,"totalSessions":0,"activeDays":0,"activeDates":[]}' ]] || + fail "unused machine prints a record the panel skips" "$record" +pass "an unused machine prints a full empty record the panel skips" for invalid in '{}' '{"default":{}}' '{"default":{"authToken":""}}' 'not-json'; do printf '%s' "$invalid" >"$credentials" - record=$("$collector" --force) + record=$(FREEBUFF_CONFIG_DIR="$signin_home" "$collector" --force) [[ $(printf '%s' "$record" | jq '.ready') == "false" ]] || fail "invalid credentials not ready: $invalid" done pass "empty, unusable, and malformed credentials are not signed in" printf '{"default":{"authToken":"test-token"}}' >"$credentials" -record=$("$collector" --force) -[[ $(printf '%s' "$record" | jq '.ready') == "true" ]] || fail "ready signed-in" -[[ $(printf '%s' "$record" | jq -r '.tierLabel') == "Free" ]] || fail "tierLabel" -[[ $(printf '%s' "$record" | jq -r '.usageStatusText') == "" ]] || fail "usageStatusText" -pass "usable token credentials report the Free tier" - -# A keychain pointer contains no token for this standalone collector. -printf '{"default":{"tokenStore":"keychain"}}' >"$credentials" -record=$("$collector" --force) -[[ $(printf '%s' "$record" | jq '.ready') == "false" ]] || fail "keychain pointer is not verified" -[[ $(printf '%s' "$record" | jq -r '.usageStatusText') == "Freebuff sign-in status unknown" ]] || fail "keychain status" -[[ $(printf '%s' "$record" | jq -r '.authHelpText') == *"cannot be verified"* ]] || fail "keychain help" -pass "keychain-backed sign-in is reported as unverifiable" +record=$(FREEBUFF_CONFIG_DIR="$signin_home" "$collector" --force) +[[ $(printf '%s' "$record" | jq -c '{ready, tierLabel, usageStatusText}') == '{"ready":true,"tierLabel":"","usageStatusText":""}' ]] || + fail "token credentials are a sign-in" "$record" +pass "token credentials make the record ready without naming a plan" + +# Freebuff moves the token to the keychain and leaves the profile behind. +printf '{"default":{"name":"a","tokenStore":"keychain"}}' >"$credentials" +record=$(FREEBUFF_CONFIG_DIR="$signin_home" "$collector" --force) +[[ $(printf '%s' "$record" | jq -c '{ready, usageStatusText, authHelpText}') == '{"ready":true,"usageStatusText":"","authHelpText":""}' ]] || + fail "a keychain sign-in is a sign-in" "$record" +pass "a keychain-backed sign-in makes the record ready" + +# Untrusted message ids: a time far in the future or nonsense falls back to +# the file's last write, and a prompt without text is no prompt. +odd_home="$test_tmp/odd" +make_chat "$odd_home/projects/p/chats/c1" <"$cache_file.new" && mv "$cache_file.new" "$cache_file" +make_chat "$projects/2026-10-01T18-00-00.000Z" <"$cache_file.new" && mv "$cache_file.new" "$cache_file" +make_chat "$projects/2026-10-01T19-00-00.000Z" < Date: Wed, 14 Oct 2026 23:26:22 -0400 Subject: [PATCH 121/546] Add a Qwen Code usage collector Qwen Code writes each session to projects//chats/.jsonl under its runtime directory: typed prompts as user records, and every billed API response as a ui_telemetry record with the model and the provider's token counts, which is what Qwen Code's own usage statistics are rebuilt from. Tokens are read by Qwen Code's own rules (the prompt count includes its cached part, and thinking is added to output only off the OpenAI-style auth types), and records a branched session repeats count once. QWEN_HOME, QWEN_RUNTIME_DIR, and settings.json's advanced.runtimeOutputDir are honored. Transcripts are append-only, so an index reads only what was appended; the scan is cached per day under flock and not cached when a transcript couldn't be read. Qwen Code reads no quota of its own, so there are no meters. Co-Authored-By: Claude Opus 5.5 --- bin/omarchy-agent-usage-qwen | 492 ++++++++++++++++++++++++++ test/shell.d/agent-usage-qwen-test.sh | 181 ++++++++++ 2 files changed, 673 insertions(+) create mode 100755 bin/omarchy-agent-usage-qwen create mode 100755 test/shell.d/agent-usage-qwen-test.sh diff --git a/bin/omarchy-agent-usage-qwen b/bin/omarchy-agent-usage-qwen new file mode 100755 index 00000000000..1402d00b5ee --- /dev/null +++ b/bin/omarchy-agent-usage-qwen @@ -0,0 +1,492 @@ +#!/usr/bin/python3 +# omarchy:summary=Print the Qwen Code usage record as JSON +# omarchy:args=[--force] [--limits-only] +# omarchy:hidden=true +"""Collect Qwen Code usage into one display-ready JSON record. + +Qwen Code records each session as JSONL under +~/.qwen/projects//chats/.jsonl (QWEN_HOME moves ~/.qwen; +QWEN_RUNTIME_DIR or settings.json's advanced.runtimeOutputDir move the +projects). Every API response it bills lands there as a ui_telemetry record +carrying the model and the provider's token counts, the same records Qwen +Code's own usage statistics are rebuilt from, and every typed prompt as a +user record. Qwen Code reads no quota of its own, so the record carries no +meters. The agents panel only ever reads the JSON this prints. +""" + +import argparse +import fcntl +import hashlib +import json +import os +import sys +import tempfile +import time +from datetime import datetime, timedelta, timezone +from pathlib import Path + +AGENT_ID = "qwen" +AGENT_NAME = "Qwen Code" + +API_RESPONSE_EVENT = "qwen-code.api_response" + +# Qwen Code's own rule (getEventTotalTokenCount): through an OpenAI-style API +# the output count already includes thinking, and through Gemini it doesn't. +THOUGHTS_IN_OUTPUT_AUTH_TYPES = ("openai", "qwen-oauth") + +# User records Qwen Code writes for text someone typed: a prompt, or a message +# sent while a turn was still running. Goal, mention, and notification records +# carry a subtype of their own. +PROMPT_SUBTYPES = (None, "mid_turn_user_message") + +# 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). --limits-only promises only fresh limits, and Qwen Code has none, so +# it may reuse a scan for 15 minutes. +SCAN_REUSE_SECONDS = 20 +LIMITS_ONLY_REUSE_SECONDS = 900 + +CACHE_SCHEMA_VERSION = 1 +INDEX_VERSION = 1 +INDEX_HEAD_BYTES = 4096 + +# Counts and times read from disk are untrusted. Anything past these bounds is +# a corrupt or crafted field, not usage. +MAX_TOKENS = 10**12 +MIN_EPOCH_SECONDS = 946_684_800 # 2000-01-01 +MAX_MODEL_NAME_LEN = 120 + + +def token_count(value): + """A token field as a whole number, or 0 when it isn't a plausible count. + + bool is an int subclass, and a string, NaN, negative, or absurd value is a + bad field: it counts as nothing rather than skewing the totals. + """ + if isinstance(value, bool) or not isinstance(value, (int, float)): + return 0 + if value != value or value <= 0 or value > MAX_TOKENS: + return 0 + return int(value) + + +def epoch_seconds(value): + """Seconds since the epoch for an ISO string, or None when the value can't + be a real moment of use.""" + if not isinstance(value, str) or not value.strip(): + return None + text = value.strip() + try: + parsed = datetime.fromisoformat(text[:-1] + "+00:00" if text.endswith("Z") else text) + except ValueError: + return None + if parsed.tzinfo is None: + parsed = parsed.astimezone() + seconds = parsed.timestamp() + # A day of slack past now covers a clock a little behind the writer's. + if seconds < MIN_EPOCH_SECONDS or seconds > time.time() + 86_400: + return None + return seconds + + +def local_day(seconds): + return datetime.fromtimestamp(seconds).strftime("%Y-%m-%d") + + +def obj(value): + return value if isinstance(value, dict) else {} + + +def model_label(raw): + """A model id fit for the record: its last path segment, or None when it + is too long to be a real id.""" + value = str(raw or "").rstrip("/").split("/")[-1] + if len(value) > MAX_MODEL_NAME_LEN: + return None + return value or "qwen" + + +def empty_bucket(): + return { + "inputTokens": 0, + "outputTokens": 0, + "cacheReadInputTokens": 0, + "cacheCreationInputTokens": 0, + } + + +# ------------------------------------------------------------- locations + + +def resolve_dir(value): + """A directory setting as Qwen Code resolves it, `~` included. A relative + one resolves against wherever Qwen Code was started, which this collector + can't know, so it is left out.""" + if not isinstance(value, str) or not value.strip(): + return None + path = Path(value.strip()).expanduser() + return path if path.is_absolute() else None + + +def qwen_home(): + return resolve_dir(os.environ.get("QWEN_HOME")) or (Path.home() / ".qwen") + + +def runtime_output_dir(home): + try: + settings = json.loads((home / "settings.json").read_text(encoding="utf-8")) + except (OSError, ValueError): + return None + return resolve_dir(obj(obj(settings).get("advanced")).get("runtimeOutputDir")) + + +def runtime_dirs(): + """Every directory whose projects/ may hold chats: QWEN_RUNTIME_DIR or the + configured runtimeOutputDir, where Qwen Code writes, and the Qwen home, + where it writes by default and where its own usage rebuild looks.""" + home = qwen_home() + dirs = [] + for candidate in (resolve_dir(os.environ.get("QWEN_RUNTIME_DIR")), runtime_output_dir(home), home): + if candidate is not None and candidate not in dirs: + dirs.append(candidate) + return dirs + + +def chat_files(): + files = [] + for base in runtime_dirs(): + try: + found = sorted((base / "projects").glob("*/chats/*.jsonl")) + except OSError: + continue + # A ledger sits beside its transcript and repeats nothing billable. + files.extend(path for path in found if not path.name.endswith(".ledger.jsonl") and path not in files) + return files + + +# ------------------------------------------------------------- transcripts + + +def parse_line(raw, session): + """One transcript line as an index row, or None. + + A prompt is ["p", key, seconds, session]; a billed response is + ["r", key, seconds, session, model, input, output, cacheRead]. The key + dedups a record a resumed or branched session carries again. + """ + if b'"user"' not in raw and API_RESPONSE_EVENT.encode() not in raw: + return None + try: + record = json.loads(raw) + except ValueError: + return None + if not isinstance(record, dict): + return None + session = str(record.get("sessionId") or session) + uuid = str(record.get("uuid") or "") + + if record.get("type") == "user": + if record.get("subtype") not in PROMPT_SUBTYPES: + return None + seconds = epoch_seconds(record.get("timestamp")) + if seconds is None or not uuid: + return None + return ["p", uuid, seconds, session] + + if record.get("type") != "system" or record.get("subtype") != "ui_telemetry": + return None + event = obj(obj(record.get("systemPayload")).get("uiEvent")) + if event.get("event.name") != API_RESPONSE_EVENT: + return None + seconds = epoch_seconds(event.get("event.timestamp")) or epoch_seconds(record.get("timestamp")) + model = model_label(event.get("model")) + if seconds is None or model is None: + return None + # The prompt count includes its cached part; a provider that reports only + # the cached part reports a fully cached prompt, as Qwen Code reads it. + cached = token_count(event.get("cached_content_token_count")) + prompt = token_count(event.get("input_token_count")) or cached + cached = min(cached, prompt) + output = token_count(event.get("output_token_count")) + if str(event.get("auth_type") or "") not in THOUGHTS_IN_OUTPUT_AUTH_TYPES: + output += token_count(event.get("thoughts_token_count")) + if prompt + output <= 0: + return None + key = str(event.get("response_id") or "") or uuid + if not key: + return None + return ["r", key, seconds, session, model, prompt - cached, output, cached] + + +def read_digest(path, length): + digest = hashlib.sha1() + with path.open("rb") as handle: + digest.update(handle.read(min(length, INDEX_HEAD_BYTES))) + return digest.hexdigest() + + +def parse_file(path, entry): + """The index entry for one transcript, reusing `entry` where the file only + grew: Qwen Code appends to a transcript and never rewrites what it wrote.""" + stat = path.stat() + if entry and entry.get("size") == stat.st_size and entry.get("mtime") == stat.st_mtime_ns: + return entry + offset = 0 + rows = [] + if ( + entry + and isinstance(entry.get("offset"), int) + and stat.st_size > entry["offset"] + and entry.get("file") == [stat.st_dev, stat.st_ino] + and entry.get("head") == read_digest(path, entry["offset"]) + ): + offset = entry["offset"] + rows = list(entry.get("rows") or []) + with path.open("rb") as handle: + handle.seek(offset) + for raw in handle: + if not raw.endswith(b"\n"): + # A line still being written; pick it up whole next time. + break + offset += len(raw) + row = parse_line(raw, path.stem) + if row is not None: + rows.append(row) + return { + "size": stat.st_size, + "mtime": stat.st_mtime_ns, + "file": [stat.st_dev, stat.st_ino], + "offset": offset, + "head": read_digest(path, offset), + "rows": rows, + } + + +def load_index(path): + try: + data = json.loads(path.read_text(encoding="utf-8")) + except Exception: + return {} + if isinstance(data, dict) and data.get("version") == INDEX_VERSION and isinstance(data.get("files"), dict): + return data["files"] + return {} + + +def summarize(rows): + now = datetime.now() + today = now.strftime("%Y-%m-%d") + recent_dates = [(now - timedelta(days=offset)).strftime("%Y-%m-%d") for offset in range(6, -1, -1)] + recent = {day: 0 for day in recent_dates} + today_prompts = 0 + today_total_tokens = 0 + today_sessions = set() + today_tokens_by_model = {} + total_prompts = 0 + total_sessions = set() + active_days = set() + model_usage = {} + seen = set() + + for row in rows: + try: + kind, key, seconds, session = row[:4] + day = local_day(float(seconds)) + if kind == "r": + model = str(row[4]) + input_tokens, output_tokens, cache_read = (int(value) for value in row[5:8]) + elif kind != "p": + continue + except (TypeError, ValueError, IndexError, OverflowError, OSError): + continue + if (kind, key) in seen: + continue + seen.add((kind, key)) + session = str(session) + total_sessions.add(session) + active_days.add(day) + if day == today: + today_sessions.add(session) + if kind == "p": + total_prompts += 1 + if day == today: + today_prompts += 1 + continue + total = input_tokens + output_tokens + cache_read + bucket = model_usage.setdefault(model, empty_bucket()) + bucket["inputTokens"] += input_tokens + bucket["outputTokens"] += output_tokens + bucket["cacheReadInputTokens"] += cache_read + if day in recent: + # recentDays.messageCount holds tokens, despite the shared legacy name. + recent[day] += total + if day == today: + today_total_tokens += total + today_tokens_by_model[model] = today_tokens_by_model.get(model, 0) + total + + return { + "todayPrompts": today_prompts, + "todaySessions": len(today_sessions), + "todayTotalTokens": today_total_tokens, + "todayTokensByModel": today_tokens_by_model, + "recentDays": [{"date": day, "messageCount": recent[day]} for day in recent_dates], + "totalPrompts": total_prompts, + "totalSessions": len(total_sessions), + # The dates travel too: merging snapshots from several machines needs + # their union, which a count alone cannot give. + "activeDays": len(active_days), + "activeDates": sorted(active_days), + "modelUsage": model_usage, + } + + +def scan(index_file): + """Every transcript's rows, summarized. Returns (stats, complete); a scan + that couldn't read a transcript is not cached.""" + previous = load_index(index_file) if index_file else {} + current = {} + complete = True + rows = [] + for path in chat_files(): + key = str(path) + try: + current[key] = parse_file(path, previous.get(key)) + except OSError as exc: + print(f"omarchy-agent-usage-qwen: skipping {path}: {type(exc).__name__}", file=sys.stderr) + complete = False + continue + rows.extend(current[key].get("rows") or []) + if index_file and current != previous: + try: + write_json(index_file, {"version": INDEX_VERSION, "files": current}) + except Exception as exc: + print(f"omarchy-agent-usage-qwen: could not write transcript index ({exc})", file=sys.stderr) + return summarize(rows), complete + + +# ---------------------------------------------------------------- cache + + +def cache_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 cache_paths(): + # The digest covers every directory the scan reads, so another home or + # runtime directory never reuses this one's numbers. + digest = hashlib.sha1("\n".join(str(path) for path in runtime_dirs()).encode("utf-8")).hexdigest()[:16] + root = cache_root() + return root / f"qwen-scan-{digest}.json", root / f"qwen-scan-{digest}.lock", root / f"qwen-index-{digest}.json" + + +def read_fresh_json(path, max_age_seconds): + if max_age_seconds <= 0: + 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, payload): + # A temp name unique to this writer: 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 + + +STATS_KEYS = ("todayPrompts", "todaySessions", "todayTotalTokens", "todayTokensByModel", "recentDays", "totalPrompts", "totalSessions", "activeDays", "activeDates", "modelUsage") + + +def read_cached_stats(cache_file, max_age_seconds): + cached = read_fresh_json(cache_file, max_age_seconds) + if not isinstance(cached, dict) or cached.get("schemaVersion") != CACHE_SCHEMA_VERSION: + return None + # today* fields only mean "today" on the day they were scanned. + if cached.get("scanDate") != datetime.now().strftime("%Y-%m-%d"): + return None + stats = cached.get("stats") + if not isinstance(stats, dict) or not all(key in stats for key in STATS_KEYS): + return None + return stats + + +def cached_local_stats(max_age): + """Local stats, with the cache as a pure optimization: any cache failure + degrades to a direct scan and a warning on stderr.""" + try: + cache_file, lock_file, index_file = cache_paths() + cached = read_cached_stats(cache_file, max_age) + if cached is not None: + return cached + with lock_file.open("w") as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + cached = read_cached_stats(cache_file, max_age) + if cached is not None: + return cached + stats, complete = scan(index_file) + # An interrupted scan still serves this run, but caching it would hide + # the missing usage from every reader until the cache expires. + if complete: + try: + write_json(cache_file, {"schemaVersion": CACHE_SCHEMA_VERSION, "scanDate": datetime.now().strftime("%Y-%m-%d"), "stats": stats}) + except Exception as exc: + print(f"omarchy-agent-usage-qwen: could not write usage cache ({exc})", file=sys.stderr) + return stats + except Exception as exc: + print(f"omarchy-agent-usage-qwen: cache unavailable ({exc}); scanning directly", file=sys.stderr) + stats, _ = scan(None) + return stats + + +def main(): + parser = argparse.ArgumentParser(description="Print the Qwen Code usage record as JSON") + # --force rescans everything. --limits-only only promises fresh limits, and + # Qwen Code reads none, so it reuses a scan far longer than a normal run. + parser.add_argument("--force", action="store_true") + parser.add_argument("--limits-only", action="store_true") + args = parser.parse_args() + + max_age = 0 if args.force else (LIMITS_ONLY_REUSE_SECONDS if args.limits_only else SCAN_REUSE_SECONDS) + stats = cached_local_stats(max_age) + + # There is no quota to read, so nothing waits on a sign-in: the status stays + # empty and local stats show alone. A machine that never ran Qwen Code + # prints a record the panel skips. + record = { + "schemaVersion": 1, + "id": AGENT_ID, + "name": AGENT_NAME, + "updatedAt": datetime.now(timezone.utc).isoformat(), + "ready": stats["totalSessions"] > 0, + "hasLocalStats": True, + "hasPromptStats": True, + "tierLabel": "", + "usageStatusText": "", + "authHelpText": "", + "limits": [], + "limitsStale": False, + "limitsFetchedAt": 0, + } + record.update(stats) + print(json.dumps(record, separators=(",", ":"))) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/test/shell.d/agent-usage-qwen-test.sh b/test/shell.d/agent-usage-qwen-test.sh new file mode 100755 index 00000000000..085e6c06b77 --- /dev/null +++ b/test/shell.d/agent-usage-qwen-test.sh @@ -0,0 +1,181 @@ +#!/bin/bash + +source "$(dirname "$0")/base-test.sh" + +require_command jq +require_command python3 + +COLLECTOR="$ROOT/bin/omarchy-agent-usage-qwen" +TMP_ROOT=$(mktemp -d) +trap 'rm -rf "$TMP_ROOT"' EXIT + +unset QWEN_HOME QWEN_RUNTIME_DIR XDG_CACHE_HOME + +today=$(date +%Y-%m-%d) +now_iso=$(date -u +%Y-%m-%dT%H:%M:%S.000Z) +two_days_ago=$(date -d '2 days ago' +%Y-%m-%d) +two_days_ago_iso=$(date -u -d '2 days ago' +%Y-%m-%dT%H:%M:%S.000Z) + +new_home() { + local home="$TMP_ROOT/$1" + mkdir -p "$home" + printf '%s\n' "$home" +} + +collect() { + local home="$1" + shift + HOME="$home" XDG_CACHE_HOME="$home/.cache" "$COLLECTOR" "$@" +} + +# prompt