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-deepseek b/bin/omarchy-agent-usage-deepseek new file mode 100755 index 00000000000..fedc514e346 --- /dev/null +++ b/bin/omarchy-agent-usage-deepseek @@ -0,0 +1,478 @@ +#!/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": 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. + + 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, 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) + 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/bin/omarchy-agent-usage-hermes b/bin/omarchy-agent-usage-hermes new file mode 100755 index 00000000000..33051b27e92 --- /dev/null +++ b/bin/omarchy-agent-usage-hermes @@ -0,0 +1,145 @@ +#!/usr/bin/python3 +# omarchy:summary=Print the Hermes Agent usage record as JSON +# 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 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. +""" + +import argparse +import json +import os +import shutil +import sqlite3 +from datetime import datetime, timedelta, timezone +from pathlib import Path + +RECENT_DAYS = 7 + + +# 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: + scan(path, totals, today, week) + read += 1 + except sqlite3.Error: + 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__": + main() diff --git a/bin/omarchy-agent-usage-minimax b/bin/omarchy-agent-usage-minimax new file mode 100755 index 00000000000..942041cea35 --- /dev/null +++ b/bin/omarchy-agent-usage-minimax @@ -0,0 +1,323 @@ +#!/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. + +The last good limits are cached, so a failed check keeps showing them, marked +stale, rather than dropping the section. +""" + +from __future__ import annotations + +import argparse +import json +import os +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 = "minimax" +AGENT_NAME = "MiniMax" +DEFAULT_BASE_URL = "https://api.minimax.io" +AUTH_HELP = "Set MINIMAX_API_KEY, or sign in to MiniMax in opencode." +TIER_LABEL = "Token Plan" + + +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": "", + "usageStatusText": "", + "authHelpText": "", + "limits": [], + "limitsStale": False, + "limitsFetchedAt": 0, + } + 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 + + +# 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")))) + + +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") 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 + + +# 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]: + if not api_key: + return base_record(usageStatusText="MiniMax unavailable", authHelpText=AUTH_HELP) + + 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) + try: + write_json(limits_cache(), {"fetchedAtMs": fetched_at, "limits": limits}) + except OSError: + pass + 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. + 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 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/bin/omarchy-agent-usage-ollama b/bin/omarchy-agent-usage-ollama new file mode 100755 index 00000000000..f4db9d81939 --- /dev/null +++ b/bin/omarchy-agent-usage-ollama @@ -0,0 +1,432 @@ +#!/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 +# 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: + 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" + # 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") + + 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, {"scanDate": today, "modelPrefix": model_prefix, "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 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": ""}) + if isinstance(weekly, dict): + percent = clamp_fraction(weekly.get("usage")) + if percent >= 0: + 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} + + +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]: + # `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 + # 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." + 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 + result["live"] = True + 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)) + 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 + # 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": 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 + record.update(stats or {}) + print(json.dumps(record, separators=(",", ":"), sort_keys=True)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/bin/omarchy-agent-usage-openrouter b/bin/omarchy-agent-usage-openrouter new file mode 100755 index 00000000000..836dc6cdea1 --- /dev/null +++ b/bin/omarchy-agent-usage-openrouter @@ -0,0 +1,600 @@ +#!/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 time +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 + # 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 + + +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": "", + "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]: + 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 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) + manager = management_key(config) + if not key and not manager: + 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(manager, 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 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) + + 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 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/bin/omarchy-agent-usage-zai b/bin/omarchy-agent-usage-zai new file mode 100755 index 00000000000..989c1da33c1 --- /dev/null +++ b/bin/omarchy-agent-usage-zai @@ -0,0 +1,676 @@ +#!/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. `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, "live": False, "fetchedAtMs": 0} + + 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)) + 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 "") + 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": result["fetchedAtMs"], + "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["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. + 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"] + 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"] + 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/manual/17-ai.md b/manual/17-ai.md index ea8a8d65eaf..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. +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/shell/plugins/agents/README.md b/shell/plugins/agents/README.md index c9daebba45e..d6bf6679bd4 100644 --- a/shell/plugins/agents/README.md +++ b/shell/plugins/agents/README.md @@ -74,6 +74,13 @@ 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 | +| `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 | +| `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 @@ -91,7 +98,11 @@ 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`, 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 @@ -101,6 +112,48 @@ 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. + +### 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, 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 +`~/.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. + +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. + +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 @@ -129,6 +182,40 @@ 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; +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, +non-estimated balance. A key with a credit limit additionally gets a +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: + +```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/assets/agy-light.svg b/shell/plugins/agents/assets/agy-light.svg new file mode 100644 index 00000000000..67ced694468 --- /dev/null +++ b/shell/plugins/agents/assets/agy-light.svg @@ -0,0 +1 @@ +Antigravity diff --git a/shell/plugins/agents/assets/agy.svg b/shell/plugins/agents/assets/agy.svg new file mode 100644 index 00000000000..3743be8edf4 --- /dev/null +++ b/shell/plugins/agents/assets/agy.svg @@ -0,0 +1 @@ +Antigravity 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/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/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 7ca52484b12..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, and Fireworks 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": { @@ -22,7 +22,14 @@ "claude": { "enabled": true }, "codex": { "enabled": true }, "grok": { "enabled": true }, - "fireworks": { "enabled": true } + "fireworks": { "enabled": true }, + "hermes": { "enabled": true }, + "agy": { "enabled": true }, + "zai": { "enabled": true }, + "ollama": { "enabled": true }, + "minimax": { "enabled": true }, + "openrouter": { "enabled": true }, + "deepseek": { "enabled": true } }, "refreshIntervalSec": 900, "syncMode": "Off", diff --git a/test/shell.d/agent-usage-agy-limits-test.sh b/test/shell.d/agent-usage-agy-limits-test.sh new file mode 100755 index 00000000000..1198a4356e8 --- /dev/null +++ b/test/shell.d/agent-usage-agy-limits-test.sh @@ -0,0 +1,205 @@ +#!/bin/bash + +source "$(dirname "$0")/base-test.sh" + +require_command jq +require_command python3 + +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. + +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" + +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-deepseek-scanner-test.sh b/test/shell.d/agent-usage-deepseek-scanner-test.sh new file mode 100755 index 00000000000..91226a78aa1 --- /dev/null +++ b/test/shell.d/agent-usage-deepseek-scanner-test.sh @@ -0,0 +1,225 @@ +#!/bin/bash + +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 + +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) + ":" + .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) +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" 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..9fcd0fcfc47 --- /dev/null +++ b/test/shell.d/agent-usage-hermes-test.sh @@ -0,0 +1,97 @@ +#!/bin/bash + +source "$(dirname "$0")/base-test.sh" + +require_command jq +require_command python3 + +TEST_HOME=$(mktemp -d) +trap 'rm -rf "$TEST_HOME"' EXIT + +# 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 + +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) + ); + """ +) +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)], +) +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 +} + +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" + +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 '.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" + +# 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" + +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" 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..37ff0f769a5 --- /dev/null +++ b/test/shell.d/agent-usage-minimax-test.sh @@ -0,0 +1,152 @@ +#!/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" 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":[]}' ]] || + 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 +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") +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) +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 + +# 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") + +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, + "savedBaseUrl": saved_base_url, + "oauthKey": oauth_key, + "oauthBaseUrl": oauth_base_url, + "opencodeKey": opencode_key, + "record": record, + "stale": stale, + "reset": reset, +})) +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, 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" + +[[ $(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" 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..7437010c7bf --- /dev/null +++ b/test/shell.d/agent-usage-ollama-scanner-test.sh @@ -0,0 +1,218 @@ +#!/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 + ":" + .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" + +# 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" + +# 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. +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" + +# 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}}') +[[ $(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" +[[ $(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. +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" +[[ $(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. +[[ $(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" 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..a6b93b71f3d --- /dev/null +++ b/test/shell.d/agent-usage-openrouter-scanner-test.sh @@ -0,0 +1,238 @@ +#!/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) + ":" + .tierLabel' <<<"$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].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' +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" + +# 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" + +# 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" + +# 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" 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..36915183fac --- /dev/null +++ b/test/shell.d/agent-usage-zai-limits-test.sh @@ -0,0 +1,155 @@ +#!/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" +[[ $(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. +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" +[[ $(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. +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" <